Add latest changes from gitlab-org/gitlab@master

1 files changed, 15 insertions, 87 deletions

diff --git a/lib/gitlab/background_migration.rb b/lib/gitlab/background_migration.rb
index 0826887dd0a..3e80a17dd1e 100644
--- a/lib/gitlab/background_migration.rb
+++ b/lib/gitlab/background_migration.rb
@@ -2,8 +2,12 @@
 
 module Gitlab
   module BackgroundMigration
-    def self.queue
-      @queue ||= BackgroundMigrationWorker.sidekiq_options['queue']
+    def self.coordinator_for_database(database)
+      JobCoordinator.for_database(database)
+    end
+
+    def self.queue(database: :main)
+      coordinator_for_database(database).queue
     end
 
     # Begins stealing jobs from the background migrations queue, blocking the
@@ -16,35 +20,10 @@ module Gitlab
     # re-raises the exception.
     #
     # steal_class - The name of the class for which to steal jobs.
-    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
-
-      queues.each do |queue|
-        queue.each do |job|
-          migration_class, migration_args = job.args
-
-          next unless job.klass == 'BackgroundMigrationWorker'
-          next unless migration_class == steal_class
-          next if block_given? && !(yield job)
-
-          begin
-            perform(migration_class, migration_args) if job.delete
-          rescue Exception # rubocop:disable Lint/RescueException
-            BackgroundMigrationWorker # enqueue this migration again
-              .perform_async(migration_class, migration_args)
-
-            raise
-          end
-        end
-      end
+    # retry_dead_jobs - Flag to control whether jobs in Sidekiq::RetrySet or Sidekiq::DeadSet are retried.
+    # database - tracking database this migration executes against
+    def self.steal(steal_class, retry_dead_jobs: false, database: :main, &block)
+      coordinator_for_database(database).steal(steal_class, retry_dead_jobs: retry_dead_jobs, &block)
     end
 
     ##
@@ -55,64 +34,13 @@ module Gitlab
     #
     # arguments - The arguments to pass to the background migration's "perform"
     #             method.
-    def self.perform(class_name, arguments)
-      migration_class_for(class_name).new.perform(*arguments)
-    end
-
-    def self.remaining
-      enqueued = Sidekiq::Queue.new(self.queue)
-      scheduled = Sidekiq::ScheduledSet.new
-
-      [enqueued, scheduled].sum do |set|
-        set.count do |job|
-          job.klass == 'BackgroundMigrationWorker'
-        end
-      end
-    end
-
-    def self.exists?(migration_class, additional_queues = [])
-      enqueued = Sidekiq::Queue.new(self.queue)
-      scheduled = Sidekiq::ScheduledSet.new
-
-      enqueued_job?([enqueued, scheduled], migration_class)
-    end
-
-    def self.dead_jobs?(migration_class)
-      dead_set = Sidekiq::DeadSet.new
-
-      enqueued_job?([dead_set], migration_class)
-    end
-
-    def self.retrying_jobs?(migration_class)
-      retry_set = Sidekiq::RetrySet.new
-
-      enqueued_job?([retry_set], migration_class)
-    end
-
-    def self.migration_class_for(class_name)
-      # We don't pass class name with Gitlab::BackgroundMigration:: prefix anymore
-      # but some jobs could be already spawned so we need to have some backward compatibility period.
-      # Can be removed since 13.x
-      full_class_name_prefix_regexp = /\A(::)?Gitlab::BackgroundMigration::/
-
-      if class_name.match(full_class_name_prefix_regexp)
-        Gitlab::ErrorTracking.track_and_raise_for_dev_exception(
-          StandardError.new("Full class name is used"),
-          class_name: class_name
-        )
-
-        class_name = class_name.sub(full_class_name_prefix_regexp, '')
-      end
-
-      const_get(class_name, false)
+    # database - tracking database this migration executes against
+    def self.perform(class_name, arguments, database: :main)
+      coordinator_for_database(database).perform(class_name, arguments)
     end
 
-    def self.enqueued_job?(queues, migration_class)
-      queues.any? do |queue|
-        queue.any? do |job|
-          job.klass == 'BackgroundMigrationWorker' && job.args.first == migration_class
-        end
-      end
+    def self.exists?(migration_class, additional_queues = [], database: :main)
+      coordinator_for_database(database).exists?(migration_class, additional_queues) # rubocop:disable CodeReuse/ActiveRecord
     end
   end
 end