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. 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. 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