Prevent storage migration and rollback running at the same time

This is a small polishing on the storage migration and storage rollback
rake tasks. By aborting a migration while a rollback is already
scheduled we want to prevent unexpected consequences.

2 files changed, 32 insertions, 1 deletions

diff --git a/lib/gitlab/hashed_storage/migrator.rb b/lib/gitlab/hashed_storage/migrator.rb
index 7046b4e2a43..f5368d41629 100644
--- a/lib/gitlab/hashed_storage/migrator.rb
+++ b/lib/gitlab/hashed_storage/migrator.rb
@@ -81,8 +81,26 @@ module Gitlab
         Rails.logger.error("#{err.message} rolling-back storage of #{project.full_path} (ID=#{project.id}), trace - #{err.backtrace}")
       end
 
+      # Returns whether we have any pending storage migration
+      #
+      def migration_pending?
+        any_non_empty_queue?(::HashedStorage::MigratorWorker, ::HashedStorage::ProjectMigrateWorker)
+      end
+
+      # Returns whether we have any pending storage rollback
+      #
+      def rollback_pending?
+        any_non_empty_queue?(::HashedStorage::RollbackerWorker, ::HashedStorage::ProjectRollbackWorker)
+      end
+
       private
 
+      def any_non_empty_queue?(*workers)
+        workers.any? do |worker|
+          worker.jobs.any?
+        end
+      end
+
       # rubocop: disable CodeReuse/ActiveRecord
       def build_relation(start, finish)
         relation = Project
diff --git a/lib/tasks/gitlab/storage.rake b/lib/tasks/gitlab/storage.rake
index a2136ce1b92..954f827f716 100644
--- a/lib/tasks/gitlab/storage.rake
+++ b/lib/tasks/gitlab/storage.rake
@@ -11,6 +11,13 @@ namespace :gitlab do
       storage_migrator = Gitlab::HashedStorage::Migrator.new
       helper = Gitlab::HashedStorage::RakeHelper
 
+      if storage_migrator.rollback_pending?
+        warn "There is already a rollback operation in progress, " \
+             "running a migration at the same time may have unexpected consequences."
+
+        next
+      end
+
       if helper.range_single_item?
         project = Project.with_unmigrated_storage.find_by(id: helper.range_from)
 
@@ -56,6 +63,13 @@ namespace :gitlab do
       storage_migrator = Gitlab::HashedStorage::Migrator.new
       helper = Gitlab::HashedStorage::RakeHelper
 
+      if storage_migrator.migration_pending?
+        warn "There is already a migration operation in progress, " \
+             "running a rollback at the same time may have unexpected consequences."
+
+        next
+      end
+
       if helper.range_single_item?
         project = Project.with_storage_feature(:repository).find_by(id: helper.range_from)
 
@@ -82,7 +96,6 @@ namespace :gitlab do
       print "Enqueuing rollback of #{hashed_projects_count} projects in batches of #{helper.batch_size}"
 
       helper.project_id_batches_rollback do |start, finish|
-        puts "Start: #{start} FINISH: #{finish}"
         storage_migrator.bulk_schedule_rollback(start: start, finish: finish)
 
         print '.'