Paginate diffs using Gitaly paths

Changelog: performance

app/models/merge_request_diff.rb

@@ -389,11 +389,21 @@ class MergeRequestDiff < ApplicationRecord
 
   def diffs_in_batch(batch_page, batch_size, diff_options:)
     fetching_repository_diffs(diff_options) do |comparison|
+      reorder_diff_files!
+      diffs_batch = diffs_in_batch_collection(batch_page, batch_size, diff_options: diff_options)
+
       if comparison
-        comparison.diffs_in_batch(batch_page, batch_size, diff_options: diff_options)
+
+        if diff_options[:paths].blank? && !without_files? && diffs_batch.diff_file_paths.present?
+          diff_options.merge!(
+            paths: diffs_batch.diff_file_paths,
+            pagination_data: diffs_batch.pagination_data
+          )
+        end
+
+        comparison.diffs(diff_options)
       else
-        reorder_diff_files!
-        diffs_in_batch_collection(batch_page, batch_size, diff_options: diff_options)
+        diffs_batch
       end
     end
   end

lib/gitlab/diff/file_collection/base.rb

@@ -19,6 +19,7 @@ module Gitlab
 
           @diffable = diffable
           @include_stats = diff_options.delete(:include_stats)
+          @pagination_data = diff_options.delete(:pagination_data)
           @project = project
           @diff_options = diff_options
           @diff_refs = diff_refs
@@ -47,11 +48,7 @@ module Gitlab
         end
 
         def pagination_data
-          {
-            current_page: nil,
-            next_page: nil,
-            total_pages: nil
-          }
+          @pagination_data || empty_pagination_data
         end
 
         # This mutates `diff_files` lines.
@@ -90,6 +87,14 @@ module Gitlab
 
         private
 
+        def empty_pagination_data
+          {
+            current_page: nil,
+            next_page: nil,
+            total_pages: nil
+          }
+        end
+
         def diff_stats_collection
           strong_memoize(:diff_stats) do
             next unless fetch_diff_stats?

spec/models/merge_request_diff_spec.rb

@@ -418,8 +418,8 @@ RSpec.describe MergeRequestDiff do
       shared_examples_for 'fetching full diffs' do
         it 'returns diffs from repository comparison' do
           expect_next_instance_of(Compare) do |comparison|
-            expect(comparison).to receive(:diffs_in_batch)
-              .with(1, 10, diff_options: diff_options)
+            expect(comparison).to receive(:diffs)
+              .with(diff_options)
               .and_call_original
           end
 
@@ -448,13 +448,13 @@ RSpec.describe MergeRequestDiff do
         end
 
         it_behaves_like 'fetching full diffs'
-      end
 
-      context 'when diff_options include ignore_whitespace_change' do
-        it_behaves_like 'fetching full diffs' do
+        context 'when diff_options include ignore_whitespace_change' do
           let(:diff_options) do
             { ignore_whitespace_change: true }
           end
+
+          it_behaves_like 'fetching full diffs'
         end
       end
 
@@ -485,6 +485,35 @@ RSpec.describe MergeRequestDiff do
             'files/whitespace'
           ])
         end
+
+        context 'when diff_options include ignore_whitespace_change' do
+          let(:diff_options) do
+            { ignore_whitespace_change: true }
+          end
+
+          it 'returns a Gitlab::Diff::FileCollection::Compare with paginated diffs' do
+            diffs = diff_with_commits.diffs_in_batch(1, 10, diff_options: diff_options)
+
+            expect(diffs).to be_a(Gitlab::Diff::FileCollection::Compare)
+            expect(diffs.diff_files.size).to eq 10
+            expect(diffs.pagination_data).to eq(current_page: 1, next_page: 2, total_pages: 2)
+          end
+
+          context 'with gradual load enabled' do
+            before do
+              stub_feature_flags(diffs_gradual_load: true)
+            end
+
+            it 'returns pagination data from MergeRequestDiffBatch' do
+              diffs = diff_with_commits.diffs_in_batch(1, 10, diff_options: diff_options)
+              file_count = diff_with_commits.merge_request_diff_files.count
+
+              expect(diffs).to be_a(Gitlab::Diff::FileCollection::Compare)
+              expect(diffs.diff_files.size).to eq 10
+              expect(diffs.pagination_data).to eq(current_page: nil, next_page: nil, total_pages: file_count)
+            end
+          end
+        end
       end
     end