Fix cleanup rake task limit

Methods defined in rake tasks are defined globally.

changelogs/unreleased/mk-fix-cleanup-rake-task-limit.yml

+---
+title: Fix gitlab:cleanup:orphan_job_artifact_files rake task limit
+merge_request: 52056
+author:
+type: fixed

lib/gitlab/cleanup/orphan_job_artifact_files.rb

@@ -12,10 +12,9 @@ module Gitlab
       VALID_NICENESS_LEVELS = %w{none realtime best-effort idle}.freeze
 
       attr_accessor :batch, :total_found, :total_cleaned
-      attr_reader :limit, :dry_run, :niceness, :logger
+      attr_reader :dry_run, :niceness, :logger
 
-      def initialize(limit: nil, dry_run: true, niceness: nil, logger: nil)
-        @limit = limit
+      def initialize(dry_run: true, niceness: nil, logger: nil)
         @dry_run = dry_run
         @niceness = (niceness || DEFAULT_NICENESS).downcase
         @logger = logger || Gitlab::AppLogger
@@ -31,7 +30,11 @@ module Gitlab
           batch << artifact_file
 
           clean_batch! if batch.full?
-          break if limit_reached?
+
+          if limit_reached?
+            log_info("Exiting due to reaching limit of #{limit}.")
+            break
+          end
         end
 
         clean_batch!
@@ -128,6 +131,10 @@ module Gitlab
       def log_error(msg, params = {})
         logger.error(msg)
       end
+
+      def limit
+        ENV['LIMIT']&.to_i
+      end
     end
   end
 end

lib/gitlab/cleanup/orphan_lfs_file_references.rb

@@ -5,15 +5,14 @@ module Gitlab
     class OrphanLfsFileReferences
       include Gitlab::Utils::StrongMemoize
 
-      attr_reader :project, :dry_run, :logger, :limit
+      attr_reader :project, :dry_run, :logger
 
       DEFAULT_REMOVAL_LIMIT = 1000
 
-      def initialize(project, dry_run: true, logger: nil, limit: nil)
+      def initialize(project, dry_run: true, logger: nil)
         @project = project
         @dry_run = dry_run
         @logger = logger || Gitlab::AppLogger
-        @limit = limit
       end
 
       def run!
@@ -67,6 +66,10 @@ module Gitlab
       def log_info(msg)
         logger.info("#{'[DRY RUN] ' if dry_run}#{msg}")
       end
+
+      def limit
+        ENV['LIMIT']&.to_i
+      end
     end
   end
 end

lib/tasks/gitlab/cleanup.rake

@@ -56,7 +56,7 @@ namespace :gitlab do
     task orphan_job_artifact_files: :gitlab_environment do
       warn_user_is_not_gitlab
 
-      cleaner = Gitlab::Cleanup::OrphanJobArtifactFiles.new(limit: limit, dry_run: dry_run?, niceness: niceness, logger: logger)
+      cleaner = Gitlab::Cleanup::OrphanJobArtifactFiles.new(dry_run: dry_run?, niceness: niceness, logger: logger)
       cleaner.run!
 
       if dry_run?
@@ -78,8 +78,7 @@ namespace :gitlab do
       cleaner = Gitlab::Cleanup::OrphanLfsFileReferences.new(
         project,
         dry_run: dry_run?,
-        logger: logger,
-        limit: limit
+        logger: logger
       )
 
       cleaner.run!
@@ -162,10 +161,6 @@ namespace :gitlab do
       ENV['DEBUG'].present?
     end
 
-    def limit
-      ENV['LIMIT']&.to_i
-    end
-
     def niceness
       ENV['NICENESS'].presence
     end

lib/tasks/gitlab/snippets.rake

@@ -13,7 +13,7 @@ namespace :gitlab do
         raise "Please supply the list of ids through the SNIPPET_IDS env var"
       end
 
-      raise "Invalid limit value" if limit == 0
+      raise "Invalid limit value" if snippet_task_limit == 0
 
       if migration_running?
         raise "There are already snippet migrations running. Please wait until they are finished."
@@ -41,8 +41,8 @@ namespace :gitlab do
         end
       end
 
-      if ids.size > limit
-        raise "The number of ids provided is higher than #{limit}. You can update this limit by using the env var `LIMIT`"
+      if ids.size > snippet_task_limit
+        raise "The number of ids provided is higher than #{snippet_task_limit}. You can update this limit by using the env var `LIMIT`"
       end
 
       ids
@@ -68,14 +68,14 @@ namespace :gitlab do
     #   bundle exec rake gitlab:snippets:list_non_migrated LIMIT=50
     desc 'GitLab | Show non migrated snippets'
     task list_non_migrated: :environment do
-      raise "Invalid limit value" if limit == 0
+      raise "Invalid limit value" if snippet_task_limit == 0
 
       non_migrated_count = non_migrated_snippets.count
       if non_migrated_count == 0
         puts "All snippets have been successfully migrated"
       else
         puts "There are #{non_migrated_count} snippets that haven't been migrated. Showing a batch of ids of those snippets:\n"
-        puts non_migrated_snippets.limit(limit).pluck(:id).join(',')
+        puts non_migrated_snippets.limit(snippet_task_limit).pluck(:id).join(',')
       end
     end
 
@@ -84,7 +84,7 @@ namespace :gitlab do
     end
 
     # There are problems with the specs if we memoize this value
-    def limit
+    def snippet_task_limit
       ENV['LIMIT'] ? ENV['LIMIT'].to_i : DEFAULT_LIMIT
     end
   end

spec/lib/gitlab/cleanup/orphan_job_artifact_files_spec.rb

@@ -42,7 +42,8 @@ RSpec.describe Gitlab::Cleanup::OrphanJobArtifactFiles do
   end
 
   it 'stops when limit is reached' do
-    cleanup = described_class.new(limit: 1)
+    stub_env('LIMIT', 1)
+    cleanup = described_class.new
 
     mock_artifacts_found(cleanup, 'tmp/foo/bar/1', 'tmp/foo/bar/2')
 

spec/tasks/gitlab/cleanup_rake_spec.rb

@@ -109,8 +109,7 @@ RSpec.describe 'gitlab:cleanup rake tasks' do
       it 'passes dry_run correctly' do
         expect(Gitlab::Cleanup::OrphanJobArtifactFiles)
           .to receive(:new)
-          .with(limit: anything,
-                dry_run: false,
+          .with(dry_run: false,
                 niceness: anything,
                 logger: anything)
           .and_call_original
@@ -145,7 +144,6 @@ RSpec.describe 'gitlab:cleanup rake tasks' do
         expect(Gitlab::Cleanup::OrphanLfsFileReferences)
           .to receive(:new)
           .with(project,
-                limit: anything,
                 dry_run: false,
                 logger: anything)
           .and_call_original
@@ -153,24 +151,6 @@ RSpec.describe 'gitlab:cleanup rake tasks' do
         rake_task
       end
     end
-
-    context 'with LIMIT set to 100' do
-      before do
-        stub_env('LIMIT', '100')
-      end
-
-      it 'passes limit as integer' do
-        expect(Gitlab::Cleanup::OrphanLfsFileReferences)
-          .to receive(:new)
-          .with(project,
-                limit: 100,
-                dry_run: true,
-                logger: anything)
-          .and_call_original
-
-        rake_task
-      end
-    end
   end
 
   describe 'gitlab:cleanup:orphan_lfs_files' do