Merge branch '226924-reduce-keep-around-commits' into 'master'

Reduce Gitaly calls for keeping around refs of published notes

See merge request gitlab-org/gitlab!68337

app/models/diff_note.rb

@@ -22,7 +22,7 @@ class DiffNote < Note
   validate :verify_supported, unless: :importing?
 
   before_validation :set_line_code, if: :on_text?, unless: :importing?
-  after_save :keep_around_commits, unless: :importing?
+  after_save :keep_around_commits, unless: -> { importing? || skip_keep_around_commits }
 
   NoteDiffFileCreationError = Class.new(StandardError)
 
@@ -115,6 +115,20 @@ class DiffNote < Note
     position&.multiline?
   end
 
+  def shas
+    [
+      self.original_position.base_sha,
+      self.original_position.start_sha,
+      self.original_position.head_sha
+    ].tap do |a|
+      if self.position != self.original_position
+        a << self.position.base_sha
+        a << self.position.start_sha
+        a << self.position.head_sha
+      end
+    end
+  end
+
   private
 
   def enqueue_diff_file_creation_job
@@ -173,18 +187,6 @@ class DiffNote < Note
   end
 
   def keep_around_commits
-    shas = [
-      self.original_position.base_sha,
-      self.original_position.start_sha,
-      self.original_position.head_sha
-    ]
-
-    if self.position != self.original_position
-      shas << self.position.base_sha
-      shas << self.position.start_sha
-      shas << self.position.head_sha
-    end
-
     repository.keep_around(*shas)
   end
 

app/models/note.rb

@@ -48,6 +48,9 @@ class Note < ApplicationRecord
   # Attribute used to store the attributes that have been changed by quick actions.
   attr_accessor :commands_changes
 
+  # Attribute used to determine whether keep_around_commits will be skipped for diff notes.
+  attr_accessor :skip_keep_around_commits
+
   default_value_for :system, false
 
   attr_mentionable :note, pipeline: :note

app/services/draft_notes/publish_service.rb

@@ -34,22 +34,24 @@ module DraftNotes
 
       created_notes = draft_notes.map do |draft_note|
         draft_note.review = review
-        create_note_from_draft(draft_note, skip_capture_diff_note_position: true)
+        create_note_from_draft(draft_note, skip_capture_diff_note_position: true, skip_keep_around_commits: true)
       end
 
       capture_diff_note_positions(created_notes)
+      keep_around_commits(created_notes)
       draft_notes.delete_all
       set_reviewed
       notification_service.async.new_review(review)
       MergeRequests::ResolvedDiscussionNotificationService.new(project: project, current_user: current_user).execute(merge_request)
     end
 
-    def create_note_from_draft(draft, skip_capture_diff_note_position: false)
+    def create_note_from_draft(draft, skip_capture_diff_note_position: false, skip_keep_around_commits: false)
       # Make sure the diff file is unfolded in order to find the correct line
       # codes.
       draft.diff_file&.unfold_diff_lines(draft.original_position)
 
-      note = Notes::CreateService.new(draft.project, draft.author, draft.publish_params).execute(
+      note_params = draft.publish_params.merge(skip_keep_around_commits: skip_keep_around_commits)
+      note = Notes::CreateService.new(draft.project, draft.author, note_params).execute(
         skip_capture_diff_note_position: skip_capture_diff_note_position
       )
 
@@ -86,5 +88,17 @@ module DraftNotes
         capture_service.execute(note.discussion) if note.diff_note? && note.start_of_discussion?
       end
     end
+
+    def keep_around_commits(notes)
+      shas = notes.flat_map do |note|
+        note.shas if note.diff_note?
+      end.uniq
+
+      # We are allowing this since gitaly call will be created for each sha and
+      # even though they're unique, there will still be multiple Gitaly calls.
+      Gitlab::GitalyClient.allow_n_plus_1_calls do
+        project.repository.keep_around(*shas)
+      end
+    end
   end
 end

spec/models/diff_note_spec.rb

@@ -558,4 +558,31 @@ RSpec.describe DiffNote do
 
     it { is_expected.to eq('note') }
   end
+
+  describe '#shas' do
+    it 'returns list of SHAs based on original_position' do
+      expect(subject.shas).to match_array([
+        position.base_sha,
+        position.start_sha,
+        position.head_sha
+      ])
+    end
+
+    context 'when position changes' do
+      before do
+        subject.position = new_position
+      end
+
+      it 'includes the new position SHAs' do
+        expect(subject.shas).to match_array([
+          position.base_sha,
+          position.start_sha,
+          position.head_sha,
+          new_position.base_sha,
+          new_position.start_sha,
+          new_position.head_sha
+        ])
+      end
+    end
+  end
 end

spec/services/draft_notes/publish_service_spec.rb

@@ -33,7 +33,8 @@ RSpec.describe DraftNotes::PublishService do
     end
 
     it 'does not skip notification', :sidekiq_might_not_need_inline do
-      expect(Notes::CreateService).to receive(:new).with(project, user, drafts.first.publish_params).and_call_original
+      note_params = drafts.first.publish_params.merge(skip_keep_around_commits: false)
+      expect(Notes::CreateService).to receive(:new).with(project, user, note_params).and_call_original
       expect_next_instance_of(NotificationService) do |notification_service|
         expect(notification_service).to receive(:new_note)
       end
@@ -127,12 +128,17 @@ RSpec.describe DraftNotes::PublishService do
       publish
     end
 
-    context 'capturing diff notes positions' do
+    context 'capturing diff notes positions and keeping around commits' do
       before do
         # Need to execute this to ensure that we'll be able to test creation of
         # DiffNotePosition records as that only happens when the `MergeRequest#merge_ref_head`
         # is present. This service creates that for the specified merge request.
         MergeRequests::MergeToRefService.new(project: project, current_user: user).execute(merge_request)
+
+        # Need to re-stub this and call original as we are stubbing
+        # `Gitlab::Git::KeepAround#execute` in spec_helper for performance reason.
+        # Enabling it here so we can test the Gitaly calls it makes.
+        allow(Gitlab::Git::KeepAround).to receive(:execute).and_call_original
       end
 
       it 'creates diff_note_positions for diff notes' do
@@ -143,11 +149,26 @@ RSpec.describe DraftNotes::PublishService do
         expect(notes.last.diff_note_positions).to be_any
       end
 
+      it 'keeps around the commits of each published note' do
+        publish
+
+        repository = project.repository
+        notes = merge_request.notes.order(id: :asc)
+
+        notes.first.shas.each do |sha|
+          expect(repository.ref_exists?("refs/keep-around/#{sha}")).to be_truthy
+        end
+
+        notes.last.shas.each do |sha|
+          expect(repository.ref_exists?("refs/keep-around/#{sha}")).to be_truthy
+        end
+      end
+
       it 'does not requests a lot from Gitaly', :request_store do
         # NOTE: This should be reduced as we work on reducing Gitaly calls.
         # Gitaly requests shouldn't go above this threshold as much as possible
         # as it may add more to the Gitaly N+1 issue we are experiencing.
-        expect { publish }.to change { Gitlab::GitalyClient.get_request_count }.by(11)
+        expect { publish }.to change { Gitlab::GitalyClient.get_request_count }.by(21)
       end
     end