Cache diff highlight in discussions

9 files changed, 250 insertions, 4 deletions

diff --git a/app/controllers/projects/merge_requests_controller.rb b/app/controllers/projects/merge_requests_controller.rb
index da9316d5f22..87377efa3d3 100644
--- a/app/controllers/projects/merge_requests_controller.rb
+++ b/app/controllers/projects/merge_requests_controller.rb
@@ -220,6 +220,12 @@ class Projects::MergeRequestsController < Projects::MergeRequests::ApplicationCo
     head :ok
   end
 
+  def discussions
+    merge_request.discussions_diffs.preload_all
+
+    super
+  end
+
   protected
 
   alias_method :subscribable_resource, :merge_request
diff --git a/app/models/concerns/discussion_on_diff.rb b/app/models/concerns/discussion_on_diff.rb
index 266c37fa3a1..8e434053edc 100644
--- a/app/models/concerns/discussion_on_diff.rb
+++ b/app/models/concerns/discussion_on_diff.rb
@@ -9,7 +9,7 @@ module DiscussionOnDiff
   included do
     delegate  :line_code,
               :original_line_code,
-              :diff_file,
+              :note_diff_file,
               :diff_line,
               :active?,
               :created_at_diff?,
@@ -59,6 +59,14 @@ module DiscussionOnDiff
     prev_lines
   end
 
+  def diff_file
+    if for_merge_request? && context_noteable && note_diff_file
+      context_noteable.discussions_diffs.find_by_id(note_diff_file.id)
+    else
+      first_note.diff_file
+    end
+  end
+
   def line_code_in_diffs(diff_refs)
     if active?(diff_refs)
       line_code
diff --git a/app/models/merge_request.rb b/app/models/merge_request.rb
index 8052a54c504..5daf7eb0b29 100644
--- a/app/models/merge_request.rb
+++ b/app/models/merge_request.rb
@@ -409,6 +409,17 @@ class MergeRequest < ActiveRecord::Base
     merge_request_diffs.where.not(id: merge_request_diff.id)
   end
 
+  def discussions_diffs
+    strong_memoize(:discussions_diffs) do
+      diffs_collection = NoteDiffFile
+        .where(diff_note: discussion_notes)
+        .includes(diff_note: :project)
+        .to_a
+
+      Gitlab::DiscussionsDiff::FileCollection.new(diffs_collection)
+    end
+  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 27aef7adc48..3eded465657 100644
--- a/app/models/note_diff_file.rb
+++ b/app/models/note_diff_file.rb
@@ -3,6 +3,8 @@
 class NoteDiffFile < ActiveRecord::Base
   include DiffFile
 
+  delegate :original_position, :project, to: :diff_note
+
   belongs_to :diff_note, inverse_of: :note_diff_file
 
   validates :diff_note, presence: true
diff --git a/lib/gitlab/diff/file.rb b/lib/gitlab/diff/file.rb
index a85c5386d98..faedfab8ba0 100644
--- a/lib/gitlab/diff/file.rb
+++ b/lib/gitlab/diff/file.rb
@@ -3,7 +3,7 @@
 module Gitlab
   module Diff
     class File
-      attr_reader :diff, :repository, :diff_refs, :fallback_diff_refs
+      attr_reader :diff, :repository, :diff_refs, :fallback_diff_refs, :unique_identifier
 
       delegate :new_file?, :deleted_file?, :renamed_file?,
         :old_path, :new_path, :a_mode, :b_mode, :mode_changed?,
@@ -22,12 +22,20 @@ module Gitlab
         DiffViewer::Image
       ].sort_by { |v| v.binary? ? 0 : 1 }.freeze
 
-      def initialize(diff, repository:, diff_refs: nil, fallback_diff_refs: nil, stats: nil)
+      def initialize(
+        diff,
+        repository:,
+        diff_refs: nil,
+        fallback_diff_refs: nil,
+        stats: nil,
+        unique_identifier: nil)
+
         @diff = diff
         @stats = stats
         @repository = repository
         @diff_refs = diff_refs
         @fallback_diff_refs = fallback_diff_refs
+        @unique_identifier = unique_identifier
         @unfolded = false
 
         # Ensure items are collected in the the batch
@@ -67,7 +75,12 @@ module Gitlab
       def line_for_position(pos)
         return nil unless pos.position_type == 'text'
 
-        diff_lines.find { |line| line.old_line == pos.old_line && line.new_line == pos.new_line }
+        # The `pos` is more likely to be at the end of the Array
+        # given the persisted discussion notes are a diff hunk, not
+        # the entire diff.
+        diff_lines
+          .reverse_each
+          .find { |line| line.old_line == pos.old_line && line.new_line == pos.new_line }
       end
 
       def position_for_line_code(code)
diff --git a/lib/gitlab/discussions_diff/file_collection.rb b/lib/gitlab/discussions_diff/file_collection.rb
new file mode 100644
index 00000000000..384e9c2ba81
--- /dev/null
+++ b/lib/gitlab/discussions_diff/file_collection.rb
@@ -0,0 +1,91 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module DiscussionsDiff
+    class FileCollection
+      include Gitlab::Utils::StrongMemoize
+
+      def initialize(collection)
+        @collection = collection
+      end
+
+      def find_by_id(id)
+        diff_files_indexed_by_id[id]
+      end
+
+      def preload_all
+        write_empty_cache_keys
+        preload_highlighted_lines
+      end
+
+      private
+
+      def preload_highlighted_lines
+        cached_content = read_cache
+
+        diffs.zip(cached_content).each do |diff, content|
+          next unless content
+
+          # The data structure being loaded here is in control
+          # of the system and cannot be sploited by
+          # user data.
+          content = Marshal.load(content) # rubocop:disable Security/MarshalLoad
+
+          highlighted_diff_lines = content.map do |line|
+            Gitlab::Diff::Line.init_from_hash(line)
+          end
+
+          diff.highlighted_diff_lines = highlighted_diff_lines
+        end
+      end
+
+      def read_cache
+        strong_memoize(:read_all) do
+          HighlightCache.read_multiple(ids)
+        end
+      end
+
+      def write_empty_cache_keys
+        cached_content = read_cache
+
+        uncached_ids = ids.select.each_with_index { |_, i| cached_content[i].nil? }
+        mapping = highlighted_lines_by_id(uncached_ids)
+
+        HighlightCache.write_multiple(mapping)
+      end
+
+      def ids
+        @collection.map(&:id)
+      end
+
+      def diff_files_indexed_by_id
+        strong_memoize(:diff_files_indexed_by_id) do
+          diffs.index_by { |diff| diff.unique_identifier }
+        end
+      end
+
+      def diffs
+        strong_memoize(:diffs) do
+          @collection.map { |diff| decorate_diff(diff) }
+        end
+      end
+
+      def highlighted_lines_by_id(ids)
+        diff_files_indexed_by_id.select { |id, _| ids.include?(id) }.map do |id, file|
+          raw_content = file.highlighted_diff_lines.map(&:to_hash)
+
+          [id, Marshal.dump(raw_content)]
+        end
+      end
+
+      def decorate_diff(diff)
+        raw_diff = Gitlab::Git::Diff.new(diff.to_hash)
+
+        Gitlab::Diff::File.new(raw_diff,
+                               repository: diff.project.repository,
+                               diff_refs: diff.original_position.diff_refs,
+                               unique_identifier: diff.id)
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/discussions_diff/highlight_cache.rb b/lib/gitlab/discussions_diff/highlight_cache.rb
new file mode 100644
index 00000000000..85a129fdef3
--- /dev/null
+++ b/lib/gitlab/discussions_diff/highlight_cache.rb
@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+#
+module Gitlab
+  module DiscussionsDiff
+    class HighlightCache
+      class << self
+        # Sets multiple keys to a given value.
+        #
+        # mapping - A Hash mapping the cache keys to their values.
+        # Values are normally expected to be arrays of Hashes,
+        # therefore, Marshal is used.
+        def write_multiple(mapping)
+          Redis::Cache.with do |redis|
+            redis.multi do |multi|
+              mapping.each do |raw_key, value|
+                key = cache_key_for(raw_key)
+
+                multi.set(key, value, ex: 1.week)
+              end
+            end
+          end
+        end
+
+        # Reads multiple cache keys at once.
+        #
+        # raw_keys - An array of unique cache keys.
+        # Cache namespaces are not included.
+        def read_multiple(raw_keys)
+          return [] if raw_keys.empty?
+
+          keys = raw_keys.map { |id| cache_key_for(id) }
+
+          Redis::Cache.with do |redis|
+            redis.mget(keys)
+          end
+        end
+
+        def cache_key_for(raw_key)
+          "#{cache_key_prefix}:#{raw_key}"
+        end
+
+        private
+
+        def cache_key_prefix
+          "#{Redis::Cache::CACHE_NAMESPACE}:#{Diff::Line::SERIALIZE_KEYS}:discussion-highlight"
+        end
+      end
+    end
+  end
+end
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..193bcd34180
--- /dev/null
+++ b/spec/lib/gitlab/discussions_diff/file_collection_spec.rb
@@ -0,0 +1,35 @@
+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 }
+
+  let(:subject) { described_class.new([note_diff_file_a, note_diff_file_b]) }
+
+  describe '#preload_all', :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.preload_all
+    end
+
+    it 'preloaded diff files have @highlighted_diff_lines' do
+      subject.preload_all
+
+      diff_file = subject.find_by_id(note_diff_file_a.id)
+
+      expect(diff_file.instance_variable_get(:@highlighted_diff_lines))
+        .to all(be_a(Gitlab::Diff::Line))
+    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..c9c86139ac1
--- /dev/null
+++ b/spec/lib/gitlab/discussions_diff/highlight_cache_spec.rb
@@ -0,0 +1,30 @@
+require 'spec_helper'
+
+describe Gitlab::DiscussionsDiff::HighlightCache, :clean_gitlab_redis_cache do
+  describe '#write_multiple' do
+    it 'sets multiple keys' do
+      mapping = { 'foo' => 10, 'bar' => 20 }
+
+      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_s)
+      end
+    end
+  end
+
+  describe '#read_multiple' do
+    it 'reads multiple keys' do
+      mapping = { 'foo' => 10, 'bar' => 20 }
+
+      described_class.write_multiple(mapping)
+
+      found = described_class.read_multiple(mapping.keys)
+
+      expect(found).to eq(%w(10 20))
+    end
+  end
+end