Merge branch '343055-put-bg-migrations-in-correct-queue' into 'master'

Select worker to schedule BG migrations

See merge request gitlab-org/gitlab!73306

lib/gitlab/background_migration/job_coordinator.rb

@@ -36,6 +36,8 @@ module Gitlab
 
       attr_reader :worker_class
 
+      delegate :minimum_interval, :perform_in, to: :worker_class
+
       def queue
         @queue ||= worker_class.sidekiq_options['queue']
       end

lib/gitlab/database/migrations/background_migration_helpers.rb

 # frozen_string_literal: true
-
 module Gitlab
   module Database
     module Migrations
@@ -45,11 +44,11 @@ module Gitlab
           raise "#{model_class} does not have an ID column of #{primary_column_name} to use for batch ranges" unless model_class.column_names.include?(primary_column_name.to_s)
           raise "#{primary_column_name} is not an integer column" unless model_class.columns_hash[primary_column_name.to_s].type == :integer
 
+          job_coordinator = coordinator_for_tracking_database
+
           # To not overload the worker too much we enforce a minimum interval both
           # when scheduling and performing jobs.
-          if delay_interval < BackgroundMigrationWorker.minimum_interval
-            delay_interval = BackgroundMigrationWorker.minimum_interval
-          end
+          delay_interval = [delay_interval, job_coordinator.minimum_interval].max
 
           final_delay = 0
           batch_counter = 0
@@ -60,14 +59,14 @@ module Gitlab
 
             start_id, end_id = relation.pluck(min, max).first
 
-            # `BackgroundMigrationWorker.bulk_perform_in` schedules all jobs for
+            # `SingleDatabaseWorker.bulk_perform_in` schedules all jobs for
             # the same time, which is not helpful in most cases where we wish to
             # spread the work over time.
             final_delay = initial_delay + delay_interval * index
             full_job_arguments = [start_id, end_id] + other_job_arguments
 
             track_in_database(job_class_name, full_job_arguments) if track_jobs
-            migrate_in(final_delay, job_class_name, full_job_arguments)
+            migrate_in(final_delay, job_class_name, full_job_arguments, coordinator: job_coordinator)
 
             batch_counter += 1
           end
@@ -91,9 +90,11 @@ module Gitlab
         # delay_interval - The duration between each job's scheduled time
         # batch_size - The maximum number of jobs to fetch to memory from the database.
         def requeue_background_migration_jobs_by_range_at_intervals(job_class_name, delay_interval, batch_size: BATCH_SIZE, initial_delay: 0)
+          job_coordinator = coordinator_for_tracking_database
+
           # To not overload the worker too much we enforce a minimum interval both
           # when scheduling and performing jobs.
-          delay_interval = [delay_interval, BackgroundMigrationWorker.minimum_interval].max
+          delay_interval = [delay_interval, job_coordinator.minimum_interval].max
 
           final_delay = 0
           job_counter = 0
@@ -103,7 +104,7 @@ module Gitlab
             job_batch.each do |job|
               final_delay = initial_delay + delay_interval * job_counter
 
-              migrate_in(final_delay, job_class_name, job.arguments)
+              migrate_in(final_delay, job_class_name, job.arguments, coordinator: job_coordinator)
 
               job_counter += 1
             end
@@ -132,17 +133,20 @@ module Gitlab
         # This method does not garauntee that all jobs completed successfully.
         # It can only be used if the previous background migration used the queue_background_migration_jobs_by_range_at_intervals helper.
         def finalize_background_migration(class_name, delete_tracking_jobs: ['succeeded'])
+          job_coordinator = coordinator_for_tracking_database
+
           # Empty the sidekiq queue.
-          Gitlab::BackgroundMigration.steal(class_name)
+          job_coordinator.steal(class_name)
 
           # Process pending tracked jobs.
           jobs = Gitlab::Database::BackgroundMigrationJob.pending.for_migration_class(class_name)
+
           jobs.find_each do |job|
-            BackgroundMigrationWorker.new.perform(job.class_name, job.arguments)
+            job_coordinator.perform(job.class_name, job.arguments)
           end
 
           # Empty the sidekiq queue.
-          Gitlab::BackgroundMigration.steal(class_name)
+          job_coordinator.steal(class_name)
 
           # Delete job tracking rows.
           delete_job_tracking(class_name, status: delete_tracking_jobs) if delete_tracking_jobs
@@ -152,36 +156,14 @@ module Gitlab
           Rails.env.test? || Rails.env.development?
         end
 
-        def migrate_async(*args)
-          with_migration_context do
-            BackgroundMigrationWorker.perform_async(*args)
-          end
-        end
-
-        def migrate_in(*args)
-          with_migration_context do
-            BackgroundMigrationWorker.perform_in(*args)
-          end
-        end
-
-        def bulk_migrate_in(*args)
-          with_migration_context do
-            BackgroundMigrationWorker.bulk_perform_in(*args)
-          end
-        end
-
-        def bulk_migrate_async(*args)
+        def migrate_in(*args, coordinator: coordinator_for_tracking_database)
           with_migration_context do
-            BackgroundMigrationWorker.bulk_perform_async(*args)
-          end
+            coordinator.perform_in(*args)
           end
-
-        def with_migration_context(&block)
-          Gitlab::ApplicationContext.with_context(caller_id: self.class.to_s, &block)
         end
 
         def delete_queued_jobs(class_name)
-          Gitlab::BackgroundMigration.steal(class_name) do |job|
+          coordinator_for_tracking_database.steal(class_name) do |job|
             job.delete
 
             false
@@ -196,9 +178,21 @@ module Gitlab
 
         private
 
+        def with_migration_context(&block)
+          Gitlab::ApplicationContext.with_context(caller_id: self.class.to_s, &block)
+        end
+
         def track_in_database(class_name, arguments)
           Gitlab::Database::BackgroundMigrationJob.create!(class_name: class_name, arguments: arguments)
         end
+
+        def coordinator_for_tracking_database
+          Gitlab::BackgroundMigration.coordinator_for_database(tracking_database)
+        end
+
+        def tracking_database
+          Gitlab::BackgroundMigration::DEFAULT_TRACKING_DATABASE
+        end
       end
     end
   end

spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb

@@ -7,6 +7,11 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
     ActiveRecord::Migration.new.extend(described_class)
   end
 
+  shared_examples_for 'helpers that enqueue background migrations' do |worker_class, tracking_database|
+    before do
+      allow(model).to receive(:tracking_database).and_return(tracking_database)
+    end
+
     describe '#queue_background_migration_jobs_by_range_at_intervals' do
       context 'when the model has an ID column' do
         let!(:id1) { create(:user).id }
@@ -39,15 +44,30 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
           end
         end
 
+        context 'when the delay_interval is smaller than the minimum' do
+          it 'sets the delay_interval to the minimum value' do
+            Sidekiq::Testing.fake! do
+              final_delay = model.queue_background_migration_jobs_by_range_at_intervals(User, 'FooJob', 1.minute, batch_size: 2)
+
+              expect(worker_class.jobs[0]['args']).to eq(['FooJob', [id1, id2]])
+              expect(worker_class.jobs[0]['at']).to eq(2.minutes.from_now.to_f)
+              expect(worker_class.jobs[1]['args']).to eq(['FooJob', [id3, id3]])
+              expect(worker_class.jobs[1]['at']).to eq(4.minutes.from_now.to_f)
+
+              expect(final_delay.to_f).to eq(4.minutes.to_f)
+            end
+          end
+        end
+
         context 'with batch_size option' do
           it 'queues jobs correctly' do
             Sidekiq::Testing.fake! do
               model.queue_background_migration_jobs_by_range_at_intervals(User, 'FooJob', 10.minutes, batch_size: 2)
 
-            expect(BackgroundMigrationWorker.jobs[0]['args']).to eq(['FooJob', [id1, id2]])
-            expect(BackgroundMigrationWorker.jobs[0]['at']).to eq(10.minutes.from_now.to_f)
-            expect(BackgroundMigrationWorker.jobs[1]['args']).to eq(['FooJob', [id3, id3]])
-            expect(BackgroundMigrationWorker.jobs[1]['at']).to eq(20.minutes.from_now.to_f)
+              expect(worker_class.jobs[0]['args']).to eq(['FooJob', [id1, id2]])
+              expect(worker_class.jobs[0]['at']).to eq(10.minutes.from_now.to_f)
+              expect(worker_class.jobs[1]['args']).to eq(['FooJob', [id3, id3]])
+              expect(worker_class.jobs[1]['at']).to eq(20.minutes.from_now.to_f)
             end
           end
         end
@@ -57,8 +77,8 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
             Sidekiq::Testing.fake! do
               model.queue_background_migration_jobs_by_range_at_intervals(User, 'FooJob', 10.minutes)
 
-            expect(BackgroundMigrationWorker.jobs[0]['args']).to eq(['FooJob', [id1, id3]])
-            expect(BackgroundMigrationWorker.jobs[0]['at']).to eq(10.minutes.from_now.to_f)
+              expect(worker_class.jobs[0]['args']).to eq(['FooJob', [id1, id3]])
+              expect(worker_class.jobs[0]['at']).to eq(10.minutes.from_now.to_f)
             end
           end
         end
@@ -68,8 +88,8 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
             Sidekiq::Testing.fake! do
               model.queue_background_migration_jobs_by_range_at_intervals(User, 'FooJob', 10.minutes, other_job_arguments: [1, 2])
 
-            expect(BackgroundMigrationWorker.jobs[0]['args']).to eq(['FooJob', [id1, id3, 1, 2]])
-            expect(BackgroundMigrationWorker.jobs[0]['at']).to eq(10.minutes.from_now.to_f)
+              expect(worker_class.jobs[0]['args']).to eq(['FooJob', [id1, id3, 1, 2]])
+              expect(worker_class.jobs[0]['at']).to eq(10.minutes.from_now.to_f)
             end
           end
         end
@@ -79,8 +99,8 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
             Sidekiq::Testing.fake! do
               model.queue_background_migration_jobs_by_range_at_intervals(User, 'FooJob', 10.minutes, other_job_arguments: [1, 2], initial_delay: 10.minutes)
 
-            expect(BackgroundMigrationWorker.jobs[0]['args']).to eq(['FooJob', [id1, id3, 1, 2]])
-            expect(BackgroundMigrationWorker.jobs[0]['at']).to eq(20.minutes.from_now.to_f)
+              expect(worker_class.jobs[0]['args']).to eq(['FooJob', [id1, id3, 1, 2]])
+              expect(worker_class.jobs[0]['at']).to eq(20.minutes.from_now.to_f)
             end
           end
         end
@@ -93,7 +113,7 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
                   other_job_arguments: [1, 2], track_jobs: true)
               end.to change { Gitlab::Database::BackgroundMigrationJob.count }.from(0).to(1)
 
-            expect(BackgroundMigrationWorker.jobs.size).to eq(1)
+              expect(worker_class.jobs.size).to eq(1)
 
               tracked_job = Gitlab::Database::BackgroundMigrationJob.first
 
@@ -111,7 +131,7 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
                 model.queue_background_migration_jobs_by_range_at_intervals(User, 'FooJob', 10.minutes, other_job_arguments: [1, 2])
               end.not_to change { Gitlab::Database::BackgroundMigrationJob.count }
 
-            expect(BackgroundMigrationWorker.jobs.size).to eq(1)
+              expect(worker_class.jobs.size).to eq(1)
             end
           end
         end
@@ -137,10 +157,10 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
             final_delay = model.queue_background_migration_jobs_by_range_at_intervals(ContainerExpirationPolicy, 'FooJob', 10.minutes, batch_size: 2, primary_column_name: :project_id)
 
             expect(final_delay.to_f).to eq(20.minutes.to_f)
-          expect(BackgroundMigrationWorker.jobs[0]['args']).to eq(['FooJob', [id1, id2]])
-          expect(BackgroundMigrationWorker.jobs[0]['at']).to eq(10.minutes.from_now.to_f)
-          expect(BackgroundMigrationWorker.jobs[1]['args']).to eq(['FooJob', [id3, id3]])
-          expect(BackgroundMigrationWorker.jobs[1]['at']).to eq(20.minutes.from_now.to_f)
+            expect(worker_class.jobs[0]['args']).to eq(['FooJob', [id1, id2]])
+            expect(worker_class.jobs[0]['at']).to eq(10.minutes.from_now.to_f)
+            expect(worker_class.jobs[1]['args']).to eq(['FooJob', [id3, id3]])
+            expect(worker_class.jobs[1]['at']).to eq(20.minutes.from_now.to_f)
           end
         end
 
@@ -202,10 +222,10 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
       it 'queues pending jobs' do
         subject
 
-      expect(BackgroundMigrationWorker.jobs[0]['args']).to eq([job_class_name, [1, 2]])
-      expect(BackgroundMigrationWorker.jobs[0]['at']).to be_nil
-      expect(BackgroundMigrationWorker.jobs[1]['args']).to eq([job_class_name, [3, 4]])
-      expect(BackgroundMigrationWorker.jobs[1]['at']).to eq(10.minutes.from_now.to_f)
+        expect(worker_class.jobs[0]['args']).to eq([job_class_name, [1, 2]])
+        expect(worker_class.jobs[0]['at']).to be_nil
+        expect(worker_class.jobs[1]['args']).to eq([job_class_name, [3, 4]])
+        expect(worker_class.jobs[1]['at']).to eq(10.minutes.from_now.to_f)
       end
 
       context 'with batch_size option' do
@@ -218,10 +238,10 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
         it 'queues pending jobs' do
           subject
 
-        expect(BackgroundMigrationWorker.jobs[0]['args']).to eq([job_class_name, [1, 2]])
-        expect(BackgroundMigrationWorker.jobs[0]['at']).to be_nil
-        expect(BackgroundMigrationWorker.jobs[1]['args']).to eq([job_class_name, [3, 4]])
-        expect(BackgroundMigrationWorker.jobs[1]['at']).to eq(10.minutes.from_now.to_f)
+          expect(worker_class.jobs[0]['args']).to eq([job_class_name, [1, 2]])
+          expect(worker_class.jobs[0]['at']).to be_nil
+          expect(worker_class.jobs[1]['args']).to eq([job_class_name, [3, 4]])
+          expect(worker_class.jobs[1]['at']).to eq(10.minutes.from_now.to_f)
         end
 
         it 'retrieve jobs in batches' do
@@ -246,10 +266,10 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
         it 'queues pending jobs' do
           subject
 
-        expect(BackgroundMigrationWorker.jobs[0]['args']).to eq([job_class_name, [1, 2]])
-        expect(BackgroundMigrationWorker.jobs[0]['at']).to eq(3.minutes.from_now.to_f)
-        expect(BackgroundMigrationWorker.jobs[1]['args']).to eq([job_class_name, [3, 4]])
-        expect(BackgroundMigrationWorker.jobs[1]['at']).to eq(13.minutes.from_now.to_f)
+          expect(worker_class.jobs[0]['args']).to eq([job_class_name, [1, 2]])
+          expect(worker_class.jobs[0]['at']).to eq(3.minutes.from_now.to_f)
+          expect(worker_class.jobs[1]['args']).to eq([job_class_name, [3, 4]])
+          expect(worker_class.jobs[1]['at']).to eq(13.minutes.from_now.to_f)
         end
 
         context 'when nothing is queued' do
@@ -262,79 +282,135 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
       end
     end
 
-  describe '#perform_background_migration_inline?' do
-    it 'returns true in a test environment' do
-      stub_rails_env('test')
+    describe '#finalized_background_migration' do
+      let(:coordinator) { Gitlab::BackgroundMigration::JobCoordinator.new(worker_class) }
 
-      expect(model.perform_background_migration_inline?).to eq(true)
+      let!(:tracked_pending_job) { create(:background_migration_job, class_name: job_class_name, status: :pending, arguments: [1]) }
+      let!(:tracked_successful_job) { create(:background_migration_job, class_name: job_class_name, status: :succeeded, arguments: [2]) }
+      let!(:job_class_name) { 'TestJob' }
+
+      let!(:job_class) do
+        Class.new do
+          def perform(*arguments)
+            Gitlab::Database::BackgroundMigrationJob.mark_all_as_succeeded('TestJob', arguments)
+          end
+        end
       end
 
-    it 'returns true in a development environment' do
-      stub_rails_env('development')
+      before do
+        allow(Gitlab::BackgroundMigration).to receive(:coordinator_for_database)
+          .with('main').and_return(coordinator)
 
-      expect(model.perform_background_migration_inline?).to eq(true)
-    end
+        expect(coordinator).to receive(:migration_class_for)
+          .with(job_class_name).at_least(:once) { job_class }
 
-    it 'returns false in a production environment' do
-      stub_rails_env('production')
+        Sidekiq::Testing.disable! do
+          worker_class.perform_async(job_class_name, [1, 2])
+          worker_class.perform_async(job_class_name, [3, 4])
+          worker_class.perform_in(10, job_class_name, [5, 6])
+          worker_class.perform_in(20, job_class_name, [7, 8])
+        end
+      end
 
-      expect(model.perform_background_migration_inline?).to eq(false)
+      it_behaves_like 'finalized tracked background migration', worker_class do
+        before do
+          model.finalize_background_migration(job_class_name)
         end
       end
 
-  describe '#migrate_async' do
-    it 'calls BackgroundMigrationWorker.perform_async' do
-      expect(BackgroundMigrationWorker).to receive(:perform_async).with("Class", "hello", "world")
+      context 'when removing all tracked job records' do
+        let!(:job_class) do
+          Class.new do
+            def perform(*arguments)
+              # Force pending jobs to remain pending
+            end
+          end
+        end
 
-      model.migrate_async("Class", "hello", "world")
+        before do
+          model.finalize_background_migration(job_class_name, delete_tracking_jobs: %w[pending succeeded])
         end
 
-    it 'pushes a context with the current class name as caller_id' do
-      expect(Gitlab::ApplicationContext).to receive(:with_context).with(caller_id: model.class.to_s)
+        it_behaves_like 'finalized tracked background migration', worker_class
+        it_behaves_like 'removed tracked jobs', 'pending'
+        it_behaves_like 'removed tracked jobs', 'succeeded'
+      end
 
-      model.migrate_async('Class', 'hello', 'world')
+      context 'when retaining all tracked job records' do
+        before do
+          model.finalize_background_migration(job_class_name, delete_tracking_jobs: false)
         end
+
+        it_behaves_like 'finalized background migration', worker_class
+        include_examples 'retained tracked jobs', 'succeeded'
       end
 
-  describe '#migrate_in' do
-    it 'calls BackgroundMigrationWorker.perform_in' do
-      expect(BackgroundMigrationWorker).to receive(:perform_in).with(10.minutes, 'Class', 'Hello', 'World')
+      context 'during retry race condition' do
+        let!(:job_class) do
+          Class.new do
+            class << self
+              attr_accessor :worker_class
 
-      model.migrate_in(10.minutes, 'Class', 'Hello', 'World')
+              def queue_items_added
+                @queue_items_added ||= []
+              end
             end
 
-    it 'pushes a context with the current class name as caller_id' do
-      expect(Gitlab::ApplicationContext).to receive(:with_context).with(caller_id: model.class.to_s)
-
-      model.migrate_in(10.minutes, 'Class', 'Hello', 'World')
+            def worker_class
+              self.class.worker_class
             end
+
+            def queue_items_added
+              self.class.queue_items_added
             end
 
-  describe '#bulk_migrate_async' do
-    it 'calls BackgroundMigrationWorker.bulk_perform_async' do
-      expect(BackgroundMigrationWorker).to receive(:bulk_perform_async).with([%w(Class hello world)])
+            def perform(*arguments)
+              Gitlab::Database::BackgroundMigrationJob.mark_all_as_succeeded('TestJob', arguments)
 
-      model.bulk_migrate_async([%w(Class hello world)])
+              # Mock another process pushing queue jobs.
+              if self.class.queue_items_added.count < 10
+                Sidekiq::Testing.disable! do
+                  queue_items_added << worker_class.perform_async('TestJob', [Time.current])
+                  queue_items_added << worker_class.perform_in(10, 'TestJob', [Time.current])
+                end
+              end
+            end
+          end
         end
 
-    it 'pushes a context with the current class name as caller_id' do
-      expect(Gitlab::ApplicationContext).to receive(:with_context).with(caller_id: model.class.to_s)
+        it_behaves_like 'finalized tracked background migration', worker_class do
+          before do
+            # deliberately set the worker class on our test job since it won't be pulled from the surrounding scope
+            job_class.worker_class = worker_class
 
-      model.bulk_migrate_async([%w(Class hello world)])
+            model.finalize_background_migration(job_class_name, delete_tracking_jobs: ['succeeded'])
+          end
+        end
       end
     end
 
-  describe '#bulk_migrate_in' do
-    it 'calls BackgroundMigrationWorker.bulk_perform_in_' do
-      expect(BackgroundMigrationWorker).to receive(:bulk_perform_in).with(10.minutes, [%w(Class hello world)])
+    describe '#migrate_in' do
+      it 'calls perform_in for the correct worker' do
+        expect(worker_class).to receive(:perform_in).with(10.minutes, 'Class', 'Hello', 'World')
 
-      model.bulk_migrate_in(10.minutes, [%w(Class hello world)])
+        model.migrate_in(10.minutes, 'Class', 'Hello', 'World')
       end
 
       it 'pushes a context with the current class name as caller_id' do
         expect(Gitlab::ApplicationContext).to receive(:with_context).with(caller_id: model.class.to_s)
 
-      model.bulk_migrate_in(10.minutes, [%w(Class hello world)])
+        model.migrate_in(10.minutes, 'Class', 'Hello', 'World')
+      end
+
+      context 'when a specific coordinator is given' do
+        let(:coordinator) { Gitlab::BackgroundMigration::JobCoordinator.for_tracking_database('main') }
+
+        it 'uses that coordinator' do
+          expect(coordinator).to receive(:perform_in).with(10.minutes, 'Class', 'Hello', 'World').and_call_original
+          expect(worker_class).to receive(:perform_in).with(10.minutes, 'Class', 'Hello', 'World')
+
+          model.migrate_in(10.minutes, 'Class', 'Hello', 'World', coordinator: coordinator)
+        end
       end
     end
 
@@ -343,10 +419,12 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
       let(:job2) { double }
 
       it 'deletes all queued jobs for the given background migration' do
-      expect(Gitlab::BackgroundMigration).to receive(:steal).with('BackgroundMigrationClassName') do |&block|
+        expect_next_instance_of(Gitlab::BackgroundMigration::JobCoordinator) do |coordinator|
+          expect(coordinator).to receive(:steal).with('BackgroundMigrationClassName') do |&block|
             expect(block.call(job1)).to be(false)
             expect(block.call(job2)).to be(false)
           end
+        end
 
         expect(job1).to receive(:delete)
         expect(job2).to receive(:delete)
@@ -354,97 +432,29 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
         model.delete_queued_jobs('BackgroundMigrationClassName')
       end
     end
-
-  describe '#finalized_background_migration' do
-    let(:job_coordinator) { Gitlab::BackgroundMigration::JobCoordinator.new(BackgroundMigrationWorker) }
-
-    let!(:job_class_name) { 'TestJob' }
-    let!(:job_class) { Class.new }
-    let!(:job_perform_method) do
-      ->(*arguments) do
-        Gitlab::Database::BackgroundMigrationJob.mark_all_as_succeeded(
-          # Value is 'TestJob' defined by :job_class_name in the let! above.
-          # Scoping prohibits us from directly referencing job_class_name.
-          RSpec.current_example.example_group_instance.job_class_name,
-          arguments
-        )
-      end
-    end
-
-    let!(:tracked_pending_job) { create(:background_migration_job, class_name: job_class_name, status: :pending, arguments: [1]) }
-    let!(:tracked_successful_job) { create(:background_migration_job, class_name: job_class_name, status: :succeeded, arguments: [2]) }
-
-    before do
-      job_class.define_method(:perform, job_perform_method)
-
-      allow(Gitlab::BackgroundMigration).to receive(:coordinator_for_database)
-        .with('main').and_return(job_coordinator)
-
-      expect(job_coordinator).to receive(:migration_class_for)
-        .with(job_class_name).at_least(:once) { job_class }
-
-      Sidekiq::Testing.disable! do
-        BackgroundMigrationWorker.perform_async(job_class_name, [1, 2])
-        BackgroundMigrationWorker.perform_async(job_class_name, [3, 4])
-        BackgroundMigrationWorker.perform_in(10, job_class_name, [5, 6])
-        BackgroundMigrationWorker.perform_in(20, job_class_name, [7, 8])
-      end
   end
 
-    it_behaves_like 'finalized tracked background migration' do
-      before do
-        model.finalize_background_migration(job_class_name)
-      end
+  context 'when the migration is running against the main database' do
+    it_behaves_like 'helpers that enqueue background migrations', BackgroundMigrationWorker, 'main'
   end
 
-    context 'when removing all tracked job records' do
-      # Force pending jobs to remain pending.
-      let!(:job_perform_method) { ->(*arguments) { } }
-
-      before do
-        model.finalize_background_migration(job_class_name, delete_tracking_jobs: %w[pending succeeded])
-      end
+  describe '#perform_background_migration_inline?' do
+    it 'returns true in a test environment' do
+      stub_rails_env('test')
 
-      it_behaves_like 'finalized tracked background migration'
-      it_behaves_like 'removed tracked jobs', 'pending'
-      it_behaves_like 'removed tracked jobs', 'succeeded'
+      expect(model.perform_background_migration_inline?).to eq(true)
     end
 
-    context 'when retaining all tracked job records' do
-      before do
-        model.finalize_background_migration(job_class_name, delete_tracking_jobs: false)
-      end
+    it 'returns true in a development environment' do
+      stub_rails_env('development')
 
-      it_behaves_like 'finalized background migration'
-      include_examples 'retained tracked jobs', 'succeeded'
+      expect(model.perform_background_migration_inline?).to eq(true)
     end
 
-    context 'during retry race condition' do
-      let(:queue_items_added) { [] }
-      let!(:job_perform_method) do
-        ->(*arguments) do
-          Gitlab::Database::BackgroundMigrationJob.mark_all_as_succeeded(
-            RSpec.current_example.example_group_instance.job_class_name,
-            arguments
-          )
-
-          # Mock another process pushing queue jobs.
-          queue_items_added = RSpec.current_example.example_group_instance.queue_items_added
-          if queue_items_added.count < 10
-            Sidekiq::Testing.disable! do
-              job_class_name = RSpec.current_example.example_group_instance.job_class_name
-              queue_items_added << BackgroundMigrationWorker.perform_async(job_class_name, [Time.current])
-              queue_items_added << BackgroundMigrationWorker.perform_in(10, job_class_name, [Time.current])
-            end
-          end
-        end
-      end
+    it 'returns false in a production environment' do
+      stub_rails_env('production')
 
-      it_behaves_like 'finalized tracked background migration' do
-        before do
-          model.finalize_background_migration(job_class_name, delete_tracking_jobs: ['succeeded'])
-        end
-      end
+      expect(model.perform_background_migration_inline?).to eq(false)
     end
   end
 

spec/support/shared_examples/lib/gitlab/database/background_migration_job_shared_examples.rb

@@ -22,19 +22,19 @@ RSpec.shared_examples 'marks background migration job records' do
   end
 end
 
-RSpec.shared_examples 'finalized background migration' do
+RSpec.shared_examples 'finalized background migration' do |worker_class|
   it 'processed the scheduled sidekiq queue' do
     queued = Sidekiq::ScheduledSet
       .new
       .select do |scheduled|
-        scheduled.klass == 'BackgroundMigrationWorker' &&
+        scheduled.klass == worker_class.name &&
         scheduled.args.first == job_class_name
       end
     expect(queued.size).to eq(0)
   end
 
   it 'processed the async sidekiq queue' do
-    queued = Sidekiq::Queue.new('BackgroundMigrationWorker')
+    queued = Sidekiq::Queue.new(worker_class.name)
       .select { |scheduled| scheduled.klass == job_class_name }
     expect(queued.size).to eq(0)
   end
@@ -42,8 +42,8 @@ RSpec.shared_examples 'finalized background migration' do
   include_examples 'removed tracked jobs', 'pending'
 end
 
-RSpec.shared_examples 'finalized tracked background migration' do
-  include_examples 'finalized background migration'
+RSpec.shared_examples 'finalized tracked background migration' do |worker_class|
+  include_examples 'finalized background migration', worker_class
   include_examples 'removed tracked jobs', 'succeeded'
 end