Use migration database to enqueue bg migrations

Update the background migration helpers to enqueue background jobs
for the target tracking database. For now we default to using main for
all migrations, since we don't have additional databases yet.

This changes not only enqueuing operations, but other related functions
like stealing scheduled jobs from the queue.

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