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)
+            coordinator.perform_in(*args)
           end
         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,444 +7,454 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
     ActiveRecord::Migration.new.extend(described_class)
   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 }
-      let!(:id2) { create(:user).id }
-      let!(:id3) { create(:user).id }
+  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
 
-      around do |example|
-        freeze_time { example.run }
-      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 }
+        let!(:id2) { create(:user).id }
+        let!(:id3) { create(:user).id }
 
-      before do
-        User.class_eval do
-          include EachBatch
+        around do |example|
+          freeze_time { example.run }
         end
-      end
-
-      it 'returns the final expected delay' do
-        Sidekiq::Testing.fake! do
-          final_delay = model.queue_background_migration_jobs_by_range_at_intervals(User, 'FooJob', 10.minutes, batch_size: 2)
 
-          expect(final_delay.to_f).to eq(20.minutes.to_f)
+        before do
+          User.class_eval do
+            include EachBatch
+          end
         end
-      end
 
-      it 'returns zero when nothing gets queued' do
-        Sidekiq::Testing.fake! do
-          final_delay = model.queue_background_migration_jobs_by_range_at_intervals(User.none, 'FooJob', 10.minutes)
+        it 'returns the final expected delay' do
+          Sidekiq::Testing.fake! do
+            final_delay = model.queue_background_migration_jobs_by_range_at_intervals(User, 'FooJob', 10.minutes, batch_size: 2)
 
-          expect(final_delay).to eq(0)
+            expect(final_delay.to_f).to eq(20.minutes.to_f)
+          end
         end
-      end
 
-      context 'with batch_size option' do
-        it 'queues jobs correctly' do
+        it 'returns zero when nothing gets queued' do
           Sidekiq::Testing.fake! do
-            model.queue_background_migration_jobs_by_range_at_intervals(User, 'FooJob', 10.minutes, batch_size: 2)
+            final_delay = model.queue_background_migration_jobs_by_range_at_intervals(User.none, 'FooJob', 10.minutes)
 
-            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(final_delay).to eq(0)
           end
         end
-      end
 
-      context 'without 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)
+        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(BackgroundMigrationWorker.jobs[0]['args']).to eq(['FooJob', [id1, id3]])
-            expect(BackgroundMigrationWorker.jobs[0]['at']).to eq(10.minutes.from_now.to_f)
+              expect(final_delay.to_f).to eq(4.minutes.to_f)
+            end
           end
         end
-      end
 
-      context 'with other_job_arguments option' do
-        it 'queues jobs correctly' do
-          Sidekiq::Testing.fake! do
-            model.queue_background_migration_jobs_by_range_at_intervals(User, 'FooJob', 10.minutes, other_job_arguments: [1, 2])
+        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, 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, 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
-      end
 
-      context 'with initial_delay option' do
-        it 'queues jobs correctly' 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)
+        context 'without 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)
 
-            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]])
+              expect(worker_class.jobs[0]['at']).to eq(10.minutes.from_now.to_f)
+            end
           end
         end
-      end
 
-      context 'with track_jobs option' do
-        it 'creates a record for each job in the database' do
-          Sidekiq::Testing.fake! do
-            expect do
-              model.queue_background_migration_jobs_by_range_at_intervals(User, '::FooJob', 10.minutes,
-                other_job_arguments: [1, 2], track_jobs: true)
-            end.to change { Gitlab::Database::BackgroundMigrationJob.count }.from(0).to(1)
+        context 'with other_job_arguments option' do
+          it 'queues jobs correctly' 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.size).to eq(1)
+              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
 
-            tracked_job = Gitlab::Database::BackgroundMigrationJob.first
+        context 'with initial_delay option' do
+          it 'queues jobs correctly' 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(tracked_job.class_name).to eq('FooJob')
-            expect(tracked_job.arguments).to eq([id1, id3, 1, 2])
-            expect(tracked_job).to be_pending
+              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
-      end
 
-      context 'without track_jobs option' do
-        it 'does not create records in the database' do
-          Sidekiq::Testing.fake! do
-            expect 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 }
+        context 'with track_jobs option' do
+          it 'creates a record for each job in the database' do
+            Sidekiq::Testing.fake! do
+              expect do
+                model.queue_background_migration_jobs_by_range_at_intervals(User, '::FooJob', 10.minutes,
+                  other_job_arguments: [1, 2], track_jobs: true)
+              end.to change { Gitlab::Database::BackgroundMigrationJob.count }.from(0).to(1)
+
+              expect(worker_class.jobs.size).to eq(1)
+
+              tracked_job = Gitlab::Database::BackgroundMigrationJob.first
 
-            expect(BackgroundMigrationWorker.jobs.size).to eq(1)
+              expect(tracked_job.class_name).to eq('FooJob')
+              expect(tracked_job.arguments).to eq([id1, id3, 1, 2])
+              expect(tracked_job).to be_pending
+            end
           end
         end
-      end
-    end
 
-    context 'when the model specifies a primary_column_name' do
-      let!(:id1) { create(:container_expiration_policy).id }
-      let!(:id2) { create(:container_expiration_policy).id }
-      let!(:id3) { create(:container_expiration_policy).id }
+        context 'without track_jobs option' do
+          it 'does not create records in the database' do
+            Sidekiq::Testing.fake! do
+              expect 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 }
 
-      around do |example|
-        freeze_time { example.run }
+              expect(worker_class.jobs.size).to eq(1)
+            end
+          end
+        end
       end
 
-      before do
-        ContainerExpirationPolicy.class_eval do
-          include EachBatch
+      context 'when the model specifies a primary_column_name' do
+        let!(:id1) { create(:container_expiration_policy).id }
+        let!(:id2) { create(:container_expiration_policy).id }
+        let!(:id3) { create(:container_expiration_policy).id }
+
+        around do |example|
+          freeze_time { example.run }
+        end
+
+        before do
+          ContainerExpirationPolicy.class_eval do
+            include EachBatch
+          end
         end
-      end
 
-      it 'returns the final expected delay', :aggregate_failures do
-        Sidekiq::Testing.fake! do
-          final_delay = model.queue_background_migration_jobs_by_range_at_intervals(ContainerExpirationPolicy, 'FooJob', 10.minutes, batch_size: 2, primary_column_name: :project_id)
+        it 'returns the final expected delay', :aggregate_failures do
+          Sidekiq::Testing.fake! 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(final_delay.to_f).to eq(20.minutes.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
 
-      context "when the primary_column_name is not an integer" do
-        it 'raises error' do
-          expect do
-            model.queue_background_migration_jobs_by_range_at_intervals(ContainerExpirationPolicy, 'FooJob', 10.minutes, primary_column_name: :enabled)
-          end.to raise_error(StandardError, /is not an integer column/)
+        context "when the primary_column_name is not an integer" do
+          it 'raises error' do
+            expect do
+              model.queue_background_migration_jobs_by_range_at_intervals(ContainerExpirationPolicy, 'FooJob', 10.minutes, primary_column_name: :enabled)
+            end.to raise_error(StandardError, /is not an integer column/)
+          end
         end
-      end
 
-      context "when the primary_column_name does not exist" do
-        it 'raises error' do
-          expect do
-            model.queue_background_migration_jobs_by_range_at_intervals(ContainerExpirationPolicy, 'FooJob', 10.minutes, primary_column_name: :foo)
-          end.to raise_error(StandardError, /does not have an ID column of foo/)
+        context "when the primary_column_name does not exist" do
+          it 'raises error' do
+            expect do
+              model.queue_background_migration_jobs_by_range_at_intervals(ContainerExpirationPolicy, 'FooJob', 10.minutes, primary_column_name: :foo)
+            end.to raise_error(StandardError, /does not have an ID column of foo/)
+          end
         end
       end
-    end
 
-    context "when the model doesn't have an ID or primary_column_name column" do
-      it 'raises error (for now)' do
-        expect do
-          model.queue_background_migration_jobs_by_range_at_intervals(ProjectAuthorization, 'FooJob', 10.seconds)
-        end.to raise_error(StandardError, /does not have an ID/)
+      context "when the model doesn't have an ID or primary_column_name column" do
+        it 'raises error (for now)' do
+          expect do
+            model.queue_background_migration_jobs_by_range_at_intervals(ProjectAuthorization, 'FooJob', 10.seconds)
+          end.to raise_error(StandardError, /does not have an ID/)
+        end
       end
     end
-  end
 
-  describe '#requeue_background_migration_jobs_by_range_at_intervals' do
-    let!(:job_class_name) { 'TestJob' }
-    let!(:pending_job_1) { create(:background_migration_job, class_name: job_class_name, status: :pending, arguments: [1, 2]) }
-    let!(:pending_job_2) { create(:background_migration_job, class_name: job_class_name, status: :pending, arguments: [3, 4]) }
-    let!(:successful_job_1) { create(:background_migration_job, class_name: job_class_name, status: :succeeded, arguments: [5, 6]) }
-    let!(:successful_job_2) { create(:background_migration_job, class_name: job_class_name, status: :succeeded, arguments: [7, 8]) }
+    describe '#requeue_background_migration_jobs_by_range_at_intervals' do
+      let!(:job_class_name) { 'TestJob' }
+      let!(:pending_job_1) { create(:background_migration_job, class_name: job_class_name, status: :pending, arguments: [1, 2]) }
+      let!(:pending_job_2) { create(:background_migration_job, class_name: job_class_name, status: :pending, arguments: [3, 4]) }
+      let!(:successful_job_1) { create(:background_migration_job, class_name: job_class_name, status: :succeeded, arguments: [5, 6]) }
+      let!(:successful_job_2) { create(:background_migration_job, class_name: job_class_name, status: :succeeded, arguments: [7, 8]) }
 
-    around do |example|
-      freeze_time do
-        Sidekiq::Testing.fake! do
-          example.run
+      around do |example|
+        freeze_time do
+          Sidekiq::Testing.fake! do
+            example.run
+          end
         end
       end
-    end
-
-    subject { model.requeue_background_migration_jobs_by_range_at_intervals(job_class_name, 10.minutes) }
-
-    it 'returns the expected duration' do
-      expect(subject).to eq(20.minutes)
-    end
 
-    context 'when nothing is queued' do
-      subject { model.requeue_background_migration_jobs_by_range_at_intervals('FakeJob', 10.minutes) }
+      subject { model.requeue_background_migration_jobs_by_range_at_intervals(job_class_name, 10.minutes) }
 
-      it 'returns expected duration of zero when nothing gets queued' do
-        expect(subject).to eq(0)
+      it 'returns the expected duration' do
+        expect(subject).to eq(20.minutes)
       end
-    end
-
-    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)
-    end
 
-    context 'with batch_size option' do
-      subject { model.requeue_background_migration_jobs_by_range_at_intervals(job_class_name, 10.minutes, batch_size: 1) }
+      context 'when nothing is queued' do
+        subject { model.requeue_background_migration_jobs_by_range_at_intervals('FakeJob', 10.minutes) }
 
-      it 'returns the expected duration' do
-        expect(subject).to eq(20.minutes)
+        it 'returns expected duration of zero when nothing gets queued' do
+          expect(subject).to eq(0)
+        end
       end
 
       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
-        jobs = double('jobs')
-        expect(Gitlab::Database::BackgroundMigrationJob).to receive(:pending) { jobs }
-        allow(jobs).to receive(:where).with(class_name: job_class_name) { jobs }
-        expect(jobs).to receive(:each_batch).with(of: 1)
-
-        subject
-      end
-    end
+      context 'with batch_size option' do
+        subject { model.requeue_background_migration_jobs_by_range_at_intervals(job_class_name, 10.minutes, batch_size: 1) }
 
-    context 'with initial_delay option' do
-      let_it_be(:initial_delay) { 3.minutes }
+        it 'returns the expected duration' do
+          expect(subject).to eq(20.minutes)
+        end
 
-      subject { model.requeue_background_migration_jobs_by_range_at_intervals(job_class_name, 10.minutes, initial_delay: initial_delay) }
+        it 'queues pending jobs' do
+          subject
 
-      it 'returns the expected duration' do
-        expect(subject).to eq(23.minutes)
-      end
+          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 'queues pending jobs' do
-        subject
+        it 'retrieve jobs in batches' do
+          jobs = double('jobs')
+          expect(Gitlab::Database::BackgroundMigrationJob).to receive(:pending) { jobs }
+          allow(jobs).to receive(:where).with(class_name: job_class_name) { jobs }
+          expect(jobs).to receive(:each_batch).with(of: 1)
 
-        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)
+          subject
+        end
       end
 
-      context 'when nothing is queued' do
-        subject { model.requeue_background_migration_jobs_by_range_at_intervals('FakeJob', 10.minutes) }
+      context 'with initial_delay option' do
+        let_it_be(:initial_delay) { 3.minutes }
 
-        it 'returns expected duration of zero when nothing gets queued' do
-          expect(subject).to eq(0)
+        subject { model.requeue_background_migration_jobs_by_range_at_intervals(job_class_name, 10.minutes, initial_delay: initial_delay) }
+
+        it 'returns the expected duration' do
+          expect(subject).to eq(23.minutes)
         end
-      end
-    end
-  end
 
-  describe '#perform_background_migration_inline?' do
-    it 'returns true in a test environment' do
-      stub_rails_env('test')
+        it 'queues pending jobs' do
+          subject
 
-      expect(model.perform_background_migration_inline?).to eq(true)
-    end
+          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
 
-    it 'returns true in a development environment' do
-      stub_rails_env('development')
+        context 'when nothing is queued' do
+          subject { model.requeue_background_migration_jobs_by_range_at_intervals('FakeJob', 10.minutes) }
 
-      expect(model.perform_background_migration_inline?).to eq(true)
+          it 'returns expected duration of zero when nothing gets queued' do
+            expect(subject).to eq(0)
+          end
+        end
+      end
     end
 
-    it 'returns false in a production environment' do
-      stub_rails_env('production')
+    describe '#finalized_background_migration' do
+      let(:coordinator) { Gitlab::BackgroundMigration::JobCoordinator.new(worker_class) }
 
-      expect(model.perform_background_migration_inline?).to eq(false)
-    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]) }
+      let!(:job_class_name) { 'TestJob' }
 
-  describe '#migrate_async' do
-    it 'calls BackgroundMigrationWorker.perform_async' do
-      expect(BackgroundMigrationWorker).to receive(:perform_async).with("Class", "hello", "world")
+      let!(:job_class) do
+        Class.new do
+          def perform(*arguments)
+            Gitlab::Database::BackgroundMigrationJob.mark_all_as_succeeded('TestJob', arguments)
+          end
+        end
+      end
 
-      model.migrate_async("Class", "hello", "world")
-    end
+      before do
+        allow(Gitlab::BackgroundMigration).to receive(:coordinator_for_database)
+          .with('main').and_return(coordinator)
 
-    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)
+        expect(coordinator).to receive(:migration_class_for)
+          .with(job_class_name).at_least(:once) { job_class }
 
-      model.migrate_async('Class', 'hello', 'world')
-    end
-  end
+        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
 
-  describe '#migrate_in' do
-    it 'calls BackgroundMigrationWorker.perform_in' do
-      expect(BackgroundMigrationWorker).to receive(:perform_in).with(10.minutes, 'Class', 'Hello', 'World')
+      it_behaves_like 'finalized tracked background migration', worker_class do
+        before do
+          model.finalize_background_migration(job_class_name)
+        end
+      end
 
-      model.migrate_in(10.minutes, 'Class', 'Hello', 'World')
-    end
+      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
 
-    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)
+        before do
+          model.finalize_background_migration(job_class_name, delete_tracking_jobs: %w[pending succeeded])
+        end
 
-      model.migrate_in(10.minutes, 'Class', 'Hello', 'World')
-    end
-  end
+        it_behaves_like 'finalized tracked background migration', worker_class
+        it_behaves_like 'removed tracked jobs', 'pending'
+        it_behaves_like 'removed tracked jobs', 'succeeded'
+      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)])
+      context 'when retaining all tracked job records' do
+        before do
+          model.finalize_background_migration(job_class_name, delete_tracking_jobs: false)
+        end
 
-      model.bulk_migrate_async([%w(Class hello world)])
-    end
+        it_behaves_like 'finalized background migration', worker_class
+        include_examples 'retained tracked jobs', '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)
+      context 'during retry race condition' do
+        let!(:job_class) do
+          Class.new do
+            class << self
+              attr_accessor :worker_class
 
-      model.bulk_migrate_async([%w(Class hello world)])
-    end
-  end
+              def queue_items_added
+                @queue_items_added ||= []
+              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)])
+            def worker_class
+              self.class.worker_class
+            end
 
-      model.bulk_migrate_in(10.minutes, [%w(Class hello world)])
-    end
+            def queue_items_added
+              self.class.queue_items_added
+            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)
+            def perform(*arguments)
+              Gitlab::Database::BackgroundMigrationJob.mark_all_as_succeeded('TestJob', arguments)
 
-      model.bulk_migrate_in(10.minutes, [%w(Class hello world)])
-    end
-  end
+              # 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
 
-  describe '#delete_queued_jobs' do
-    let(:job1) { double }
-    let(:job2) { double }
+        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
 
-    it 'deletes all queued jobs for the given background migration' do
-      expect(Gitlab::BackgroundMigration).to receive(:steal).with('BackgroundMigrationClassName') do |&block|
-        expect(block.call(job1)).to be(false)
-        expect(block.call(job2)).to be(false)
+            model.finalize_background_migration(job_class_name, delete_tracking_jobs: ['succeeded'])
+          end
+        end
       end
-
-      expect(job1).to receive(:delete)
-      expect(job2).to receive(:delete)
-
-      model.delete_queued_jobs('BackgroundMigrationClassName')
     end
-  end
 
-  describe '#finalized_background_migration' do
-    let(:job_coordinator) { Gitlab::BackgroundMigration::JobCoordinator.new(BackgroundMigrationWorker) }
+    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')
 
-    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
-        )
+        model.migrate_in(10.minutes, 'Class', 'Hello', 'World')
       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]) }
+      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)
 
-    before do
-      job_class.define_method(:perform, job_perform_method)
+        model.migrate_in(10.minutes, 'Class', 'Hello', 'World')
+      end
 
-      allow(Gitlab::BackgroundMigration).to receive(:coordinator_for_database)
-        .with('main').and_return(job_coordinator)
+      context 'when a specific coordinator is given' do
+        let(:coordinator) { Gitlab::BackgroundMigration::JobCoordinator.for_tracking_database('main') }
 
-      expect(job_coordinator).to receive(:migration_class_for)
-        .with(job_class_name).at_least(:once) { job_class }
+        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')
 
-      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])
+          model.migrate_in(10.minutes, 'Class', 'Hello', 'World', coordinator: coordinator)
+        end
       end
     end
 
-    it_behaves_like 'finalized tracked background migration' do
-      before do
-        model.finalize_background_migration(job_class_name)
+    describe '#delete_queued_jobs' do
+      let(:job1) { double }
+      let(:job2) { double }
+
+      it 'deletes all queued jobs for the given background migration' do
+        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)
+
+        model.delete_queued_jobs('BackgroundMigrationClassName')
       end
     end
+  end
 
-    context 'when removing all tracked job records' do
-      # Force pending jobs to remain pending.
-      let!(:job_perform_method) { ->(*arguments) { } }
+  context 'when the migration is running against the main database' do
+    it_behaves_like 'helpers that enqueue background migrations', BackgroundMigrationWorker, 'main'
+  end
 
-      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