Consume remaining MigrateApproverToApprovalRulesInBatch

Allow `steal` to handle dead jobs.

2 files changed, 43 insertions, 4 deletions

diff --git a/lib/gitlab/background_migration.rb b/lib/gitlab/background_migration.rb
index 6cf40e2d4ca..5251e0fadf9 100644
--- a/lib/gitlab/background_migration.rb
+++ b/lib/gitlab/background_migration.rb
@@ -16,11 +16,18 @@ module Gitlab
     # re-raises the exception.
     #
     # steal_class - The name of the class for which to steal jobs.
-    def self.steal(steal_class)
-      enqueued = Sidekiq::Queue.new(self.queue)
-      scheduled = Sidekiq::ScheduledSet.new
+    def self.steal(steal_class, retry_dead_jobs: false)
+      queues = [
+        Sidekiq::ScheduledSet.new,
+        Sidekiq::Queue.new(self.queue)
+      ]
+
+      if retry_dead_jobs
+        queues << Sidekiq::RetrySet.new
+        queues << Sidekiq::DeadSet.new
+      end
 
-      [scheduled, enqueued].each do |queue|
+      queues.each do |queue|
         queue.each do |job|
           migration_class, migration_args = job.args
 
diff --git a/spec/lib/gitlab/background_migration_spec.rb b/spec/lib/gitlab/background_migration_spec.rb
index 8a83b76fd94..7d3d8a949ef 100644
--- a/spec/lib/gitlab/background_migration_spec.rb
+++ b/spec/lib/gitlab/background_migration_spec.rb
@@ -104,6 +104,38 @@ describe Gitlab::BackgroundMigration do
         end
       end
     end
+
+    context 'when retry_dead_jobs is true', :sidekiq, :redis do
+      let(:retry_queue) do
+        [double(args: ['Object', [3]], queue: described_class.queue, delete: true)]
+      end
+      let(:dead_queue) do
+        [double(args: ['Object', [4]], queue: described_class.queue, delete: true)]
+      end
+
+      before do
+        allow(Sidekiq::RetrySet).to receive(:new).and_return(retry_queue)
+        allow(Sidekiq::DeadSet).to receive(:new).and_return(dead_queue)
+      end
+
+      it 'steals from the dead and retry queue' do
+        Sidekiq::Testing.disable! do
+          expect(described_class).to receive(:perform)
+            .with('Object', [1]).ordered
+          expect(described_class).to receive(:perform)
+            .with('Object', [2]).ordered
+          expect(described_class).to receive(:perform)
+            .with('Object', [3]).ordered
+          expect(described_class).to receive(:perform)
+            .with('Object', [4]).ordered
+
+          BackgroundMigrationWorker.perform_async('Object', [2])
+          BackgroundMigrationWorker.perform_in(10.minutes, 'Object', [1])
+
+          described_class.steal('Object', retry_dead_jobs: true)
+        end
+      end
+    end
   end
 
   describe '.perform' do