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.

changelogs/unreleased/58739-hashed-storage-prevent-a-migration-and-rollback-running-at-the-same-time.yml

+---
+title: 'Hashed Storage: Prevent a migration and rollback running at the same time'
+merge_request: 25976
+author:
+type: changed

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

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 '.'

spec/lib/gitlab/hashed_storage/migrator_spec.rb

+# frozen_string_literal: true
+
 require 'spec_helper'
-describe Gitlab::HashedStorage::Migrator do
+describe Gitlab::HashedStorage::Migrator, :sidekiq do
   describe '#bulk_schedule_migration' do
     it 'schedules job to HashedStorage::MigratorWorker' do
       Sidekiq::Testing.fake! do
@@ -182,4 +184,52 @@ describe Gitlab::HashedStorage::Migrator do
       end
     end
   end
+
+  describe 'migration_pending?' do
+    set(:project) { create(:project, :empty_repo) }
+
+    it 'returns true when there are MigratorWorker jobs scheduled' do
+      Sidekiq::Testing.fake! do
+        ::HashedStorage::MigratorWorker.perform_async(1, 5)
+
+        expect(subject.migration_pending?).to be_truthy
+      end
+    end
+
+    it 'returns true when there are ProjectMigrateWorker jobs scheduled' do
+      Sidekiq::Testing.fake! do
+        ::HashedStorage::ProjectMigrateWorker.perform_async(1)
+
+        expect(subject.migration_pending?).to be_truthy
+      end
+    end
+
+    it 'returns false when queues are empty' do
+      expect(subject.migration_pending?).to be_falsey
+    end
+  end
+
+  describe 'rollback_pending?' do
+    set(:project) { create(:project, :empty_repo) }
+
+    it 'returns true when there are RollbackerWorker jobs scheduled' do
+      Sidekiq::Testing.fake! do
+        ::HashedStorage::RollbackerWorker.perform_async(1, 5)
+
+        expect(subject.rollback_pending?).to be_truthy
+      end
+    end
+
+    it 'returns true when there are jobs scheduled' do
+      Sidekiq::Testing.fake! do
+        ::HashedStorage::ProjectRollbackWorker.perform_async(1)
+
+        expect(subject.rollback_pending?).to be_truthy
+      end
+    end
+
+    it 'returns false when queues are empty' do
+      expect(subject.rollback_pending?).to be_falsey
+    end
+  end
 end

spec/tasks/gitlab/storage_rake_spec.rb

 require 'rake_helper'
-describe 'rake gitlab:storage:*' do
+describe 'rake gitlab:storage:*', :sidekiq do
   before do
     Rake.application.rake_require 'tasks/gitlab/storage'
@@ -43,9 +43,7 @@ describe 'rake gitlab:storage:*' do
     end
   end
-  describe 'gitlab:storage:migrate_to_hashed' do
-    let(:task) { 'gitlab:storage:migrate_to_hashed' }
-
+  shared_examples "make sure database is writable" do
     context 'read-only database' do
       it 'does nothing' do
         expect(Gitlab::Database).to receive(:read_only?).and_return(true)
@@ -55,48 +53,68 @@ describe 'rake gitlab:storage:*' do
         expect { run_rake_task(task) }.to output(/This task requires database write access. Exiting./).to_stderr
       end
     end
-    context '0 legacy projects' do
-      it 'does nothing' do
-        expect(::HashedStorage::MigratorWorker).not_to receive(:perform_async)
+  end
+  shared_examples "handles custom BATCH env var" do |worker_klass|
+    context 'in batches of 1' do
+      before do
+        stub_env('BATCH' => 1)
+      end
+
+      it "enqueues one #{worker_klass} per project" do
+        projects.each do |project|
+          expect(worker_klass).to receive(:perform_async).with(project.id, project.id)
+        end
         run_rake_task(task)
       end
     end
-    context '3 legacy projects' do
-      let(:projects) { create_list(:project, 3, :legacy_storage) }
-      context 'in batches of 1' do
-        before do
-          stub_env('BATCH' => 1)
+    context 'in batches of 2' do
+      before do
+        stub_env('BATCH' => 2)
+      end
+      it "enqueues one #{worker_klass} per 2 projects" do
+        projects.map(&:id).sort.each_slice(2) do |first, last|
+          last ||= first
+          expect(worker_klass).to receive(:perform_async).with(first, last)
         end
-        it 'enqueues one HashedStorage::MigratorWorker per project' do
-          projects.each do |project|
-            expect(::HashedStorage::MigratorWorker).to receive(:perform_async).with(project.id, project.id)
-          end
-
-          run_rake_task(task)
-        end
+        run_rake_task(task)
       end
-      context 'in batches of 2' do
-        before do
-          stub_env('BATCH' => 2)
-        end
-        it 'enqueues one HashedStorage::MigratorWorker per 2 projects' do
-          projects.map(&:id).sort.each_slice(2) do |first, last|
-            last ||= first
-            expect(::HashedStorage::MigratorWorker).to receive(:perform_async).with(first, last)
-          end
-          run_rake_task(task)
+    end
+  end
+  describe 'gitlab:storage:migrate_to_hashed' do
+    let(:task) { 'gitlab:storage:migrate_to_hashed' }
+    context 'with rollback already scheduled' do
+      it 'does nothing' do
+        Sidekiq::Testing.fake! do
+          ::HashedStorage::RollbackerWorker.perform_async(1, 5)
+
+          expect(Project).not_to receive(:with_unmigrated_storage)
+          expect { run_rake_task(task) }.to output(/There is already a rollback operation in progress/).to_stderr
         end
       end
     end
+    context 'with 0 legacy projects' do
+      it 'does nothing' do
+        expect(::HashedStorage::MigratorWorker).not_to receive(:perform_async)
+
+        run_rake_task(task)
+      end
+    end
+
+    context 'with 3 legacy projects' do
+      let(:projects) { create_list(:project, 3, :legacy_storage) }
+
+      it_behaves_like "handles custom BATCH env var", ::HashedStorage::MigratorWorker
+    end
+
     context 'with same id in range' do
       it 'displays message when project cant be found' do
         stub_env('ID_FROM', 99999)
@@ -123,6 +141,38 @@ describe 'rake gitlab:storage:*' do
     end
   end
+  describe 'gitlab:storage:rollback_to_legacy' do
+    let(:task) { 'gitlab:storage:rollback_to_legacy' }
+
+    it_behaves_like 'make sure database is writable'
+
+    context 'with migration already scheduled' do
+      it 'does nothing' do
+        Sidekiq::Testing.fake! do
+          ::HashedStorage::MigratorWorker.perform_async(1, 5)
+
+          expect(Project).not_to receive(:with_unmigrated_storage)
+
+          expect { run_rake_task(task) }.to output(/There is already a migration operation in progress/).to_stderr
+        end
+      end
+    end
+
+    context 'with 0 hashed projects' do
+      it 'does nothing' do
+        expect(::HashedStorage::RollbackerWorker).not_to receive(:perform_async)
+
+        run_rake_task(task)
+      end
+    end
+
+    context 'with 3 hashed projects' do
+      let(:projects) { create_list(:project, 3) }
+
+      it_behaves_like "handles custom BATCH env var", ::HashedStorage::RollbackerWorker
+    end
+  end
+
   describe 'gitlab:storage:legacy_projects' do
     it_behaves_like 'rake entities summary', 'projects', 'Legacy' do
       let(:task) { 'gitlab:storage:legacy_projects' }