From 2a0ead2c4730110c898127660fe6600155539f0c Mon Sep 17 00:00:00 2001
From: Grzegorz Bizon <grzesiek.bizon@gmail.com>
Date: Thu, 13 Jul 2017 11:44:52 +0200
Subject: [PATCH] Implement draining scheduled sets of background migrations

---
 lib/gitlab/background_migration.rb           | 19 ++++++++++++-------
 spec/lib/gitlab/background_migration_spec.rb | 14 ++++++++++++++
 spec/support/sidekiq.rb                      |  4 ++++
 3 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/lib/gitlab/background_migration.rb b/lib/gitlab/background_migration.rb
index 6ddffe70da4..73688421e74 100644
--- a/lib/gitlab/background_migration.rb
+++ b/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,16 +9,20 @@ 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
 
-      queue.each do |job|
-        migration_class, migration_args = job.args
+      [scheduled, enqueued].each do |queue|
+        queue.each do |job|
+          migration_class, migration_args = job.args
 
-        next unless migration_class == steal_class
+          next unless job.queue == self.queue
+          next unless migration_class == steal_class
 
-        perform(migration_class, migration_args)
+          perform(migration_class, migration_args)
 
-        job.delete
+          job.delete
+        end
       end
     end
 
@@ -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
diff --git a/spec/lib/gitlab/background_migration_spec.rb b/spec/lib/gitlab/background_migration_spec.rb
index 93ffaab65ce..d823c5342ae 100644
--- a/spec/lib/gitlab/background_migration_spec.rb
+++ b/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
diff --git a/spec/support/sidekiq.rb b/spec/support/sidekiq.rb
index 5478fea4e64..d143014692d 100644
--- a/spec/support/sidekiq.rb
+++ b/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