Use diff cache for diffs batch loading

It uses the existing cache for the
MergeRequestDiffBatch collection class.

app/controllers/projects/merge_requests/diffs_controller.rb

@@ -28,6 +28,7 @@ class Projects::MergeRequests::DiffsController < Projects::MergeRequests::Applic
     positions = @merge_request.note_positions_for_paths(diffs.diff_file_paths, current_user)
 
     diffs.unfold_diff_files(positions.unfoldable)
+    diffs.write_cache
 
     options = {
       merge_request: @merge_request,

lib/gitlab/diff/file_collection/merge_request_diff.rb

@@ -4,45 +4,6 @@ module Gitlab
   module Diff
     module FileCollection
       class MergeRequestDiff < MergeRequestDiffBase
-        include Gitlab::Utils::StrongMemoize
-
-        def diff_files
-          strong_memoize(:diff_files) do
-            diff_files = super
-
-            diff_files.each { |diff_file| cache.decorate(diff_file) }
-
-            diff_files
-          end
-        end
-
-        override :write_cache
-        def write_cache
-          cache.write_if_empty
-        end
-
-        override :clear_cache
-        def clear_cache
-          cache.clear
-        end
-
-        def cache_key
-          cache.key
-        end
-
-        def real_size
-          @merge_request_diff.real_size
-        end
-
-        private
-
-        def cache
-          @cache ||= if Feature.enabled?(:hset_redis_diff_caching, project)
-                       Gitlab::Diff::HighlightCache.new(self)
-                     else
-                       Gitlab::Diff::DeprecatedHighlightCache.new(self)
-                     end
-        end
       end
     end
   end

lib/gitlab/diff/file_collection/merge_request_diff_base.rb

@@ -15,6 +15,44 @@ module Gitlab
             diff_refs: merge_request_diff.diff_refs,
             fallback_diff_refs: merge_request_diff.fallback_diff_refs)
         end
+
+        def diff_files
+          strong_memoize(:diff_files) do
+            diff_files = super
+
+            diff_files.each { |diff_file| cache.decorate(diff_file) }
+
+            diff_files
+          end
+        end
+
+        override :write_cache
+        def write_cache
+          cache.write_if_empty
+        end
+
+        override :clear_cache
+        def clear_cache
+          cache.clear
+        end
+
+        def cache_key
+          cache.key
+        end
+
+        def real_size
+          @merge_request_diff.real_size
+        end
+
+        private
+
+        def cache
+          @cache ||= if Feature.enabled?(:hset_redis_diff_caching, project)
+                       Gitlab::Diff::HighlightCache.new(self)
+                     else
+                       Gitlab::Diff::DeprecatedHighlightCache.new(self)
+                     end
+        end
       end
     end
   end

spec/controllers/projects/merge_requests/diffs_controller_spec.rb

@@ -25,6 +25,16 @@ describe Projects::MergeRequests::DiffsController do
     end
   end
 
+  shared_examples 'cached diff collection' do
+    it 'ensures diff highlighting cache writing' do
+      expect_next_instance_of(Gitlab::Diff::HighlightCache) do |cache|
+        expect(cache).to receive(:write_if_empty).once
+      end
+
+      go
+    end
+  end
+
   shared_examples 'persisted preferred diff view cookie' do
     context 'with view param' do
       before do
@@ -112,6 +122,7 @@ describe Projects::MergeRequests::DiffsController do
     end
 
     it_behaves_like 'persisted preferred diff view cookie'
+    it_behaves_like 'cached diff collection'
   end
 
   describe 'GET diffs_metadata' do
@@ -404,6 +415,7 @@ describe Projects::MergeRequests::DiffsController do
 
     it_behaves_like 'forked project with submodules'
     it_behaves_like 'persisted preferred diff view cookie'
+    it_behaves_like 'cached diff collection'
 
     context 'diff unfolding' do
       let!(:unfoldable_diff_note) do

spec/lib/gitlab/diff/file_collection/merge_request_diff_batch_spec.rb

@@ -123,4 +123,8 @@ describe Gitlab::Diff::FileCollection::MergeRequestDiffBatch do
                           collection_default_args)
     end
   end
+
+  it_behaves_like 'cacheable diff collection' do
+    let(:cacheable_files_count) { batch_size }
+  end
 end

spec/lib/gitlab/diff/file_collection/merge_request_diff_spec.rb

@@ -4,7 +4,8 @@ require 'spec_helper'
 
 describe Gitlab::Diff::FileCollection::MergeRequestDiff do
   let(:merge_request) { create(:merge_request) }
-  let(:subject) { described_class.new(merge_request.merge_request_diff, diff_options: nil) }
+  let(:diffable) { merge_request.merge_request_diff }
+  let(:subject) { described_class.new(diffable, diff_options: nil) }
   let(:diff_files) { subject.diff_files }
 
   describe '#diff_files' do
@@ -52,6 +53,10 @@ describe Gitlab::Diff::FileCollection::MergeRequestDiff do
     let(:stub_path) { '.gitignore' }
   end
 
+  it_behaves_like 'cacheable diff collection' do
+    let(:cacheable_files_count) { diffable.size.to_i }
+  end
+
   it 'returns a valid instance of a DiffCollection' do
     expect(diff_files).to be_a(Gitlab::Git::DiffCollection)
   end

spec/support/shared_examples/diff_file_collections.rb

@@ -57,3 +57,45 @@ shared_examples 'unfoldable diff' do
     subject.unfold_diff_files([position])
   end
 end
+
+shared_examples 'cacheable diff collection' do
+  let(:cache) { instance_double(Gitlab::Diff::HighlightCache) }
+
+  before do
+    expect(Gitlab::Diff::HighlightCache).to receive(:new).with(subject) { cache }
+  end
+
+  describe '#write_cache' do
+    it 'calls Gitlab::Diff::HighlightCache#write_if_empty' do
+      expect(cache).to receive(:write_if_empty).once
+
+      subject.write_cache
+    end
+  end
+
+  describe '#clear_cache' do
+    it 'calls Gitlab::Diff::HighlightCache#clear' do
+      expect(cache).to receive(:clear).once
+
+      subject.clear_cache
+    end
+  end
+
+  describe '#cache_key' do
+    it 'calls Gitlab::Diff::HighlightCache#key' do
+      expect(cache).to receive(:key).once
+
+      subject.cache_key
+    end
+  end
+
+  describe '#diff_files' do
+    it 'calls Gitlab::Diff::HighlightCache#decorate' do
+      expect(cache).to receive(:decorate)
+        .with(instance_of(Gitlab::Diff::File))
+        .exactly(cacheable_files_count).times
+
+      subject.diff_files
+    end
+  end
+end