Make migrations reschedulable

5 files changed, 92 insertions, 27 deletions

diff --git a/lib/gitlab/background_migration/delete_diff_files.rb b/lib/gitlab/background_migration/delete_diff_files.rb
index 11851c23ee3..4cb068eb71a 100644
--- a/lib/gitlab/background_migration/delete_diff_files.rb
+++ b/lib/gitlab/background_migration/delete_diff_files.rb
@@ -27,7 +27,7 @@ module Gitlab
 
       private
 
-      def should_reschedule?
+      def need_reschedule?
         wait_for_deadtuple_vacuum?(MergeRequestDiffFile.table_name)
       end
 
diff --git a/lib/gitlab/background_migration/helpers/reschedulable.rb b/lib/gitlab/background_migration/helpers/reschedulable.rb
index 3b653059858..6087181609e 100644
--- a/lib/gitlab/background_migration/helpers/reschedulable.rb
+++ b/lib/gitlab/background_migration/helpers/reschedulable.rb
@@ -14,7 +14,7 @@ module Gitlab
         extend ActiveSupport::Concern
 
         def reschedule_if_needed(args, &block)
-          if should_reschedule?
+          if need_reschedule?
             BackgroundMigrationWorker.perform_in(vacuum_wait_time, self.class.name.demodulize, args)
           else
             yield
@@ -22,9 +22,9 @@ module Gitlab
         end
 
         # Override this on base class if you need a different reschedule condition
-        def should_reschedule?
-          raise NotImplementedError, "#{self.class} does not implement #{__method__}"
-        end
+        # def need_reschedule?
+        #   raise NotImplementedError, "#{self.class} does not implement #{__method__}"
+        # end
 
         def wait_for_deadtuple_vacuum?(table_name)
           return false unless Gitlab::Database.postgresql?
diff --git a/lib/gitlab/background_migration/sync_issues_state_id.rb b/lib/gitlab/background_migration/sync_issues_state_id.rb
index 33b997c8533..f95c433c64c 100644
--- a/lib/gitlab/background_migration/sync_issues_state_id.rb
+++ b/lib/gitlab/background_migration/sync_issues_state_id.rb
@@ -4,19 +4,29 @@
 module Gitlab
   module BackgroundMigration
     class SyncIssuesStateId
+      include Helpers::Reschedulable
+
       def perform(start_id, end_id)
         Rails.logger.info("Issues - Populating state_id: #{start_id} - #{end_id}")
 
-        ActiveRecord::Base.connection.execute <<~SQL
-          UPDATE issues
-          SET state_id =
-            CASE state
-            WHEN 'opened' THEN 1
-            WHEN 'closed' THEN 2
-            END
-          WHERE state_id IS NULL
-          AND id BETWEEN #{start_id} AND #{end_id}
-        SQL
+        reschedule_if_needed([start_id, end_id]) do
+          ActiveRecord::Base.connection.execute <<~SQL
+            UPDATE issues
+            SET state_id =
+              CASE state
+              WHEN 'opened' THEN 1
+              WHEN 'closed' THEN 2
+              END
+            WHERE state_id IS NULL
+            AND id BETWEEN #{start_id} AND #{end_id}
+          SQL
+        end
+      end
+
+      private
+
+      def need_reschedule?
+        wait_for_deadtuple_vacuum?('issues')
       end
     end
   end
diff --git a/lib/gitlab/background_migration/sync_merge_requests_state_id.rb b/lib/gitlab/background_migration/sync_merge_requests_state_id.rb
index 923ceaeec54..5d92553f577 100644
--- a/lib/gitlab/background_migration/sync_merge_requests_state_id.rb
+++ b/lib/gitlab/background_migration/sync_merge_requests_state_id.rb
@@ -4,21 +4,31 @@
 module Gitlab
   module BackgroundMigration
     class SyncMergeRequestsStateId
+      include Helpers::Reschedulable
+
       def perform(start_id, end_id)
         Rails.logger.info("Merge Requests - Populating state_id: #{start_id} - #{end_id}")
 
-        ActiveRecord::Base.connection.execute <<~SQL
-          UPDATE merge_requests
-          SET state_id =
-            CASE state
-            WHEN 'opened' THEN 1
-            WHEN 'closed' THEN 2
-            WHEN 'merged' THEN 3
-            WHEN 'locked' THEN 4
-            END
-          WHERE state_id IS NULL
-          AND id BETWEEN #{start_id} AND #{end_id}
-        SQL
+        reschedule_if_needed([start_id, end_id]) do
+          ActiveRecord::Base.connection.execute <<~SQL
+            UPDATE merge_requests
+            SET state_id =
+              CASE state
+              WHEN 'opened' THEN 1
+              WHEN 'closed' THEN 2
+              WHEN 'merged' THEN 3
+              WHEN 'locked' THEN 4
+              END
+            WHERE state_id IS NULL
+            AND id BETWEEN #{start_id} AND #{end_id}
+          SQL
+        end
+      end
+
+      private
+
+      def need_reschedule?
+        wait_for_deadtuple_vacuum?('issues')
       end
     end
   end
diff --git a/spec/migrations/schedule_sync_issuables_state_id_spec.rb b/spec/migrations/schedule_sync_issuables_state_id_spec.rb
index a926ee38387..a746fb68c30 100644
--- a/spec/migrations/schedule_sync_issuables_state_id_spec.rb
+++ b/spec/migrations/schedule_sync_issuables_state_id_spec.rb
@@ -15,7 +15,40 @@ describe ScheduleSyncIssuablesStateId, :migration do
     @project = projects.create!(namespace_id: @group.id)
   end
 
+  shared_examples 'rescheduling migrations' do
+    before do
+      Sidekiq::Worker.clear_all
+    end
+
+    it 'reschedules migrations when should wait for dead tuple vacuum' do
+      worker = worker_class.new
+
+      Sidekiq::Testing.fake! do
+        allow(worker).to receive(:wait_for_deadtuple_vacuum?) { true }
+
+        worker.perform(resource_1.id, resource_2.id)
+
+        expect(worker_class.name.demodulize).to be_scheduled_delayed_migration(5.minutes, resource_1.id, resource_2.id)
+        expect(BackgroundMigrationWorker.jobs.size).to eq(1)
+      end
+    end
+  end
+
   describe '#up' do
+    # it 'correctly schedules diff file deletion workers' do
+    #   Sidekiq::Testing.fake! do
+    #     Timecop.freeze do
+    #       described_class.new.perform
+
+    #       expect(described_class::MIGRATION).to be_scheduled_delayed_migration(5.minutes, [1, 4, 5])
+
+    #       expect(described_class::MIGRATION).to be_scheduled_delayed_migration(10.minutes, [6])
+
+    #       expect(BackgroundMigrationWorker.jobs.size).to eq(2)
+    #     end
+    #   end
+    # end
+
     context 'issues' do
       it 'migrates state column to integer' do
         opened_issue = issues.create!(description: 'first', state: 'opened')
@@ -30,6 +63,12 @@ describe ScheduleSyncIssuablesStateId, :migration do
         expect(invalid_state_issue.reload.state_id).to be_nil
         expect(nil_state_issue.reload.state_id).to be_nil
       end
+
+      it_behaves_like 'rescheduling migrations' do
+        let(:worker_class) { Gitlab::BackgroundMigration::SyncIssuesStateId }
+        let(:resource_1) { issues.create!(description: 'first', state: 'opened') }
+        let(:resource_2) { issues.create!(description: 'second', state: 'closed') }
+      end
     end
 
     context 'merge requests' do
@@ -48,6 +87,12 @@ describe ScheduleSyncIssuablesStateId, :migration do
         expect(locked_merge_request.reload.state_id).to eq(MergeRequest.available_states[:locked])
         expect(invalid_state_merge_request.reload.state_id).to be_nil
       end
+
+      it_behaves_like 'rescheduling migrations' do
+        let(:worker_class) { Gitlab::BackgroundMigration::SyncMergeRequestsStateId }
+        let(:resource_1) { merge_requests.create!(state: 'opened', target_project_id: @project.id, target_branch: 'feature1', source_branch: 'master') }
+        let(:resource_2) { merge_requests.create!(state: 'closed', target_project_id: @project.id, target_branch: 'feature2', source_branch: 'master') }
+      end
     end
   end
 end