Implement draining scheduled sets of background migrations

lib/gitlab/background_migration.rb

@@ -1,7 +1,7 @@
 module Gitlab
   module BackgroundMigration
     def self.queue
-      BackgroundMigrationWorker.sidekiq_options['queue']
+      @queue ||= BackgroundMigrationWorker.sidekiq_options['queue']
     end
 
     # Begins stealing jobs from the background migrations queue, blocking the
@@ -9,11 +9,14 @@ module Gitlab
     #
     # steal_class - The name of the class for which to steal jobs.
     def self.steal(steal_class)
-      queue = Sidekiq::Queue.new(self.queue)
+      enqueued = Sidekiq::Queue.new(self.queue)
+      scheduled = Sidekiq::ScheduledSet.new
 
+      [scheduled, enqueued].each do |queue|
         queue.each do |job|
           migration_class, migration_args = job.args
 
+          next unless job.queue == self.queue
           next unless migration_class == steal_class
 
           perform(migration_class, migration_args)
@@ -21,6 +24,7 @@ module Gitlab
           job.delete
         end
       end
+    end
 
     # class_name - The name of the background migration class as defined in the
     #              Gitlab::BackgroundMigration namespace.
@@ -28,6 +32,7 @@ module Gitlab
     # arguments - The arguments to pass to the background migration's "perform"
     #             method.
     def self.perform(class_name, arguments)
+      puts class_name
       const_get(class_name).new.perform(*arguments)
     end
   end

spec/lib/gitlab/background_migration_spec.rb

@@ -34,6 +34,20 @@ describe Gitlab::BackgroundMigration do
         described_class.steal('Bar')
       end
     end
+
+    context 'when there are scheduled jobs present', :sidekiq, :redis do
+      it 'steals all jobs from the schedule sets' do
+        Sidekiq::Testing.disable! do
+          BackgroundMigrationWorker.perform_in(10.minutes, 'Object')
+          expect(Sidekiq::ScheduledSet.new).to be_one
+          expect(described_class).to receive(:perform).with('Object', any_args)
+
+          described_class.steal('Object')
+
+          expect(Sidekiq::ScheduledSet.new).to be_none
+        end
+      end
+    end
   end
 
   describe '.perform' do

spec/support/sidekiq.rb

@@ -8,4 +8,8 @@ RSpec.configure do |config|
   config.after(:each, :sidekiq) do
     Sidekiq::Worker.clear_all
   end
+
+  config.after(:each, :sidekiq, :redis) do
+    Sidekiq.redis { |redis| redis.flushdb }
+  end
 end