Reschedule DeleteDiffFiles until there is none left to remove

db/post_migrate/20180619121030_enqueue_delete_diff_files_workers.rb

@@ -3,14 +3,24 @@ class EnqueueDeleteDiffFilesWorkers < ActiveRecord::Migration
 
   DOWNTIME = false
   MIGRATION = 'DeleteDiffFiles'.freeze
+  TMP_INDEX = 'tmp_partial_diff_id_with_files_index'.freeze
 
   disable_ddl_transaction!
 
   def up
+    unless index_exists_by_name?(:merge_request_diffs, TMP_INDEX)
+      add_concurrent_index(:merge_request_diffs, :id, where: "(state NOT IN ('without_files', 'empty'))", name: TMP_INDEX)
+    end
+
     BackgroundMigrationWorker.perform_async(MIGRATION)
+
+    # We don't remove the index since it's going to be used on DeleteDiffFiles
+    # worker. We should remove it in an upcoming release.
   end
 
   def down
-    # no-op
+    if index_exists_by_name?(:merge_request_diffs, TMP_INDEX)
+      remove_concurrent_index_by_name(:merge_request_diffs, TMP_INDEX)
+    end
   end
 end

lib/gitlab/background_migration/delete_diff_files.rb

@@ -19,51 +19,60 @@ module Gitlab
         include EachBatch
       end
 
-      BATCH = 5_000
+      DIFF_ROWS_LIMIT = 5_000
       DEAD_TUPLES_THRESHOLD = 50_000
       VACUUM_WAIT_TIME = 5.minutes
 
       def perform
-        diffs_with_files = MergeRequestDiff
-          .joins(:merge_request)
-          .where("merge_requests.state = 'merged'")
-          .where('merge_requests.latest_merge_request_diff_id IS NOT NULL')
-          .where('merge_requests.latest_merge_request_diff_id != merge_request_diffs.id')
-          .where("merge_request_diffs.state NOT IN ('without_files', 'empty')")
-
-        diffs_with_files.each_batch(of: BATCH) do |batch, index|
-          wait_deadtuple_vacuum(index)
-          prune_diff_files(batch, index)
+        rescheduling do
+          prune_diff_files(diffs_collection.limit(DIFF_ROWS_LIMIT))
         end
       end
 
-      def wait_deadtuple_vacuum(index)
-        db_klass = Gitlab::Database
+      def should_wait_deadtuple_vacuum?
+        return false unless Gitlab::Database.postgresql?
 
-        if defined?(db_klass) && db_klass.respond_to?(:postgresql?) && db_klass.postgresql?
-          while diff_files_dead_tuples_count >= DEAD_TUPLES_THRESHOLD
-            log_info("Dead tuple threshold hit on merge_request_diff_files (#{index}th batch): " \
-                     "#{diff_files_dead_tuples_count}, waiting 5 minutes")
-            sleep VACUUM_WAIT_TIME
-          end
-        end
+        diff_files_dead_tuples_count >= DEAD_TUPLES_THRESHOLD
       end
 
       private
 
+      def rescheduling(&block)
+        # We should reschedule until deadtuples get in a desirable
+        # state (e.g. < 50_000). That may take move than one reschedule.
+        #
+        if should_wait_deadtuple_vacuum?
+          reschedule
+          return
+        end
+
+        block.call
+
+        reschedule if diffs_collection.limit(1).count > 0
+      end
+
+      def reschedule
+        BackgroundMigrationWorker.perform_in(VACUUM_WAIT_TIME, self.class.name.demodulize)
+      end
+
+      def diffs_collection
+        MergeRequestDiff
+          .joins(:merge_request)
+          .where("merge_requests.state = 'merged'")
+          .where('merge_requests.latest_merge_request_diff_id IS NOT NULL')
+          .where('merge_requests.latest_merge_request_diff_id != merge_request_diffs.id')
+          .where("merge_request_diffs.state NOT IN ('without_files', 'empty')")
+      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]
 
-        if dead_tuple.present?
-          dead_tuple['n_dead_tup'].to_i
-        else
-          0
-        end
+        dead_tuple&.fetch('n_dead_tup', 0).to_i
       end
 
-      def prune_diff_files(batch, index)
+      def prune_diff_files(batch)
         diff_ids = batch.pluck(:id)
 
         removed = 0
@@ -76,7 +85,7 @@ module Gitlab
             .delete_all
         end
 
-        log_info("#{index}th batch - Removed #{removed} merge_request_diff_files rows, "\
+        log_info("Removed #{removed} merge_request_diff_files rows, "\
                  "updated #{updated} merge_request_diffs rows")
       end
 

spec/lib/gitlab/background_migration/delete_diff_files_spec.rb

@@ -65,10 +65,27 @@ describe Gitlab::BackgroundMigration::DeleteDiffFiles, :migration, schema: 20180
         .not_to change { merge_request_diff.merge_request_diff_files.count }
         .from(20)
     end
+
+    it 'reschedules itself when should_wait_deadtuple_vacuum' do
+      Sidekiq::Testing.fake! do
+        worker = described_class.new
+
+        allow(worker).to receive(:should_wait_deadtuple_vacuum?) { true }
+
+        expect(BackgroundMigrationWorker)
+          .to receive(:perform_in)
+          .with(described_class::VACUUM_WAIT_TIME, 'DeleteDiffFiles')
+          .and_call_original
+
+        worker.perform
+
+        expect(BackgroundMigrationWorker.jobs.size).to eq(1)
+      end
+    end
   end
 
-  describe '#wait_deadtuple_vacuum' do
-    it 'sleeps process for VACUUM_WAIT_TIME when hitting DEAD_TUPLES_THRESHOLD', :postgresql do
+  describe '#should_wait_deadtuple_vacuum?' do
+    it 'returns true when hitting merge_request_diff_files hits DEAD_TUPLES_THRESHOLD', :postgresql do
       worker = described_class.new
       threshold_query_result = [{ "n_dead_tup" => described_class::DEAD_TUPLES_THRESHOLD.to_s }]
       normal_query_result = [{ "n_dead_tup" => '3' }]
@@ -78,9 +95,7 @@ describe Gitlab::BackgroundMigration::DeleteDiffFiles, :migration, schema: 20180
         .with(/SELECT n_dead_tup */)
         .and_return(threshold_query_result, normal_query_result)
 
-      expect(worker).to receive(:sleep).with(described_class::VACUUM_WAIT_TIME).once
-
-      worker.wait_deadtuple_vacuum(1)
+      expect(worker.should_wait_deadtuple_vacuum?).to be(true)
     end
   end
 end