Make Gitlab::BackgroundMigration.remaining work across queues

With worker routing, a Sidekiq queue can have jobs from multiple
workers. This meant `Gitlab::BackgroundMigration.remaining` needed to
account for that possibility. We do that in the same way as we already
did with the scheduled, retry, and dead sets: iterate through and look
for matching records.

We also fix other methods in this class (like `.steal`) while we're
doing this.

doc/update/index.md

@@ -86,35 +86,19 @@ To check the status of [batched background migrations](../user/admin_area/monito
 
 **For Omnibus installations**
 
-If using GitLab 12.9 and newer, also run:
+You can also run:
 
 ```shell
 sudo gitlab-rails runner -e production 'puts Gitlab::BackgroundMigration.remaining'
 ```
 
-If using GitLab 12.8 and older, run the following using a [Rails console](../administration/operations/rails_console.md#starting-a-rails-console-session):
-
-```ruby
-puts Sidekiq::Queue.new("background_migration").size
-Sidekiq::ScheduledSet.new.select { |r| r.klass == 'BackgroundMigrationWorker' }.size
-```
-
 **For installations from source**
 
-If using GitLab 12.9 and newer, run:
-
 ```shell
 cd /home/git/gitlab
 sudo -u git -H bundle exec rails runner -e production 'puts Gitlab::BackgroundMigration.remaining'
 ```
 
-If using GitLab 12.8 and older, run the following using a [Rails console](../administration/operations/rails_console.md#starting-a-rails-console-session):
-
-```ruby
-Sidekiq::Queue.new("background_migration").size
-Sidekiq::ScheduledSet.new.select { |r| r.klass == 'BackgroundMigrationWorker' }.size
-```
-
 ### What do I do if my background migrations are stuck?
 
 WARNING:

lib/gitlab/background_migration.rb

@@ -31,7 +31,7 @@ module Gitlab
         queue.each do |job|
           migration_class, migration_args = job.args
 
-          next unless job.queue == self.queue
+          next unless job.klass == 'BackgroundMigrationWorker'
           next unless migration_class == steal_class
           next if block_given? && !(yield job)
 
@@ -60,11 +60,14 @@ module Gitlab
     end
 
     def self.remaining
-      scheduled = Sidekiq::ScheduledSet.new.count do |job|
-        job.queue == self.queue
-      end
+      enqueued = Sidekiq::Queue.new(self.queue)
+      scheduled = Sidekiq::ScheduledSet.new
 
-      scheduled + Sidekiq::Queue.new(self.queue).size
+      [enqueued, scheduled].sum do |set|
+        set.count do |job|
+          job.klass == 'BackgroundMigrationWorker'
+        end
+      end
     end
 
     def self.exists?(migration_class, additional_queues = [])
@@ -105,13 +108,11 @@ module Gitlab
     end
 
     def self.enqueued_job?(queues, migration_class)
-      queues.each do |queue|
-        queue.each do |job|
-          return true if job.queue == self.queue && job.args.first == migration_class
+      queues.any? do |queue|
+        queue.any? do |job|
+          job.klass == 'BackgroundMigrationWorker' && job.args.first == migration_class
         end
       end
-
-      false
     end
   end
 end

spec/lib/gitlab/background_migration_spec.rb

@@ -13,7 +13,11 @@ RSpec.describe Gitlab::BackgroundMigration do
   describe '.steal' do
     context 'when there are enqueued jobs present' do
       let(:queue) do
-        [double(args: ['Foo', [10, 20]], queue: described_class.queue)]
+        [
+          double(args: ['Foo', [10, 20]], klass: 'BackgroundMigrationWorker'),
+          double(args: ['Bar', [20, 30]], klass: 'BackgroundMigrationWorker'),
+          double(args: ['Foo', [20, 30]], klass: 'MergeWorker')
+        ]
       end
 
       before do
@@ -45,7 +49,7 @@ RSpec.describe Gitlab::BackgroundMigration do
 
           expect(queue[0]).not_to receive(:delete)
 
-          described_class.steal('Bar')
+          described_class.steal('Baz')
         end
 
         context 'when a custom predicate is given' do
@@ -72,8 +76,8 @@ RSpec.describe Gitlab::BackgroundMigration do
         let(:migration) { spy(:migration) }
 
         let(:queue) do
-          [double(args: ['Foo', [10, 20]], queue: described_class.queue),
-           double(args: ['Foo', [20, 30]], queue: described_class.queue)]
+          [double(args: ['Foo', [10, 20]], klass: 'BackgroundMigrationWorker'),
+           double(args: ['Foo', [20, 30]], klass: 'BackgroundMigrationWorker')]
         end
 
         before do
@@ -128,11 +132,11 @@ RSpec.describe Gitlab::BackgroundMigration do
 
     context 'when retry_dead_jobs is true', :redis do
       let(:retry_queue) do
-        [double(args: ['Object', [3]], queue: described_class.queue, delete: true)]
+        [double(args: ['Object', [3]], klass: 'BackgroundMigrationWorker', delete: true)]
       end
 
       let(:dead_queue) do
-        [double(args: ['Object', [4]], queue: described_class.queue, delete: true)]
+        [double(args: ['Object', [4]], klass: 'BackgroundMigrationWorker', delete: true)]
       end
 
       before do
@@ -187,20 +191,22 @@ RSpec.describe Gitlab::BackgroundMigration do
 
   describe '.remaining', :redis do
     context 'when there are jobs remaining' do
-      let(:queue) { Array.new(12) }
-
       before do
-        allow(Sidekiq::Queue).to receive(:new)
-          .with(described_class.queue)
-          .and_return(Array.new(12))
-
         Sidekiq::Testing.disable! do
-          BackgroundMigrationWorker.perform_in(10.minutes, 'Foo')
+          MergeWorker.perform_async('Foo')
+          MergeWorker.perform_in(10.minutes, 'Foo')
+
+          5.times do
+            BackgroundMigrationWorker.perform_async('Foo')
+          end
+          3.times do
+            BackgroundMigrationWorker.perform_in(10.minutes, 'Foo')
+          end
         end
       end
 
       it 'returns the enqueued jobs plus the scheduled jobs' do
-        expect(described_class.remaining).to eq(13)
+        expect(described_class.remaining).to eq(8)
       end
     end
 
@@ -211,16 +217,13 @@ RSpec.describe Gitlab::BackgroundMigration do
     end
   end
 
-  describe '.exists?' do
+  describe '.exists?', :redis do
     context 'when there are enqueued jobs present' do
-      let(:queue) do
-        [double(args: ['Foo', [10, 20]], queue: described_class.queue)]
-      end
-
       before do
-        allow(Sidekiq::Queue).to receive(:new)
-          .with(described_class.queue)
-          .and_return(queue)
+        Sidekiq::Testing.disable! do
+          MergeWorker.perform_async('Bar')
+          BackgroundMigrationWorker.perform_async('Foo')
+        end
       end
 
       it 'returns true if specific job exists' do
@@ -232,19 +235,14 @@ RSpec.describe Gitlab::BackgroundMigration do
       end
     end
 
-    context 'when there are scheduled jobs present', :redis do
+    context 'when there are scheduled jobs present' do
       before do
         Sidekiq::Testing.disable! do
+          MergeWorker.perform_in(10.minutes, 'Bar')
           BackgroundMigrationWorker.perform_in(10.minutes, 'Foo')
-
-          expect(Sidekiq::ScheduledSet.new).to be_one
         end
       end
 
-      after do
-        Sidekiq::ScheduledSet.new.clear
-      end
-
       it 'returns true if specific job exists' do
         expect(described_class.exists?('Foo')).to eq(true)
       end
@@ -257,7 +255,10 @@ RSpec.describe Gitlab::BackgroundMigration do
 
   describe '.dead_jobs?' do
     let(:queue) do
-      [double(args: ['Foo', [10, 20]], queue: described_class.queue)]
+      [
+        double(args: ['Foo', [10, 20]], klass: 'BackgroundMigrationWorker'),
+        double(args: ['Bar'], klass: 'MergeWorker')
+      ]
     end
 
     context 'when there are dead jobs present' do
@@ -277,7 +278,10 @@ RSpec.describe Gitlab::BackgroundMigration do
 
   describe '.retrying_jobs?' do
     let(:queue) do
-      [double(args: ['Foo', [10, 20]], queue: described_class.queue)]
+      [
+        double(args: ['Foo', [10, 20]], klass: 'BackgroundMigrationWorker'),
+        double(args: ['Bar'], klass: 'MergeWorker')
+      ]
     end
 
     context 'when there are dead jobs present' do