Reduce Gitaly calls for keeping around refs of published notes

We are currently keeping around commits of each note being
published even though they are all the same commits. That results
to duplicate `ref_exists` calls.

For example, if there are 6 draft notes set to be published on a
single MR version, there will be around 36 calls for `ref_exists`.

This fix reduces those calls by getting the unique shas to be kept.
Based on local testing, this reduced the `ref_exists` calls from
36 to 4.

Changelog: performance

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