Remove cleaned up OIDs from database and cache

3 files changed, 51 insertions, 4 deletions

diff --git a/app/models/merge_request_diff.rb b/app/models/merge_request_diff.rb
index 5f5d92bc2f0..f45bd0e03de 100644
--- a/app/models/merge_request_diff.rb
+++ b/app/models/merge_request_diff.rb
@@ -51,6 +51,10 @@ class MergeRequestDiff < ApplicationRecord
     joins(:merge_request_diff_commits).where(merge_request_diff_commits: { sha: sha }).reorder(nil)
   end
 
+  scope :by_project_id, -> (project_id) do
+    joins(:merge_request).where(merge_requests: { target_project_id: project_id })
+  end
+
   scope :recent, -> { order(id: :desc).limit(100) }
   scope :files_in_database, -> { where(stored_externally: [false, nil]) }
 
diff --git a/app/models/note_diff_file.rb b/app/models/note_diff_file.rb
index 9afb94c869a..fcc9e2b3fd8 100644
--- a/app/models/note_diff_file.rb
+++ b/app/models/note_diff_file.rb
@@ -7,6 +7,10 @@ class NoteDiffFile < ApplicationRecord
     joins(:diff_note).where("resolved_at IS NULL OR noteable_type = 'Commit'")
   end
 
+  scope :referencing_sha, -> (oids, project_id:) do
+    joins(:diff_note).where(notes: { project_id: project_id, commit_id: oids })
+  end
+
   delegate :original_position, :project, to: :diff_note
 
   belongs_to :diff_note, inverse_of: :note_diff_file
diff --git a/app/services/projects/cleanup_service.rb b/app/services/projects/cleanup_service.rb
index 12103ea34b5..5972bfd4071 100644
--- a/app/services/projects/cleanup_service.rb
+++ b/app/services/projects/cleanup_service.rb
@@ -18,9 +18,6 @@ module Projects
     # per rewritten object, with the old and new SHAs space-separated. It can be
     # used to update or remove content that references the objects that BFG has
     # altered
-    #
-    # Currently, only the project repository is modified by this service, but we
-    # may wish to modify other data sources in the future.
     def execute
       apply_bfg_object_map!
 
@@ -41,10 +38,52 @@ module Projects
       raise NoUploadError unless project.bfg_object_map.exists?
 
       project.bfg_object_map.open do |io|
-        repository_cleaner.apply_bfg_object_map(io)
+        repository_cleaner.apply_bfg_object_map_stream(io) do |response|
+          cleanup_diffs(response)
+        end
+      end
+    end
+
+    def cleanup_diffs(response)
+      old_commit_shas = extract_old_commit_shas(response.entries)
+
+      ActiveRecord::Base.transaction do
+        cleanup_merge_request_diffs(old_commit_shas)
+        cleanup_note_diff_files(old_commit_shas)
       end
     end
 
+    def extract_old_commit_shas(batch)
+      batch.lazy.select { |entry| entry.type == :COMMIT }.map(&:old_oid).force
+    end
+
+    def cleanup_merge_request_diffs(old_commit_shas)
+      merge_request_diffs = MergeRequestDiff
+        .by_project_id(project.id)
+        .by_commit_sha(old_commit_shas)
+
+      # It's important to run the ActiveRecord callbacks here
+      merge_request_diffs.destroy_all # rubocop:disable Cop/DestroyAll
+
+      # TODO: ensure the highlight cache is removed immediately. It's too hard
+      # to calculate the Redis keys at present.
+      #
+      # https://gitlab.com/gitlab-org/gitlab-ce/issues/61115
+    end
+
+    def cleanup_note_diff_files(old_commit_shas)
+      # Pluck the IDs instead of running the query twice to ensure we clear the
+      # cache for exactly the note diffs we remove
+      ids = NoteDiffFile
+        .referencing_sha(old_commit_shas, project_id: project.id)
+        .pluck_primary_key
+
+      NoteDiffFile.id_in(ids).delete_all
+
+      # A highlighted version of the diff is stored in redis. Remove it now.
+      Gitlab::DiscussionsDiff::HighlightCache.clear_multiple(ids)
+    end
+
     def repository_cleaner
       @repository_cleaner ||= Gitlab::Git::RepositoryCleaner.new(repository.raw)
     end