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.

2 files changed, 143 insertions, 0 deletions

diff --git a/lib/gitlab/discussions_diff/file_collection.rb b/lib/gitlab/discussions_diff/file_collection.rb
new file mode 100644
index 00000000000..4ab7314f509
--- /dev/null
+++ b/lib/gitlab/discussions_diff/file_collection.rb
@@ -0,0 +1,76 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module DiscussionsDiff
+    class FileCollection
+      include Gitlab::Utils::StrongMemoize
+
+      def initialize(collection)
+        @collection = collection
+      end
+
+      # Returns a Gitlab::Diff::File with the given ID (`unique_identifier` in
+      # Gitlab::Diff::File).
+      def find_by_id(id)
+        diff_files_indexed_by_id[id]
+      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.
+      #
+      # - 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)
+        cached_content = read_cache(ids)
+
+        uncached_ids = ids.select.each_with_index { |_, i| cached_content[i].nil? }
+        mapping = highlighted_lines_by_ids(uncached_ids)
+
+        HighlightCache.write_multiple(mapping)
+
+        diffs = diff_files_indexed_by_id.values_at(*ids)
+
+        diffs.zip(cached_content).each do |diff, cached_lines|
+          next unless diff && cached_lines
+
+          diff.highlighted_diff_lines = cached_lines
+        end
+      end
+
+      def read_cache(ids)
+        HighlightCache.read_multiple(ids)
+      end
+
+      def diff_files_indexed_by_id
+        strong_memoize(:diff_files_indexed_by_id) do
+          diff_files.index_by(&:unique_identifier)
+        end
+      end
+
+      def diff_files
+        strong_memoize(:diff_files) do
+          @collection.map(&:raw_diff_file)
+        end
+      end
+
+      # Processes the diff lines highlighting for diff files matching the given
+      # IDs.
+      #
+      # Returns a Hash with { id => [Array of Gitlab::Diff::line], ...]
+      def highlighted_lines_by_ids(ids)
+        diff_files_indexed_by_id.slice(*ids).each_with_object({}) do |(id, file), hash|
+          hash[id] = file.highlighted_diff_lines.map(&:to_hash)
+        end
+      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..270cfb89488
--- /dev/null
+++ b/lib/gitlab/discussions_diff/highlight_cache.rb
@@ -0,0 +1,67 @@
+# frozen_string_literal: true
+#
+module Gitlab
+  module DiscussionsDiff
+    class HighlightCache
+      class << self
+        VERSION = 1
+        EXPIRATION = 1.week
+
+        # Sets multiple keys to a given value. The value
+        # is serialized as JSON.
+        #
+        # mapping - Write multiple cache values at once
+        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.to_json, ex: EXPIRATION)
+              end
+            end
+          end
+        end
+
+        # Reads multiple cache keys at once.
+        #
+        # raw_keys - An Array of unique cache keys, without namespaces.
+        #
+        # It returns a list of deserialized diff lines. Ex.:
+        # [[Gitlab::Diff::Line, ...], [Gitlab::Diff::Line]]
+        def read_multiple(raw_keys)
+          return [] if raw_keys.empty?
+
+          keys = raw_keys.map { |id| cache_key_for(id) }
+
+          content =
+            Redis::Cache.with do |redis|
+              redis.mget(keys)
+            end
+
+          content.map! do |lines|
+            next unless lines
+
+            JSON.parse(lines).map! do |line|
+              line = line.with_indifferent_access
+              rich_text = line[:rich_text]
+              line[:rich_text] = rich_text&.html_safe
+
+              Gitlab::Diff::Line.init_from_hash(line)
+            end
+          end
+        end
+
+        def cache_key_for(raw_key)
+          "#{cache_key_prefix}:#{raw_key}"
+        end
+
+        private
+
+        def cache_key_prefix
+          "#{Redis::Cache::CACHE_NAMESPACE}:#{VERSION}:discussion-highlight"
+        end
+      end
+    end
+  end
+end