Do not include stats when calling diff_batches/metadata

This call was making a gitaly call, which was causing signifcant
performance impact. We use these stats in the file, but they can
be replaced by iterating over the files instead.

Issue: gitlab.com/gitlab-org/gitlab/-/issues/209786
Merge Request: gitlab.com/gitlab-org/gitlab/-/merge_requests/33037

changelogs/unreleased/209786-improve-performance-of-diffs_batch-and-diffs_metadata-actions-for-.yml

+---
+title: Improve the performance for loading large diffs on a Merge request
+merge_request: 33037
+author:
+type: performance

lib/gitlab/diff/file.rb

@@ -230,12 +230,16 @@ module Gitlab
       end
 
       def added_lines
+        strong_memoize(:added_lines) do
           @stats&.additions || diff_lines.count(&:added?)
         end
+      end
 
       def removed_lines
+        strong_memoize(:removed_lines) do
           @stats&.deletions || diff_lines.count(&:removed?)
         end
+      end
 
       def file_identifier
         "#{file_path}-#{new_file?}-#{deleted_file?}-#{renamed_file?}"

lib/gitlab/diff/file_collection/base.rb

@@ -88,15 +88,18 @@ module Gitlab
 
         def diff_stats_collection
           strong_memoize(:diff_stats) do
-            # There are scenarios where we don't need to request Diff Stats,
-            # when caching for instance.
-            next unless @include_stats
-            next unless diff_refs
+            next unless fetch_diff_stats?
 
             @repository.diff_stats(diff_refs.base_sha, diff_refs.head_sha)
           end
         end
 
+        def fetch_diff_stats?
+          # There are scenarios where we don't need to request Diff Stats,
+          # when caching for instance.
+          @include_stats && diff_refs
+        end
+
         def decorate_diff!(diff)
           return diff if diff.is_a?(File)
 

lib/gitlab/diff/file_collection/merge_request_diff_base.rb

@@ -20,7 +20,7 @@ module Gitlab
           strong_memoize(:diff_files) do
             diff_files = super
 
-            diff_files.each { |diff_file| cache.decorate(diff_file) }
+            diff_files.each { |diff_file| highlight_cache.decorate(diff_file) }
 
             diff_files
           end
@@ -28,16 +28,14 @@ module Gitlab
 
         override :write_cache
         def write_cache
-          cache.write_if_empty
+          highlight_cache.write_if_empty
+          diff_stats_cache&.write_if_empty(diff_stats_collection)
         end
 
         override :clear_cache
         def clear_cache
-          cache.clear
-        end
-
-        def cache_key
-          cache.key
+          highlight_cache.clear
+          diff_stats_cache&.clear
         end
 
         def real_size
@@ -46,8 +44,27 @@ module Gitlab
 
         private
 
-        def cache
-          @cache ||= Gitlab::Diff::HighlightCache.new(self)
+        def highlight_cache
+          strong_memoize(:highlight_cache) do
+            Gitlab::Diff::HighlightCache.new(self)
+          end
+        end
+
+        def diff_stats_cache
+          strong_memoize(:diff_stats_cache) do
+            if Feature.enabled?(:cache_diff_stats_merge_request, project)
+              Gitlab::Diff::StatsCache.new(cachable_key: @merge_request_diff.cache_key)
+            end
+          end
+        end
+
+        override :diff_stats_collection
+        def diff_stats_collection
+          strong_memoize(:diff_stats) do
+            next unless fetch_diff_stats?
+
+            diff_stats_cache&.read || super
+          end
         end
       end
     end

lib/gitlab/diff/stats_cache.rb

+# frozen_string_literal: true
+#
+module Gitlab
+  module Diff
+    class StatsCache
+      include Gitlab::Metrics::Methods
+      include Gitlab::Utils::StrongMemoize
+
+      EXPIRATION = 1.week
+      VERSION = 1
+
+      def initialize(cachable_key:)
+        @cachable_key = cachable_key
+      end
+
+      def read
+        strong_memoize(:cached_values) do
+          content = cache.fetch(key)
+
+          next unless content
+
+          stats = content.map { |stat| Gitaly::DiffStats.new(stat) }
+
+          Gitlab::Git::DiffStatsCollection.new(stats)
+        end
+      end
+
+      def write_if_empty(stats)
+        return if cache.exist?(key)
+        return unless stats
+
+        cache.write(key, stats.as_json, expires_in: EXPIRATION)
+      end
+
+      def clear
+        cache.delete(key)
+      end
+
+      private
+
+      attr_reader :cachable_key
+
+      def cache
+        Rails.cache
+      end
+
+      def key
+        strong_memoize(:redis_key) do
+          ['diff_stats', cachable_key, VERSION].join(":")
+        end
+      end
+    end
+  end
+end

spec/lib/gitlab/diff/stats_cache_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::Diff::StatsCache, :use_clean_rails_memory_store_caching do
+  subject(:stats_cache) { described_class.new(cachable_key: cachable_key) }
+
+  let(:key) { ['diff_stats', cachable_key, described_class::VERSION].join(":") }
+  let(:cachable_key) { 'cachecachecache' }
+  let(:stat) { Gitaly::DiffStats.new(path: 'temp', additions: 10, deletions: 15) }
+  let(:stats) { Gitlab::Git::DiffStatsCollection.new([stat]) }
+  let(:cache) { Rails.cache }
+
+  describe '#read' do
+    before do
+      stats_cache.write_if_empty(stats)
+    end
+
+    it 'returns the expected stats' do
+      expect(stats_cache.read.to_json).to eq(stats.to_json)
+    end
+  end
+
+  describe '#write_if_empty' do
+    context 'when the cache already exists' do
+      before do
+        Rails.cache.write(key, true)
+      end
+
+      it 'does not write the stats' do
+        expect(cache).not_to receive(:write)
+
+        stats_cache.write_if_empty(stats)
+      end
+    end
+
+    context 'when the cache does not exist' do
+      it 'writes the stats' do
+        expect(cache)
+          .to receive(:write)
+          .with(key, stats.as_json, expires_in: described_class::EXPIRATION)
+          .and_call_original
+
+        stats_cache.write_if_empty(stats)
+
+        expect(stats_cache.read.to_a).to eq(stats.to_a)
+      end
+
+      context 'when given non utf-8 characters' do
+        let(:non_utf8_path) { '你好'.b }
+        let(:stat) { Gitaly::DiffStats.new(path: non_utf8_path, additions: 10, deletions: 15) }
+
+        it 'writes the stats' do
+          expect(cache)
+            .to receive(:write)
+            .with(key, stats.as_json, expires_in: described_class::EXPIRATION)
+            .and_call_original
+
+          stats_cache.write_if_empty(stats)
+
+          expect(stats_cache.read.to_a).to eq(stats.to_a)
+        end
+      end
+
+      context 'when given empty stats' do
+        let(:stats) { nil }
+
+        it 'does not write the stats' do
+          expect(cache).not_to receive(:write)
+
+          stats_cache.write_if_empty(stats)
+        end
+      end
+    end
+  end
+
+  describe '#clear' do
+    it 'clears cache' do
+      expect(cache).to receive(:delete).with(key)
+
+      stats_cache.clear
+    end
+  end
+end

spec/lib/gitlab/git/diff_stats_collection_spec.rb

@@ -4,11 +4,11 @@ require "spec_helper"
 
 describe Gitlab::Git::DiffStatsCollection do
   let(:stats_a) do
-    double(Gitaly::DiffStats, additions: 10, deletions: 15, path: 'foo')
+    Gitaly::DiffStats.new(additions: 10, deletions: 15, path: 'foo')
   end
 
   let(:stats_b) do
-    double(Gitaly::DiffStats, additions: 5, deletions: 1, path: 'bar')
+    Gitaly::DiffStats.new(additions: 5, deletions: 1, path: 'bar')
   end
 
   let(:diff_stats) { [stats_a, stats_b] }

spec/services/merge_requests/reload_diffs_service_spec.rb

@@ -34,10 +34,8 @@ describe MergeRequests::ReloadDiffsService, :use_clean_rails_memory_store_cachin
 
     context 'cache clearing' do
       it 'clears the cache for older diffs on the merge request' do
-        old_diff = merge_request.merge_request_diff
-        old_cache_key = old_diff.diffs_collection.cache_key
-
-        expect_any_instance_of(Redis).to receive(:del).with(old_cache_key).and_call_original
+        expect_any_instance_of(Redis).to receive(:del).once.and_call_original
+        expect(Rails.cache).to receive(:delete).once.and_call_original
 
         subject.execute
       end

spec/support/shared_examples/lib/gitlab/diff_file_collections_shared_examples.rb

@@ -10,7 +10,7 @@ RSpec.shared_examples 'diff statistics' do |test_include_stats_flag: true|
     end
   end
 
-  context 'when should request diff stats' do
+  context 'when include_stats is true' do
     it 'Repository#diff_stats is called' do
       expect(diffable.project.repository)
         .to receive(:diff_stats)
@@ -59,43 +59,134 @@ RSpec.shared_examples 'unfoldable diff' do
 end
 
 RSpec.shared_examples 'cacheable diff collection' do
-  let(:cache) { instance_double(Gitlab::Diff::HighlightCache) }
+  let(:highlight_cache) { instance_double(Gitlab::Diff::HighlightCache, write_if_empty: true, clear: nil, decorate: nil) }
+  let(:stats_cache) { instance_double(Gitlab::Diff::StatsCache, read: nil, write_if_empty: true, clear: nil) }
 
   before do
-    expect(Gitlab::Diff::HighlightCache).to receive(:new).with(subject) { cache }
+    expect(Gitlab::Diff::HighlightCache).to receive(:new).with(subject) { highlight_cache }
   end
 
   describe '#write_cache' do
     it 'calls Gitlab::Diff::HighlightCache#write_if_empty' do
-      expect(cache).to receive(:write_if_empty).once
+      expect(highlight_cache).to receive(:write_if_empty).once
+
+      subject.write_cache
+    end
+
+    context 'when the feature flag is enabled' do
+      before do
+        stub_feature_flags(cache_diff_stats_merge_request: true)
+        expect(Gitlab::Diff::StatsCache).to receive(:new).with(cachable_key: diffable.cache_key) { stats_cache }
+      end
+
+      it 'calls Gitlab::Diff::StatsCache#write_if_empty with diff stats' do
+        diff_stats = Gitlab::Git::DiffStatsCollection.new([])
+
+        expect(diffable.project.repository)
+          .to receive(:diff_stats).and_return(diff_stats)
+
+        expect(stats_cache).to receive(:write_if_empty).once.with(diff_stats)
+
+        subject.write_cache
+      end
+    end
+
+    context 'when the feature flag is disabled' do
+      before do
+        stub_feature_flags(cache_diff_stats_merge_request: false)
+      end
+
+      it 'does not call Gitlab::Diff::StatsCache#write_if_empty' do
+        expect(stats_cache).not_to receive(:write_if_empty)
 
         subject.write_cache
       end
     end
+  end
 
   describe '#clear_cache' do
     it 'calls Gitlab::Diff::HighlightCache#clear' do
-      expect(cache).to receive(:clear).once
+      expect(highlight_cache).to receive(:clear).once
+
+      subject.clear_cache
+    end
+
+    context 'when the feature flag is enabled' do
+      before do
+        stub_feature_flags(cache_diff_stats_merge_request: true)
+        expect(Gitlab::Diff::StatsCache).to receive(:new).with(cachable_key: diffable.cache_key) { stats_cache }
+      end
+
+      it 'calls Gitlab::Diff::StatsCache#clear' do
+        expect(stats_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
+    context 'when the feature flag is disabled' do
+      before do
+        stub_feature_flags(cache_diff_stats_merge_request: false)
+      end
+
+      it 'does not calls Gitlab::Diff::StatsCache#clear' do
+        expect(stats_cache).not_to receive(:clear)
 
-      subject.cache_key
+        subject.clear_cache
+      end
     end
   end
 
   describe '#diff_files' do
     it 'calls Gitlab::Diff::HighlightCache#decorate' do
-      expect(cache).to receive(:decorate)
+      expect(highlight_cache).to receive(:decorate)
         .with(instance_of(Gitlab::Diff::File))
         .exactly(cacheable_files_count).times
 
       subject.diff_files
     end
+
+    context 'when the feature swtich is enabled' do
+      before do
+        stub_feature_flags(cache_diff_stats_merge_request: true)
+        expect(Gitlab::Diff::StatsCache).to receive(:new).with(cachable_key: diffable.cache_key) { stats_cache }
+      end
+
+      context 'when there are stats cached' do
+        before do
+          allow(stats_cache).to receive(:read).and_return(Gitlab::Git::DiffStatsCollection.new([]))
+        end
+
+        it 'does not make a diff stats rpc call' do
+          expect(diffable.project.repository).not_to receive(:diff_stats)
+
+          subject.diff_files
+        end
+      end
+
+      context 'when there are no stats cached' do
+        it 'makes a diff stats rpc call' do
+          expect(diffable.project.repository)
+            .to receive(:diff_stats)
+            .with(diffable.diff_refs.base_sha, diffable.diff_refs.head_sha)
+
+          subject.diff_files
+        end
+      end
+    end
+
+    context 'when the feature switch is disabled' do
+      before do
+        stub_feature_flags(cache_diff_stats_merge_request: false)
+      end
+
+      it 'makes a diff stats rpc call' do
+        expect(diffable.project.repository)
+          .to receive(:diff_stats)
+          .with(diffable.diff_refs.base_sha, diffable.diff_refs.head_sha)
+
+        subject.diff_files
+      end
+    end
   end
 end