Cache diff highlight in discussions

This commit handles note diffs caching, which considerably improves
the performance on merge requests with lots of comments.
Important to note that the caching approach taken here is different
from `Gitlab::Diff::HighlightCache`. We do not reset the whole cache
when a new push is sent or anything else. That's because discussions
diffs are persisted and do not change.

4 files changed, 278 insertions, 0 deletions

diff --git a/spec/controllers/projects/merge_requests_controller_spec.rb b/spec/controllers/projects/merge_requests_controller_spec.rb
index 1ab227bde39..7c167420e1f 100644
--- a/spec/controllers/projects/merge_requests_controller_spec.rb
+++ b/spec/controllers/projects/merge_requests_controller_spec.rb
@@ -957,6 +957,70 @@ describe Projects::MergeRequestsController do
     end
   end
 
+  describe 'GET discussions' do
+    context 'when authenticated' do
+      before do
+        project.add_developer(user)
+        sign_in(user)
+      end
+
+      it 'returns 200' do
+        get :discussions, namespace_id: project.namespace, project_id: project, id: merge_request.iid
+
+        expect(response.status).to eq(200)
+      end
+
+      context 'highlight preloading' do
+        context 'with commit diff notes' do
+          let!(:commit_diff_note) do
+            create(:diff_note_on_commit, project: merge_request.project)
+          end
+
+          it 'preloads notes diffs highlights' 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(:find_by_id).with(note_diff_file.id).and_call_original
+            end
+
+            get :discussions, namespace_id: project.namespace, project_id: project, id: merge_request.iid
+          end
+        end
+
+        context 'with diff notes' do
+          let!(:diff_note) do
+            create(:diff_note_on_merge_request, noteable: merge_request, project: merge_request.project)
+          end
+
+          it 'preloads notes diffs highlights' 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(:find_by_id).with(note_diff_file.id).and_call_original
+            end
+
+            get :discussions, namespace_id: project.namespace, project_id: project, id: merge_request.iid
+          end
+
+          it 'does not preload highlights when diff note is resolved' do
+            Notes::ResolveService.new(diff_note.project, user).execute(diff_note)
+
+            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(:find_by_id).with(note_diff_file.id).and_call_original
+            end
+
+            get :discussions, namespace_id: project.namespace, project_id: project, id: merge_request.iid
+          end
+        end
+      end
+    end
+  end
+
   describe 'GET edit' do
     it 'responds successfully' do
       get :edit, params: { namespace_id: project.namespace, project_id: project, id: merge_request }
diff --git a/spec/lib/gitlab/discussions_diff/file_collection_spec.rb b/spec/lib/gitlab/discussions_diff/file_collection_spec.rb
new file mode 100644
index 00000000000..0489206458b
--- /dev/null
+++ b/spec/lib/gitlab/discussions_diff/file_collection_spec.rb
@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::DiscussionsDiff::FileCollection do
+  let(:merge_request) { create(:merge_request) }
+  let!(:diff_note_a) { create(:diff_note_on_merge_request, project: merge_request.project, noteable: merge_request) }
+  let!(:diff_note_b) { create(:diff_note_on_merge_request, project: merge_request.project, noteable: merge_request) }
+  let(:note_diff_file_a) { diff_note_a.note_diff_file }
+  let(:note_diff_file_b) { diff_note_b.note_diff_file }
+
+  subject { described_class.new([note_diff_file_a, note_diff_file_b]) }
+
+  describe '#load_highlight', :clean_gitlab_redis_shared_state do
+    it 'writes uncached diffs highlight' do
+      file_a_caching_content = diff_note_a.diff_file.highlighted_diff_lines.map(&:to_hash)
+      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_a.id => file_a_caching_content,
+                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])
+    end
+
+    it 'does not write cache for already cached file' do
+      subject.load_highlight([note_diff_file_a.id])
+
+      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([note_diff_file_a.id, note_diff_file_b.id])
+    end
+
+    it 'does not err when given ID does not exist in @collection' do
+      expect { subject.load_highlight([999]) }.not_to raise_error
+    end
+
+    it 'loaded diff files have highlighted lines loaded' do
+      subject.load_highlight([note_diff_file_a.id])
+
+      diff_file = subject.find_by_id(note_diff_file_a.id)
+
+      expect(diff_file.highlight_loaded?).to be(true)
+    end
+
+    it 'not loaded diff files does not have highlighted lines loaded' do
+      subject.load_highlight([note_diff_file_a.id])
+
+      diff_file = subject.find_by_id(note_diff_file_b.id)
+
+      expect(diff_file.highlight_loaded?).to be(false)
+    end
+  end
+end
diff --git a/spec/lib/gitlab/discussions_diff/highlight_cache_spec.rb b/spec/lib/gitlab/discussions_diff/highlight_cache_spec.rb
new file mode 100644
index 00000000000..fe26ebb8796
--- /dev/null
+++ b/spec/lib/gitlab/discussions_diff/highlight_cache_spec.rb
@@ -0,0 +1,102 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::DiscussionsDiff::HighlightCache, :clean_gitlab_redis_cache do
+  describe '#write_multiple' do
+    it 'sets multiple keys serializing content as JSON' do
+      mapping = {
+        3 => [
+          {
+            text: 'foo',
+            type: 'new',
+            index: 2,
+            old_pos: 10,
+            new_pos: 11,
+            line_code: 'xpto',
+            rich_text: '<blips>blops</blips>'
+          },
+          {
+            text: 'foo',
+            type: 'new',
+            index: 3,
+            old_pos: 11,
+            new_pos: 12,
+            line_code: 'xpto',
+            rich_text: '<blops>blips</blops>'
+          }
+        ]
+      }
+
+      described_class.write_multiple(mapping)
+
+      mapping.each do |key, value|
+        full_key = described_class.cache_key_for(key)
+        found = Gitlab::Redis::Cache.with { |r| r.get(full_key) }
+
+        expect(found).to eq(value.to_json)
+      end
+    end
+  end
+
+  describe '#read_multiple' do
+    it 'reads multiple keys and serializes content into Gitlab::Diff::Line objects' do
+      mapping = {
+        3 => [
+          {
+            text: 'foo',
+            type: 'new',
+            index: 2,
+            old_pos: 11,
+            new_pos: 12,
+            line_code: 'xpto',
+            rich_text: '<blips>blops</blips>'
+          },
+          {
+            text: 'foo',
+            type: 'new',
+            index: 3,
+            old_pos: 10,
+            new_pos: 11,
+            line_code: 'xpto',
+            rich_text: '<blips>blops</blips>'
+          }
+        ]
+      }
+
+      described_class.write_multiple(mapping)
+
+      found = described_class.read_multiple(mapping.keys)
+
+      expect(found.size).to eq(1)
+      expect(found.first.size).to eq(2)
+      expect(found.first).to all(be_a(Gitlab::Diff::Line))
+    end
+
+    it 'returns nil when cached key is not found' do
+      mapping = {
+        3 => [
+          {
+            text: 'foo',
+            type: 'new',
+            index: 2,
+            old_pos: 11,
+            new_pos: 12,
+            line_code: 'xpto',
+            rich_text: '<blips>blops</blips>'
+          }
+        ]
+      }
+
+      described_class.write_multiple(mapping)
+
+      found = described_class.read_multiple([2, 3])
+
+      expect(found.size).to eq(2)
+
+      expect(found.first).to eq(nil)
+      expect(found.second.size).to eq(1)
+      expect(found.second).to all(be_a(Gitlab::Diff::Line))
+    end
+  end
+end
diff --git a/spec/models/merge_request_spec.rb b/spec/models/merge_request_spec.rb
index 6793d4e8718..4cc3a6a3644 100644
--- a/spec/models/merge_request_spec.rb
+++ b/spec/models/merge_request_spec.rb
@@ -559,6 +559,57 @@ describe MergeRequest do
     end
   end
 
+  describe '#preload_discussions_diff_highlight' do
+    let(:merge_request) { create(:merge_request) }
+
+    context 'with commit diff note' do
+      let(:other_merge_request) { create(:merge_request) }
+
+      let!(:diff_note) do
+        create(:diff_note_on_commit, project: merge_request.project)
+      end
+
+      let!(:other_mr_diff_note) 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
+    end
+
+    context 'with merge request diff note' do
+      let!(:unresolved_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
+    end
+  end
+
   describe '#diff_size' do
     let(:merge_request) do
       build(:merge_request, source_branch: 'expand-collapse-files', target_branch: 'master')