Address review comments

5 files changed, 93 insertions, 70 deletions

diff --git a/lib/gitlab/background_migration/delete_diff_files.rb b/lib/gitlab/background_migration/delete_diff_files.rb
index 0066faa23dd..664ead1af44 100644
--- a/lib/gitlab/background_migration/delete_diff_files.rb
+++ b/lib/gitlab/background_migration/delete_diff_files.rb
@@ -4,8 +4,6 @@
 module Gitlab
   module BackgroundMigration
     class DeleteDiffFiles
-      include Helpers::Reschedulable
-
       class MergeRequestDiff < ActiveRecord::Base
         self.table_name = 'merge_request_diffs'
 
@@ -17,24 +15,47 @@ module Gitlab
         self.table_name = 'merge_request_diff_files'
       end
 
+      DEAD_TUPLES_THRESHOLD = 50_000
+      VACUUM_WAIT_TIME = 5.minutes
+
       def perform(ids)
         @ids = ids
 
-        reschedule_if_needed(ids) do
-          prune_diff_files
+        # We should reschedule until deadtuples get in a desirable
+        # state (e.g. < 50_000). That may take more than one reschedule.
+        #
+        if should_wait_deadtuple_vacuum?
+          reschedule
+          return
         end
+
+        prune_diff_files
+      end
+
+      def should_wait_deadtuple_vacuum?
+        return false unless Gitlab::Database.postgresql?
+
+        diff_files_dead_tuples_count >= DEAD_TUPLES_THRESHOLD
       end
 
       private
 
-      def should_reschedule?
-        wait_for_deadtuple_vacuum?(MergeRequestDiffFile.table_name)
+      def reschedule
+        BackgroundMigrationWorker.perform_in(VACUUM_WAIT_TIME, self.class.name.demodulize, [@ids])
       end
 
       def diffs_collection
         MergeRequestDiff.where(id: @ids)
       end
 
+      def diff_files_dead_tuples_count
+        dead_tuple =
+          execute_statement("SELECT n_dead_tup FROM pg_stat_all_tables "\
+                            "WHERE relname = 'merge_request_diff_files'")[0]
+
+        dead_tuple&.fetch('n_dead_tup', 0).to_i
+      end
+
       def prune_diff_files
         removed = 0
         updated = 0
@@ -48,6 +69,10 @@ module Gitlab
                  "updated #{updated} merge_request_diffs rows")
       end
 
+      def execute_statement(sql)
+        ActiveRecord::Base.connection.execute(sql)
+      end
+
       def log_info(message)
         Rails.logger.info("BackgroundMigration::DeleteDiffFiles - #{message}")
       end
diff --git a/lib/gitlab/background_migration/helpers/reschedulable.rb b/lib/gitlab/background_migration/helpers/reschedulable.rb
deleted file mode 100644
index 60ab62cd973..00000000000
--- a/lib/gitlab/background_migration/helpers/reschedulable.rb
+++ /dev/null
@@ -1,62 +0,0 @@
-# frozen_string_literal: true
-
-module Gitlab
-  module BackgroundMigration
-    module Helpers
-      # == Reschedulable helper
-      #
-      # Allows background migrations to be rescheduled if a condition is met,
-      # the condition should be overridden in classes in #should_reschedule? method.
-      #
-      # For example, check DeleteDiffFiles migration which is rescheduled if dead tuple count
-      # on DB is not acceptable.
-      #
-      module Reschedulable
-        extend ActiveSupport::Concern
-
-        # Use this method to perform the background migration and it will be rescheduled
-        # if #should_reschedule? returns true.
-        def reschedule_if_needed(*args, &block)
-          if should_reschedule?
-            BackgroundMigrationWorker.perform_in(wait_time, self.class.name.demodulize, args)
-          else
-            yield
-          end
-        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
-
-        # Override in subclass if a different dead tuple threshold
-        def dead_tuples_threshold
-          @dead_tuples_threshold ||= 50_000
-        end
-
-        # Override in subclass if a different wait time
-        def wait_time
-          @wait_time ||= 5.minutes
-        end
-
-        def execute_statement(sql)
-          ActiveRecord::Base.connection.execute(sql)
-        end
-
-        def wait_for_deadtuple_vacuum?(table_name)
-          return false unless Gitlab::Database.postgresql?
-
-          dead_tuples_count_for(table_name) >= dead_tuples_threshold
-        end
-
-        def dead_tuples_count_for(table_name)
-          dead_tuple =
-            execute_statement("SELECT n_dead_tup FROM pg_stat_all_tables "\
-                              "WHERE relname = '#{table_name}'")[0]
-
-          dead_tuple&.fetch('n_dead_tup', 0).to_i
-        end
-      end
-    end
-  end
-end
diff --git a/lib/gitlab/background_migration/reschedulable.rb b/lib/gitlab/background_migration/reschedulable.rb
new file mode 100644
index 00000000000..3d6781c07c0
--- /dev/null
+++ b/lib/gitlab/background_migration/reschedulable.rb
@@ -0,0 +1,60 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module BackgroundMigration
+    # == Reschedulable helper
+    #
+    # Allows background migrations to be rescheduled if a condition is met,
+    # the condition should be overridden in classes in #should_reschedule? method.
+    #
+    # For example, check DeleteDiffFiles migration which is rescheduled if dead tuple count
+    # on DB is not acceptable.
+    #
+    module Reschedulable
+      extend ActiveSupport::Concern
+
+      # Use this method to perform the background migration and it will be rescheduled
+      # if #should_reschedule? returns true.
+      def reschedule_if_needed(*args, &block)
+        if should_reschedule?
+          BackgroundMigrationWorker.perform_in(wait_time, self.class.name.demodulize, args)
+        else
+          yield
+        end
+      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
+
+      # Override in subclass if a different dead tuple threshold
+      def dead_tuples_threshold
+        @dead_tuples_threshold ||= 50_000
+      end
+
+      # Override in subclass if a different wait time
+      def wait_time
+        @wait_time ||= 5.minutes
+      end
+
+      def execute_statement(sql)
+        ActiveRecord::Base.connection.execute(sql)
+      end
+
+      def wait_for_deadtuple_vacuum?(table_name)
+        return false unless Gitlab::Database.postgresql?
+
+        dead_tuples_count_for(table_name) >= dead_tuples_threshold
+      end
+
+      def dead_tuples_count_for(table_name)
+        dead_tuple =
+          execute_statement("SELECT n_dead_tup FROM pg_stat_all_tables "\
+                            "WHERE relname = '#{table_name}'")[0]
+
+        dead_tuple&.fetch('n_dead_tup', 0).to_i
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/background_migration/sync_issues_state_id.rb b/lib/gitlab/background_migration/sync_issues_state_id.rb
index 50841c3fe4d..053d154cef8 100644
--- a/lib/gitlab/background_migration/sync_issues_state_id.rb
+++ b/lib/gitlab/background_migration/sync_issues_state_id.rb
@@ -4,7 +4,7 @@
 module Gitlab
   module BackgroundMigration
     class SyncIssuesStateId
-      include Helpers::Reschedulable
+      include Reschedulable
 
       def perform(start_id, end_id)
         Rails.logger.info("Issues - Populating state_id: #{start_id} - #{end_id}")
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 4cf8e993c52..a94529aaf5c 100644
--- a/lib/gitlab/background_migration/sync_merge_requests_state_id.rb
+++ b/lib/gitlab/background_migration/sync_merge_requests_state_id.rb
@@ -4,7 +4,7 @@
 module Gitlab
   module BackgroundMigration
     class SyncMergeRequestsStateId
-      include Helpers::Reschedulable
+      include Reschedulable
 
       def perform(start_id, end_id)
         Rails.logger.info("Merge Requests - Populating state_id: #{start_id} - #{end_id}")