Improve the performance for MR discussions loadosw-improve-discussions-query

It considerably improves the query for preloading MR
discussion diffs for:

1. MR diff comments (made at Diffs tab; presented at the main MR page)
2. Comments on commits being listed at the MR

Additionally, it also prevents a second similar query that was
unnecessarily being made.

10 files changed, 144 insertions, 84 deletions

diff --git a/app/controllers/projects/merge_requests_controller.rb b/app/controllers/projects/merge_requests_controller.rb
index ea1dd7d19d5..870e2ca7b8f 100644
--- a/app/controllers/projects/merge_requests_controller.rb
+++ b/app/controllers/projects/merge_requests_controller.rb
@@ -210,7 +210,7 @@ class Projects::MergeRequestsController < Projects::MergeRequests::ApplicationCo
   end
 
   def discussions
-    merge_request.preload_discussions_diff_highlight
+    merge_request.discussions_diffs.load_highlight
 
     super
   end
diff --git a/app/models/diff_note.rb b/app/models/diff_note.rb
index 861185dc222..bae442481af 100644
--- a/app/models/diff_note.rb
+++ b/app/models/diff_note.rb
@@ -108,13 +108,10 @@ class DiffNote < Note
   end
 
   def fetch_diff_file
+    return note_diff_file.raw_diff_file if note_diff_file
+
     file =
-      if note_diff_file
-        diff = Gitlab::Git::Diff.new(note_diff_file.to_hash)
-        Gitlab::Diff::File.new(diff,
-                               repository: repository,
-                               diff_refs: original_position.diff_refs)
-      elsif created_at_diff?(noteable.diff_refs)
+      if created_at_diff?(noteable.diff_refs)
         # We're able to use the already persisted diffs (Postgres) if we're
         # presenting a "current version" of the MR discussion diff.
         # So no need to make an extra Gitaly diff request for it.
@@ -126,9 +123,7 @@ class DiffNote < Note
         original_position.diff_file(repository)
       end
 
-    # Since persisted diff files already have its content "unfolded"
-    # there's no need to make it pass through the unfolding process.
-    file&.unfold_diff_lines(position) unless note_diff_file
+    file&.unfold_diff_lines(position)
 
     file
   end
diff --git a/app/models/merge_request.rb b/app/models/merge_request.rb
index 901ebcf249f..c550fc9fa8f 100644
--- a/app/models/merge_request.rb
+++ b/app/models/merge_request.rb
@@ -452,24 +452,17 @@ class MergeRequest < ApplicationRecord
     true
   end
 
-  def preload_discussions_diff_highlight
-    preloadable_files = note_diff_files.for_commit_or_unresolved
-
-    discussions_diffs.load_highlight(preloadable_files.pluck(:id))
-  end
-
   def discussions_diffs
     strong_memoize(:discussions_diffs) do
+      note_diff_files = NoteDiffFile
+        .joins(:diff_note)
+        .merge(notes.or(commit_notes))
+        .includes(diff_note: :project)
+
       Gitlab::DiscussionsDiff::FileCollection.new(note_diff_files.to_a)
     end
   end
 
-  def note_diff_files
-    NoteDiffFile
-      .where(diff_note: discussion_notes)
-      .includes(diff_note: :project)
-  end
-
   def diff_size
     # Calling `merge_request_diff.diffs.real_size` will also perform
     # highlighting, which we don't need here.
diff --git a/app/models/note_diff_file.rb b/app/models/note_diff_file.rb
index fcc9e2b3fd8..67a6d5d6d6b 100644
--- a/app/models/note_diff_file.rb
+++ b/app/models/note_diff_file.rb
@@ -3,15 +3,11 @@
 class NoteDiffFile < ApplicationRecord
   include DiffFile
 
-  scope :for_commit_or_unresolved, -> do
-    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
+  delegate :original_position, :project, :resolved_at, to: :diff_note
 
   belongs_to :diff_note, inverse_of: :note_diff_file
 
diff --git a/changelogs/unreleased/osw-improve-discussions-query.yml b/changelogs/unreleased/osw-improve-discussions-query.yml
new file mode 100644
index 00000000000..487bba458da
--- /dev/null
+++ b/changelogs/unreleased/osw-improve-discussions-query.yml
@@ -0,0 +1,5 @@
+---
+title: Considerably improve the query performance for MR discussions load
+merge_request: 32679
+author:
+type: performance
diff --git a/lib/gitlab/discussions_diff/file_collection.rb b/lib/gitlab/discussions_diff/file_collection.rb
index 4ab7314f509..6692dd76438 100644
--- a/lib/gitlab/discussions_diff/file_collection.rb
+++ b/lib/gitlab/discussions_diff/file_collection.rb
@@ -4,11 +4,16 @@ module Gitlab
   module DiscussionsDiff
     class FileCollection
       include Gitlab::Utils::StrongMemoize
+      include Enumerable
 
       def initialize(collection)
         @collection = collection
       end
 
+      def each(&block)
+        @collection.each(&block)
+      end
+
       # Returns a Gitlab::Diff::File with the given ID (`unique_identifier` in
       # Gitlab::Diff::File).
       def find_by_id(id)
@@ -16,20 +21,12 @@ module Gitlab
       end
 
       # Writes cache and preloads highlighted diff lines for
-      # object IDs, in @collection.
-      #
-      # highlightable_ids - Diff file `Array` responding to ID. The ID will be used
-      # to generate the cache key.
+      # highlightable object IDs, in @collection.
       #
       # - Highlight cache is written just for uncached diff files
       # - The cache content is not updated (there's no need to do so)
-      def load_highlight(highlightable_ids)
-        preload_highlighted_lines(highlightable_ids)
-      end
-
-      private
-
-      def preload_highlighted_lines(ids)
+      def load_highlight
+        ids = highlightable_collection_ids
         cached_content = read_cache(ids)
 
         uncached_ids = ids.select.each_with_index { |_, i| cached_content[i].nil? }
@@ -46,6 +43,12 @@ module Gitlab
         end
       end
 
+      private
+
+      def highlightable_collection_ids
+        each.with_object([]) { |file, memo| memo << file.id unless file.resolved_at }
+      end
+
       def read_cache(ids)
         HighlightCache.read_multiple(ids)
       end
@@ -57,9 +60,7 @@ module Gitlab
       end
 
       def diff_files
-        strong_memoize(:diff_files) do
-          @collection.map(&:raw_diff_file)
-        end
+        strong_memoize(:diff_files) { map(&:raw_diff_file) }
       end
 
       # Processes the diff lines highlighting for diff files matching the given
diff --git a/spec/controllers/projects/merge_requests_controller_spec.rb b/spec/controllers/projects/merge_requests_controller_spec.rb
index d0370dfaeee..58b15dda8d9 100644
--- a/spec/controllers/projects/merge_requests_controller_spec.rb
+++ b/spec/controllers/projects/merge_requests_controller_spec.rb
@@ -1255,7 +1255,7 @@ describe Projects::MergeRequestsController do
             expect_next_instance_of(Gitlab::DiscussionsDiff::FileCollection) do |collection|
               note_diff_file = commit_diff_note.note_diff_file
 
-              expect(collection).to receive(:load_highlight).with([note_diff_file.id]).and_call_original
+              expect(collection).to receive(:load_highlight).and_call_original
               expect(collection).to receive(:find_by_id).with(note_diff_file.id).and_call_original
             end
 
@@ -1272,7 +1272,7 @@ describe Projects::MergeRequestsController do
             expect_next_instance_of(Gitlab::DiscussionsDiff::FileCollection) do |collection|
               note_diff_file = diff_note.note_diff_file
 
-              expect(collection).to receive(:load_highlight).with([note_diff_file.id]).and_call_original
+              expect(collection).to receive(:load_highlight).and_call_original
               expect(collection).to receive(:find_by_id).with(note_diff_file.id).and_call_original
             end
 
@@ -1285,7 +1285,7 @@ describe Projects::MergeRequestsController do
             expect_next_instance_of(Gitlab::DiscussionsDiff::FileCollection) do |collection|
               note_diff_file = diff_note.note_diff_file
 
-              expect(collection).to receive(:load_highlight).with([]).and_call_original
+              expect(collection).to receive(:load_highlight).and_call_original
               expect(collection).to receive(:find_by_id).with(note_diff_file.id).and_call_original
             end
 
diff --git a/spec/lib/gitlab/discussions_diff/file_collection_spec.rb b/spec/lib/gitlab/discussions_diff/file_collection_spec.rb
index 0489206458b..6ef1e41450f 100644
--- a/spec/lib/gitlab/discussions_diff/file_collection_spec.rb
+++ b/spec/lib/gitlab/discussions_diff/file_collection_spec.rb
@@ -22,11 +22,13 @@ describe Gitlab::DiscussionsDiff::FileCollection do
                 note_diff_file_b.id => file_b_caching_content })
         .and_call_original
 
-      subject.load_highlight([note_diff_file_a.id, note_diff_file_b.id])
+      subject.load_highlight
     end
 
     it 'does not write cache for already cached file' do
-      subject.load_highlight([note_diff_file_a.id])
+      file_a_caching_content = diff_note_a.diff_file.highlighted_diff_lines.map(&:to_hash)
+      Gitlab::DiscussionsDiff::HighlightCache
+        .write_multiple({ note_diff_file_a.id => file_a_caching_content })
 
       file_b_caching_content = diff_note_b.diff_file.highlighted_diff_lines.map(&:to_hash)
 
@@ -35,27 +37,42 @@ describe Gitlab::DiscussionsDiff::FileCollection do
         .with({ note_diff_file_b.id => file_b_caching_content })
         .and_call_original
 
-      subject.load_highlight([note_diff_file_a.id, note_diff_file_b.id])
+      subject.load_highlight
     end
 
-    it 'does not err when given ID does not exist in @collection' do
-      expect { subject.load_highlight([999]) }.not_to raise_error
+    it 'does not write cache for resolved notes' do
+      diff_note_a.update_column(:resolved_at, Time.now)
+
+      file_b_caching_content = diff_note_b.diff_file.highlighted_diff_lines.map(&:to_hash)
+
+      expect(Gitlab::DiscussionsDiff::HighlightCache)
+        .to receive(:write_multiple)
+        .with({ note_diff_file_b.id => file_b_caching_content })
+        .and_call_original
+
+      subject.load_highlight
     end
 
     it 'loaded diff files have highlighted lines loaded' do
-      subject.load_highlight([note_diff_file_a.id])
+      subject.load_highlight
 
-      diff_file = subject.find_by_id(note_diff_file_a.id)
+      diff_file_a = subject.find_by_id(note_diff_file_a.id)
+      diff_file_b = subject.find_by_id(note_diff_file_b.id)
 
-      expect(diff_file.highlight_loaded?).to be(true)
+      expect(diff_file_a).to be_highlight_loaded
+      expect(diff_file_b).to be_highlight_loaded
     end
 
     it 'not loaded diff files does not have highlighted lines loaded' do
-      subject.load_highlight([note_diff_file_a.id])
+      diff_note_a.update_column(:resolved_at, Time.now)
+
+      subject.load_highlight
 
-      diff_file = subject.find_by_id(note_diff_file_b.id)
+      diff_file_a = subject.find_by_id(note_diff_file_a.id)
+      diff_file_b = subject.find_by_id(note_diff_file_b.id)
 
-      expect(diff_file.highlight_loaded?).to be(false)
+      expect(diff_file_a).not_to be_highlight_loaded
+      expect(diff_file_b).to be_highlight_loaded
     end
   end
 end
diff --git a/spec/models/merge_request_spec.rb b/spec/models/merge_request_spec.rb
index 11234982dd4..57bf1cbf933 100644
--- a/spec/models/merge_request_spec.rb
+++ b/spec/models/merge_request_spec.rb
@@ -652,9 +652,35 @@ describe MergeRequest do
     end
   end
 
-  describe '#preload_discussions_diff_highlight' do
+  describe '#discussions_diffs' do
     let(:merge_request) { create(:merge_request) }
 
+    shared_examples 'discussions diffs collection' do
+      it 'initializes Gitlab::DiscussionsDiff::FileCollection with correct data' do
+        note_diff_file = diff_note.note_diff_file
+
+        expect(Gitlab::DiscussionsDiff::FileCollection)
+          .to receive(:new)
+          .with([note_diff_file])
+          .and_call_original
+
+        result = merge_request.discussions_diffs
+
+        expect(result).to be_a(Gitlab::DiscussionsDiff::FileCollection)
+      end
+
+      it 'eager loads relations' do
+        result = merge_request.discussions_diffs
+
+        recorder = ActiveRecord::QueryRecorder.new do
+          result.first.diff_note
+          result.first.diff_note.project
+        end
+
+        expect(recorder.count).to be_zero
+      end
+    end
+
     context 'with commit diff note' do
       let(:other_merge_request) { create(:merge_request) }
 
@@ -666,40 +692,15 @@ describe MergeRequest do
         create(:diff_note_on_commit, project: other_merge_request.project)
       end
 
-      it 'preloads diff highlighting' do
-        expect_next_instance_of(Gitlab::DiscussionsDiff::FileCollection) do |collection|
-          note_diff_file = diff_note.note_diff_file
-
-          expect(collection)
-            .to receive(:load_highlight)
-            .with([note_diff_file.id]).and_call_original
-        end
-
-        merge_request.preload_discussions_diff_highlight
-      end
+      it_behaves_like 'discussions diffs collection'
     end
 
     context 'with merge request diff note' do
-      let!(:unresolved_diff_note) do
+      let!(:diff_note) do
         create(:diff_note_on_merge_request, project: merge_request.project, noteable: merge_request)
       end
 
-      let!(:resolved_diff_note) do
-        create(:diff_note_on_merge_request, :resolved, project: merge_request.project, noteable: merge_request)
-      end
-
-      it 'preloads diff highlighting' do
-        expect_next_instance_of(Gitlab::DiscussionsDiff::FileCollection) do |collection|
-          note_diff_file = unresolved_diff_note.note_diff_file
-
-          expect(collection)
-            .to receive(:load_highlight)
-            .with([note_diff_file.id])
-            .and_call_original
-        end
-
-        merge_request.preload_discussions_diff_highlight
-      end
+      it_behaves_like 'discussions diffs collection'
     end
   end
 
diff --git a/spec/requests/projects/merge_requests_discussions_spec.rb b/spec/requests/projects/merge_requests_discussions_spec.rb
new file mode 100644
index 00000000000..5945561aa7b
--- /dev/null
+++ b/spec/requests/projects/merge_requests_discussions_spec.rb
@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe 'merge requests discussions' do
+  # Further tests can be found at merge_requests_controller_spec.rb
+  describe 'GET /:namespace/:project/merge_requests/:iid/discussions' do
+    let(:project) { create(:project, :repository) }
+    let(:user) { project.owner }
+    let(:merge_request) { create(:merge_request_with_diffs, target_project: project, source_project: project) }
+
+    before do
+      project.add_developer(user)
+      login_as(user)
+    end
+
+    def send_request
+      get discussions_namespace_project_merge_request_path(namespace_id: project.namespace, project_id: project, id: merge_request.iid)
+    end
+
+    it 'returns 200' do
+      send_request
+
+      expect(response.status).to eq(200)
+    end
+
+    # https://docs.gitlab.com/ee/development/query_recorder.html#use-request-specs-instead-of-controller-specs
+    it 'avoids N+1 DB queries', :request_store do
+      control = ActiveRecord::QueryRecorder.new { send_request }
+
+      create(:diff_note_on_merge_request, noteable: merge_request,
+             project: merge_request.project)
+
+      expect do
+        send_request
+      end.not_to exceed_query_limit(control)
+    end
+
+    it 'limits Gitaly queries', :request_store do
+      Gitlab::GitalyClient.allow_n_plus_1_calls do
+        create_list(:diff_note_on_merge_request, 7, noteable: merge_request,
+                    project: merge_request.project)
+      end
+
+      # The creations above write into the Gitaly counts
+      Gitlab::GitalyClient.reset_counts
+
+      expect { send_request }
+        .to change { Gitlab::GitalyClient.get_request_count }.by_at_most(4)
+    end
+  end
+end