Background migration cleanup helper

A helper to ensure that sidekiq and pending tracked jobs are processed.
Optionally delete successful tracked job rows.

doc/development/background_migrations.md

@@ -174,13 +174,18 @@ roughly be as follows:
 1. Release B:
    1. Deploy code so that the application starts using the new column and stops
       scheduling jobs for newly created data.
-   1. In a post-deployment migration you'll need to ensure no jobs remain.
+   1. In a post-deployment migration use `finalize_background_migration` from
+      `BackgroundMigrationHelpers` to ensure no jobs remain. This helper will:
          1. Use `Gitlab::BackgroundMigration.steal` to process any remaining
             jobs in Sidekiq.
          1. Reschedule the migration to be run directly (i.e. not through Sidekiq)
             on any rows that weren't migrated by Sidekiq. This can happen if, for
             instance, Sidekiq received a SIGKILL, or if a particular batch failed
             enough times to be marked as dead.
+         1. Remove `Gitlab::Database::BackgroundMigrationJob` rows where
+            `status = succeeded`. To retain diagnostic information that may
+            help with future bug tracking you can skip this step by specifying
+            the `delete_tracking_jobs: false` parameter.
    1. Remove the old column.
 
 This may also require a bump to the [import/export version](../user/project/settings/import_export.md), if

lib/gitlab/database/migrations/background_migration_helpers.rb

@@ -264,6 +264,34 @@ module Gitlab
           migration
         end
 
+        # Force a background migration to complete.
+        #
+        # WARNING: This method will block the caller and move the background migration from an
+        # asynchronous migration to a synchronous migration.
+        #
+        # 1. Steal work from sidekiq and perform immediately (avoid duplicates generated by step 2).
+        # 2. Process any pending tracked jobs.
+        # 3. Steal work from sidekiq and perform immediately (clear anything left from step 2).
+        # 4. Optionally remove job tracking information.
+        #
+        # This method does not garauntee that all jobs completed successfully.
+        def finalize_background_migration(class_name, delete_tracking_jobs: ['succeeded'])
+          # Empty the sidekiq queue.
+          Gitlab::BackgroundMigration.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)
+          end
+
+          # Empty the sidekiq queue.
+          Gitlab::BackgroundMigration.steal(class_name)
+
+          # Delete job tracking rows.
+          delete_job_tracking(class_name, status: delete_tracking_jobs) if delete_tracking_jobs
+        end
+
         def perform_background_migration_inline?
           Rails.env.test? || Rails.env.development?
         end
@@ -304,6 +332,12 @@ module Gitlab
           end
         end
 
+        def delete_job_tracking(class_name, status: 'succeeded')
+          status = Array(status).map { |s| Gitlab::Database::BackgroundMigrationJob.statuses[s] }
+          jobs = Gitlab::Database::BackgroundMigrationJob.where(status: status).for_migration_class(class_name)
+          jobs.each_batch { |batch| batch.delete_all }
+        end
+
         private
 
         def track_in_database(class_name, arguments)

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

@@ -581,4 +581,101 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
       model.delete_queued_jobs('BackgroundMigrationClassName')
     end
   end
+
+  describe '#finalized_background_migration' do
+    include_context 'background migration job class'
+
+    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
+      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
+    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
+
+      it_behaves_like 'finalized tracked background migration'
+      it_behaves_like 'removed tracked jobs', 'pending'
+      it_behaves_like 'removed tracked jobs', 'succeeded'
+    end
+
+    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'
+      include_examples 'retained tracked jobs', 'succeeded'
+    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_behaves_like 'finalized tracked background migration' do
+        before do
+          model.finalize_background_migration(job_class_name, delete_tracking_jobs: ['succeeded'])
+        end
+      end
+    end
+  end
+
+  describe '#delete_job_tracking' do
+    let!(:job_class_name) { 'TestJob' }
+
+    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]) }
+
+    context 'with default status' do
+      before do
+        model.delete_job_tracking(job_class_name)
+      end
+
+      include_examples 'retained tracked jobs', 'pending'
+      include_examples 'removed tracked jobs', 'succeeded'
+    end
+
+    context 'with explicit status' do
+      before do
+        model.delete_job_tracking(job_class_name, status: %w[pending succeeded])
+      end
+
+      include_examples 'removed tracked jobs', 'pending'
+      include_examples 'removed tracked jobs', 'succeeded'
+    end
+  end
 end

spec/support/shared_contexts/lib/gitlab/database/background_migration_job_shared_context.rb

+# frozen_string_literal: true
+
+RSpec.shared_context 'background migration job class' do
+  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
+
+  before do
+    job_class.define_method(:perform, job_perform_method)
+    expect(Gitlab::BackgroundMigration).to receive(:migration_class_for).with(job_class_name).at_least(:once) { job_class }
+  end
+end

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

@@ -21,3 +21,46 @@ RSpec.shared_examples 'marks background migration job records' do
     expect(jobs_updated).to eq(1)
   end
 end
+
+RSpec.shared_examples 'finalized background migration' do
+  it 'processed the scheduled sidekiq queue' do
+    queued = Sidekiq::ScheduledSet
+      .new
+      .select do |scheduled|
+        scheduled.klass == 'BackgroundMigrationWorker' &&
+        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')
+      .select { |scheduled| scheduled.klass == job_class_name }
+    expect(queued.size).to eq(0)
+  end
+
+  include_examples 'removed tracked jobs', 'pending'
+end
+
+RSpec.shared_examples 'finalized tracked background migration' do
+  include_examples 'finalized background migration'
+  include_examples 'removed tracked jobs', 'succeeded'
+end
+
+RSpec.shared_examples 'removed tracked jobs' do |status|
+  it "removes '#{status}' tracked jobs" do
+    jobs = Gitlab::Database::BackgroundMigrationJob
+      .where(status: Gitlab::Database::BackgroundMigrationJob.statuses[status])
+      .for_migration_class(job_class_name)
+    expect(jobs).to be_empty
+  end
+end
+
+RSpec.shared_examples 'retained tracked jobs' do |status|
+  it "retains '#{status}' tracked jobs" do
+    jobs = Gitlab::Database::BackgroundMigrationJob
+      .where(status: Gitlab::Database::BackgroundMigrationJob.statuses[status])
+      .for_migration_class(job_class_name)
+    expect(jobs).to be_present
+  end
+end