Cache diff highlighting upon Merge Request creation (refactors diff caching)

app/controllers/projects/merge_requests/diffs_controller.rb

@@ -21,6 +21,8 @@ class Projects::MergeRequests::DiffsController < Projects::MergeRequests::Applic
   def render_diffs
     @environment = @merge_request.environments_for(current_user).last
+    @diffs.write_cache
+
     render json: DiffsSerializer.new(current_user: current_user).represent(@diffs, additional_attributes)
   end

app/services/merge_requests/reload_diffs_service.rb

@@ -30,7 +30,7 @@ module MergeRequests
     def clear_cache(new_diff)
       # Executing the iteration we cache highlighted diffs for each diff file of
       # MergeRequestDiff.
-      new_diff.diffs_collection.diff_files.to_a
+      new_diff.diffs_collection.write_cache
       # Remove cache for all diffs on this MR. Do not use the association on the
       # model, as that will interfere with other actions happening when
@@ -38,7 +38,7 @@ module MergeRequests
       MergeRequestDiff.where(merge_request: merge_request).each do |merge_request_diff|
         next if merge_request_diff == new_diff
-        merge_request_diff.diffs_collection.clear_cache!
+        merge_request_diff.diffs_collection.clear_cache
       end
     end
   end

app/workers/new_merge_request_worker.rb

@@ -9,6 +9,8 @@ class NewMergeRequestWorker
     EventCreateService.new.open_mr(issuable, user)
     NotificationService.new.new_merge_request(issuable, user)
+
+    issuable.diffs.write_cache
     issuable.create_cross_references!(user)
   end

changelogs/unreleased/osw-write-cache-upon-mr-creation-and-cache-refactoring.yml

+---
+title: Write diff highlighting cache upon MR creation (refactors caching)
+merge_request: 21489
+author:
+type: performance

lib/gitlab/diff/file_collection/base.rb

@@ -2,7 +2,7 @@ module Gitlab
   module Diff
     module FileCollection
       class Base
-        attr_reader :project, :diff_options, :diff_refs, :fallback_diff_refs
+        attr_reader :project, :diff_options, :diff_refs, :fallback_diff_refs, :diffable
         delegate :count, :size, :real_size, to: :diff_files
@@ -33,6 +33,14 @@ module Gitlab
           diff_files.find { |diff_file| diff_file.new_path == new_path }
         end
+        def clear_cache
+          # No-op
+        end
+
+        def write_cache
+          # No-op
+        end
+
         private
         def decorate_diff!(diff)

lib/gitlab/diff/file_collection/merge_request_diff.rb

@@ -2,6 +2,8 @@ module Gitlab
   module Diff
     module FileCollection
       class MergeRequestDiff < Base
+        extend ::Gitlab::Utils::Override
+
         def initialize(merge_request_diff, diff_options:)
           @merge_request_diff = merge_request_diff
@@ -13,70 +15,35 @@ module Gitlab
         end
         def diff_files
-          # Make sure to _not_ send any method call to Gitlab::Diff::File
-          # _before_ all of them were collected (`super`). Premature method calls will
-          # trigger N+1 RPCs to Gitaly through BatchLoader records (Blob.lazy).
-          #
           diff_files = super
-          diff_files.each { |diff_file| cache_highlight!(diff_file) if cacheable?(diff_file) }
-          store_highlight_cache
+          diff_files.each { |diff_file| cache.decorate(diff_file) }
           diff_files
         end
-        def real_size
-          @merge_request_diff.real_size
+        override :write_cache
+        def write_cache
+          cache.write_if_empty
         end
-        def clear_cache!
-          Rails.cache.delete(cache_key)
+        override :clear_cache
+        def clear_cache
+          cache.clear
         end
         def cache_key
-          [@merge_request_diff, 'highlighted-diff-files', Gitlab::Diff::Line::SERIALIZE_KEYS, diff_options]
-        end
-
-        private
-
-        def highlight_diff_file_from_cache!(diff_file, cache_diff_lines)
-          diff_file.highlighted_diff_lines = cache_diff_lines.map do |line|
-            Gitlab::Diff::Line.init_from_hash(line)
-          end
+          cache.key
         end
-        #
-        # If we find the highlighted diff files lines on the cache we replace existing diff_files lines (no highlighted)
-        # for the highlighted ones, so we just skip their execution.
-        # If the highlighted diff files lines are not cached we calculate and cache them.
-        #
-        # The content of the cache is a Hash where the key identifies the file and the values are Arrays of
-        # hashes that represent serialized diff lines.
-        #
-        def cache_highlight!(diff_file)
-          item_key = diff_file.file_identifier
-
-          if highlight_cache[item_key]
-            highlight_diff_file_from_cache!(diff_file, highlight_cache[item_key])
-          else
-            highlight_cache[item_key] = diff_file.highlighted_diff_lines.map(&:to_hash)
-          end
-        end
-
-        def highlight_cache
-          return @highlight_cache if defined?(@highlight_cache)
-
-          @highlight_cache = Rails.cache.read(cache_key) || {}
-          @highlight_cache_was_empty = @highlight_cache.empty?
-          @highlight_cache
+        def real_size
+          @merge_request_diff.real_size
         end
-        def store_highlight_cache
-          Rails.cache.write(cache_key, highlight_cache, expires_in: 1.week) if @highlight_cache_was_empty
-        end
-        def cacheable?(diff_file)
-          @merge_request_diff.present? && diff_file.text? && diff_file.diffable?
+        private
+        def cache
+          @cache ||= Gitlab::Diff::HighlightCache.new(self)
         end
       end
     end

lib/gitlab/diff/highlight_cache.rb

+# frozen_string_literal: true
+#
+module Gitlab
+  module Diff
+    class HighlightCache
+      delegate :diffable, to: :@diff_collection
+      delegate :diff_options, to: :@diff_collection
+
+      def initialize(diff_collection, backend: Rails.cache)
+        @backend = backend
+        @diff_collection = diff_collection
+      end
+
+      # - Reads from cache
+      # - Assigns DiffFile#highlighted_diff_lines for cached files
+      def decorate(diff_file)
+        if content = read_file(diff_file)
+          diff_file.highlighted_diff_lines = content.map do |line|
+            Gitlab::Diff::Line.init_from_hash(line)
+          end
+        end
+      end
+
+      # It populates a Hash in order to submit a single write to the memory
+      # cache. This avoids excessive IO generated by N+1's (1 writing for
+      # each highlighted line or file).
+      def write_if_empty
+        return if cached_content.present?
+
+        @diff_collection.diff_files.each do |diff_file|
+          next unless cacheable?(diff_file)
+
+          diff_file_id = diff_file.file_identifier
+
+          cached_content[diff_file_id] = diff_file.highlighted_diff_lines.map(&:to_hash)
+        end
+
+        cache.write(key, cached_content, expires_in: 1.week)
+      end
+
+      def clear
+        cache.delete(key)
+      end
+
+      def key
+        [diffable, 'highlighted-diff-files', Gitlab::Diff::Line::SERIALIZE_KEYS, diff_options]
+      end
+
+      private
+
+      def read_file(diff_file)
+        cached_content[diff_file.file_identifier]
+      end
+
+      def cache
+        @backend
+      end
+
+      def cached_content
+        @cached_content ||= cache.read(key) || {}
+      end
+
+      def cacheable?(diff_file)
+        diffable.present? && diff_file.text? && diff_file.diffable?
+      end
+    end
+  end
+end

spec/lib/gitlab/diff/highlight_cache_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::Diff::HighlightCache do
+  let(:merge_request) { create(:merge_request_with_diffs) }
+
+  subject(:cache) { described_class.new(merge_request.diffs, backend: backend) }
+
+  describe '#decorate' do
+    let(:backend) { double('backend').as_null_object }
+
+    # Manually creates a Diff::File object to avoid triggering the cache on
+    # the FileCollection::MergeRequestDiff
+    let(:diff_file) do
+      diffs = merge_request.diffs
+      raw_diff = diffs.diffable.raw_diffs(diffs.diff_options.merge(paths: ['CHANGELOG'])).first
+      Gitlab::Diff::File.new(raw_diff,
+                             repository: diffs.project.repository,
+                             diff_refs: diffs.diff_refs,
+                             fallback_diff_refs: diffs.fallback_diff_refs)
+    end
+
+    it 'does not calculate highlighting when reading from cache' do
+      cache.write_if_empty
+      cache.decorate(diff_file)
+
+      expect_any_instance_of(Gitlab::Diff::Highlight).not_to receive(:highlight)
+
+      diff_file.highlighted_diff_lines
+    end
+
+    it 'assigns highlighted diff lines to the DiffFile' do
+      cache.write_if_empty
+      cache.decorate(diff_file)
+
+      expect(diff_file.highlighted_diff_lines.size).to be > 5
+    end
+
+    it 'submits a single reading from the cache' do
+      cache.decorate(diff_file)
+      cache.decorate(diff_file)
+
+      expect(backend).to have_received(:read).with(cache.key).once
+    end
+  end
+
+  describe '#write_if_empty' do
+    let(:backend) { double('backend', read: {}).as_null_object }
+
+    it 'submits a single writing to the cache' do
+      cache.write_if_empty
+      cache.write_if_empty
+
+      expect(backend).to have_received(:write).with(cache.key,
+                                                    hash_including('CHANGELOG-false-false-false'),
+                                                    expires_in: 1.week).once
+    end
+  end
+
+  describe '#clear' do
+    let(:backend) { double('backend').as_null_object }
+
+    it 'clears cache' do
+      cache.clear
+
+      expect(backend).to have_received(:delete).with(cache.key)
+    end
+  end
+end

spec/services/merge_requests/reload_diffs_service_spec.rb

@@ -57,6 +57,7 @@ describe MergeRequests::ReloadDiffsService, :use_clean_rails_memory_store_cachin
         expect(Rails.cache).to receive(:delete).with(old_cache_key).and_call_original
         expect(Rails.cache).to receive(:read).with(new_cache_key).and_call_original
         expect(Rails.cache).to receive(:write).with(new_cache_key, anything, anything).and_call_original
+
         subject.execute
       end
     end