Introduce paginated diffs endpoint

Prepares the diffs_batch endpoint to return the
latest MR diff files version and the pagination data
(page, total pages, etc).

It doesn't read from cache as it's not worth today, given
the read/write can't be done in chunks.
This will be worked further at:
https://gitlab.com/gitlab-org/gitlab/issues/30550

app/controllers/projects/merge_requests/application_controller.rb

@@ -14,7 +14,11 @@ class Projects::MergeRequests::ApplicationController < Projects::ApplicationCont
   end
 
   def merge_request_includes(association)
-    association.includes(:metrics, :assignees, author: :status) # rubocop:disable CodeReuse/ActiveRecord
+    association.includes(preloadable_mr_relations) # rubocop:disable CodeReuse/ActiveRecord
+  end
+
+  def preloadable_mr_relations
+    [:metrics, :assignees, { author: :status }]
   end
 
   def merge_request_params

app/controllers/projects/merge_requests/diffs_controller.rb

@@ -5,9 +5,9 @@ class Projects::MergeRequests::DiffsController < Projects::MergeRequests::Applic
   include RendersNotes
 
   before_action :apply_diff_view_cookie!
-  before_action :commit
-  before_action :define_diff_vars
-  before_action :define_diff_comment_vars
+  before_action :commit, except: :diffs_batch
+  before_action :define_diff_vars, except: :diffs_batch
+  before_action :define_diff_comment_vars, except: :diffs_batch
 
   def show
     render_diffs
@@ -17,8 +17,29 @@ class Projects::MergeRequests::DiffsController < Projects::MergeRequests::Applic
     render_diffs
   end
 
+  def diffs_batch
+    return render_404 unless Feature.enabled?(:diffs_batch_load, @merge_request.project)
+
+    diffable = @merge_request.merge_request_diff
+
+    return render_404 unless diffable
+
+    diffs = diffable.diffs_in_batch(params[:page], params[:per_page], diff_options: diff_options)
+
+    options = {
+      merge_request: @merge_request,
+      pagination_data: diffs.pagination_data
+    }
+
+    render json: PaginatedDiffSerializer.new(current_user: current_user).represent(diffs, options)
+  end
+
   private
 
+  def preloadable_mr_relations
+    [{ source_project: :namespace }, { target_project: :namespace }]
+  end
+
   def render_diffs
     @environment = @merge_request.environments_for(current_user).last
 

app/controllers/projects/merge_requests_controller.rb

@@ -17,6 +17,9 @@ class Projects::MergeRequestsController < Projects::MergeRequests::ApplicationCo
   before_action :set_issuables_index, only: [:index]
   before_action :authenticate_user!, only: [:assign_related_issues]
   before_action :check_user_can_push_to_source_branch!, only: [:rebase]
+  before_action only: [:show] do
+    push_frontend_feature_flag(:diffs_batch_load, @project)
+  end
 
   around_action :allow_gitaly_ref_name_caching, only: [:index, :show, :discussions]
 

app/models/merge_request_diff.rb

@@ -297,6 +297,13 @@ class MergeRequestDiff < ApplicationRecord
     base_commit_sha? && head_commit_sha? && start_commit_sha?
   end
 
+  def diffs_in_batch(batch_page, batch_size, diff_options:)
+    Gitlab::Diff::FileCollection::MergeRequestDiffBatch.new(self,
+                                                            batch_page,
+                                                            batch_size,
+                                                            diff_options: diff_options)
+  end
+
   def diffs(diff_options = nil)
     if without_files? && comparison = diff_refs&.compare_in(project)
       # It should fetch the repository when diffs are cleaned by the system.

app/models/merge_request_diff_file.rb

@@ -5,6 +5,7 @@ class MergeRequestDiffFile < ApplicationRecord
   include DiffFile
 
   belongs_to :merge_request_diff, inverse_of: :merge_request_diff_files
+  alias_attribute :index, :relative_order
 
   def utf8_diff
     return '' if diff.blank?

app/serializers/paginated_diff_entity.rb

+# frozen_string_literal: true
+
+# Serializes diffs with pagination data.
+#
+# Avoid adding more keys to this serializer as processing the
+# diff file serialization is not cheap.
+#
+class PaginatedDiffEntity < Grape::Entity
+  include RequestAwareEntity
+
+  expose :diff_files do |diffs, options|
+    submodule_links = Gitlab::SubmoduleLinks.new(merge_request.project.repository)
+    DiffFileEntity.represent(diffs.diff_files, options.merge(submodule_links: submodule_links))
+  end
+
+  expose :pagination do
+    expose :current_page
+    expose :next_page
+    expose :total_pages
+    expose :next_page_href do |diffs|
+      next unless next_page
+
+      project = merge_request.target_project
+
+      diffs_batch_namespace_project_json_merge_request_path(
+        namespace_id: project.namespace.to_param,
+        project_id: project.to_param,
+        id: merge_request.iid,
+        page: next_page,
+        format: :json
+      )
+    end
+  end
+
+  private
+
+  %i[current_page next_page total_pages].each do |method|
+    define_method method do
+      pagination_data[method]
+    end
+  end
+
+  def pagination_data
+    options.fetch(:pagination_data, {})
+  end
+
+  def merge_request
+    options[:merge_request]
+  end
+end

app/serializers/paginated_diff_serializer.rb

+# frozen_string_literal: true
+
+class PaginatedDiffSerializer < BaseSerializer
+  entity PaginatedDiffEntity
+end

changelogs/unreleased/31290-mr-diffs-batch-load.yml

+---
+title: Introduce diffs_batch JSON endpoint for paginated diffs
+merge_request: 17651
+author:
+type: added

config/routes/project.rb

@@ -281,6 +281,7 @@ constraints(::Constraints::ProjectUrlConstrainer.new) do
             get :commits
             get :pipelines
             get :diffs, to: 'merge_requests/diffs#show'
+            get :diffs_batch, to: 'merge_requests/diffs#diffs_batch'
             get :widget, to: 'merge_requests/content#widget'
             get :cached_widget, to: 'merge_requests/content#cached_widget'
           end

lib/gitlab/diff/file_collection/merge_request_diff.rb

@@ -3,19 +3,7 @@
 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
-
-          super(merge_request_diff,
-            project: merge_request_diff.project,
-            diff_options: diff_options,
-            diff_refs: merge_request_diff.diff_refs,
-            fallback_diff_refs: merge_request_diff.fallback_diff_refs)
-        end
-
+      class MergeRequestDiff < MergeRequestDiffBase
         def diff_files
           diff_files = super
 

lib/gitlab/diff/file_collection/merge_request_diff_base.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Diff
+    module FileCollection
+      class MergeRequestDiffBase < Base
+        extend ::Gitlab::Utils::Override
+
+        def initialize(merge_request_diff, diff_options:)
+          @merge_request_diff = merge_request_diff
+
+          super(merge_request_diff,
+            project: merge_request_diff.project,
+            diff_options: diff_options,
+            diff_refs: merge_request_diff.diff_refs,
+            fallback_diff_refs: merge_request_diff.fallback_diff_refs)
+        end
+      end
+    end
+  end
+end

lib/gitlab/diff/file_collection/merge_request_diff_batch.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Diff
+    module FileCollection
+      # Builds a paginated diff file collection and collects pagination
+      # metadata.
+      #
+      # It doesn't handle caching yet as we're not prepared to write/read
+      # separate file keys (https://gitlab.com/gitlab-org/gitlab/issues/30550).
+      #
+      class MergeRequestDiffBatch < MergeRequestDiffBase
+        DEFAULT_BATCH_PAGE = 1
+        DEFAULT_BATCH_SIZE = 20
+
+        attr_reader :pagination_data
+
+        def initialize(merge_request_diff, batch_page, batch_size, diff_options:)
+          super(merge_request_diff, diff_options: diff_options)
+
+          batch_page ||= DEFAULT_BATCH_PAGE
+          batch_size ||= DEFAULT_BATCH_SIZE
+
+          @paginated_collection = relation.page(batch_page).per(batch_size)
+          @pagination_data = {
+            current_page: @paginated_collection.current_page,
+            next_page: @paginated_collection.next_page,
+            total_pages: @paginated_collection.total_pages
+          }
+        end
+
+        override :diffs
+        def diffs
+          strong_memoize(:diffs) do
+            @merge_request_diff.opening_external_diff do
+              # Avoiding any extra queries.
+              collection = @paginated_collection.to_a
+
+              # The offset collection and calculation is required so that we
+              # know how much has been loaded in previous batches, collapsing
+              # the current paginated set accordingly (collection limit calculation).
+              # See: https://docs.gitlab.com/ee/development/diffs.html#diff-collection-limits
+              #
+              offset_index = collection.first&.index
+              options = diff_options.dup
+
+              collection =
+                if offset_index && offset_index > 0
+                  offset_collection = relation.limit(offset_index) # rubocop:disable CodeReuse/ActiveRecord
+                  options[:offset_index] = offset_index
+                  offset_collection + collection
+                else
+                  collection
+                end
+
+              Gitlab::Git::DiffCollection.new(collection.map(&:to_hash), options)
+            end
+          end
+        end
+
+        private
+
+        def relation
+          @merge_request_diff.merge_request_diff_files
+        end
+      end
+    end
+  end
+end

lib/gitlab/git/diff_collection.rb

@@ -31,6 +31,7 @@ module Gitlab
         @limits = self.class.limits(options)
         @enforce_limits = !!options.fetch(:limits, true)
         @expanded = !!options.fetch(:expanded, true)
+        @offset_index = options.fetch(:offset_index, 0)
 
         @line_count = 0
         @byte_count = 0
@@ -128,7 +129,7 @@ module Gitlab
       def each_serialized_patch
         i = @array.length
 
-        @iterator.each do |raw|
+        @iterator.each_with_index do |raw, iterator_index|
           @empty = false
 
           if @enforce_limits && i >= max_files
@@ -154,8 +155,12 @@ module Gitlab
             break
           end
 
-          yield @array[i] = diff
-          i += 1
+          # We should not yield / memoize diffs before the offset index. Though,
+          # we still consider the limit buffers for diffs before it.
+          if iterator_index >= @offset_index
+            yield @array[i] = diff
+            i += 1
+          end
         end
       end
     end

spec/controllers/projects/merge_requests/diffs_controller_spec.rb

@@ -5,6 +5,49 @@ require 'spec_helper'
 describe Projects::MergeRequests::DiffsController do
   include ProjectForksHelper
 
+  shared_examples 'forked project with submodules' do
+    render_views
+
+    let(:project) { create(:project, :repository) }
+    let(:forked_project) { fork_project_with_submodules(project) }
+    let(:merge_request) { create(:merge_request_with_diffs, source_project: forked_project, source_branch: 'add-submodule-version-bump', target_branch: 'master', target_project: project) }
+
+    before do
+      project.add_developer(user)
+
+      merge_request.reload
+      go
+    end
+
+    it 'renders' do
+      expect(response).to be_successful
+      expect(response.body).to have_content('Subproject commit')
+    end
+  end
+
+  shared_examples 'persisted preferred diff view cookie' do
+    context 'with view param' do
+      before do
+        go(view: 'parallel')
+      end
+
+      it 'saves the preferred diff view in a cookie' do
+        expect(response.cookies['diff_view']).to eq('parallel')
+      end
+    end
+
+    context 'when the user cannot view the merge request' do
+      before do
+        project.team.truncate
+        go
+      end
+
+      it 'returns a 404' do
+        expect(response).to have_gitlab_http_status(404)
+      end
+    end
+  end
+
   let(:project) { create(:project, :repository) }
   let(:user) { create(:user) }
   let(:merge_request) { create(:merge_request_with_diffs, target_project: project, source_project: project) }
@@ -51,36 +94,10 @@ describe Projects::MergeRequests::DiffsController do
         end
       end
 
-      context 'with forked projects with submodules' do
-        render_views
-
-        let(:project) { create(:project, :repository) }
-        let(:forked_project) { fork_project_with_submodules(project) }
-        let(:merge_request) { create(:merge_request_with_diffs, source_project: forked_project, source_branch: 'add-submodule-version-bump', target_branch: 'master', target_project: project) }
-
-        before do
-          project.add_developer(user)
-
-          merge_request.reload
-          go
-        end
-
-        it 'renders' do
-          expect(response).to be_successful
-          expect(response.body).to have_content('Subproject commit')
-        end
-      end
+      it_behaves_like 'forked project with submodules'
     end
 
-    context 'with view' do
-      before do
-        go(view: 'parallel')
-      end
-
-      it 'saves the preferred diff view in a cookie' do
-        expect(response.cookies['diff_view']).to eq('parallel')
-      end
-    end
+    it_behaves_like 'persisted preferred diff view cookie'
   end
 
   describe 'GET diff_for_path' do
@@ -154,4 +171,92 @@ describe Projects::MergeRequests::DiffsController do
       end
     end
   end
+
+  describe 'GET diffs_batch' do
+    def go(extra_params = {})
+      params = {
+        namespace_id: project.namespace.to_param,
+        project_id: project,
+        id: merge_request.iid,
+        format: 'json'
+      }
+
+      get :diffs_batch, params: params.merge(extra_params)
+    end
+
+    context 'when feature is disabled' do
+      before do
+        stub_feature_flags(diffs_batch_load: false)
+      end
+
+      it 'returns 404' do
+        go
+
+        expect(response).to have_gitlab_http_status(404)
+      end
+    end
+
+    context 'when not authorized' do
+      let(:other_user) { create(:user) }
+
+      before do
+        sign_in(other_user)
+      end
+
+      it 'returns 404' do
+        go
+
+        expect(response).to have_gitlab_http_status(404)
+      end
+    end
+
+    context 'with default params' do
+      let(:expected_options) do
+        {
+          merge_request: merge_request,
+          pagination_data: {
+            current_page: 1,
+            next_page: nil,
+            total_pages: 1
+          }
+        }
+      end
+
+      it 'serializes paginated merge request diff collection' do
+        expect_next_instance_of(PaginatedDiffSerializer) do |instance|
+          expect(instance).to receive(:represent)
+            .with(an_instance_of(Gitlab::Diff::FileCollection::MergeRequestDiffBatch), expected_options)
+            .and_call_original
+        end
+
+        go
+      end
+    end
+
+    context 'with smaller diff batch params' do
+      let(:expected_options) do
+        {
+          merge_request: merge_request,
+          pagination_data: {
+            current_page: 2,
+            next_page: 3,
+            total_pages: 4
+          }
+        }
+      end
+
+      it 'serializes paginated merge request diff collection' do
+        expect_next_instance_of(PaginatedDiffSerializer) do |instance|
+          expect(instance).to receive(:represent)
+            .with(an_instance_of(Gitlab::Diff::FileCollection::MergeRequestDiffBatch), expected_options)
+            .and_call_original
+        end
+
+        go(page: 2, per_page: 5)
+      end
+    end
+
+    it_behaves_like 'forked project with submodules'
+    it_behaves_like 'persisted preferred diff view cookie'
+  end
 end

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

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::Diff::FileCollection::MergeRequestDiffBatch do
+  let(:merge_request) { create(:merge_request) }
+  let(:batch_page) { 1 }
+  let(:batch_size) { 10 }
+  let(:diffable) { merge_request.merge_request_diff }
+  let(:diff_files_relation) { diffable.merge_request_diff_files }
+
+  subject do
+    described_class.new(diffable,
+                        batch_page,
+                        batch_size,
+                        diff_options: nil)
+  end
+
+  let(:diff_files) { subject.diff_files }
+
+  describe 'initialize' do
+    it 'memoizes pagination_data' do
+      expect(subject.pagination_data).to eq(current_page: 1, next_page: 2, total_pages: 2)
+    end
+  end
+
+  describe '#diff_files' do
+    let(:batch_size) { 3 }
+    let(:paginated_rel) { diff_files_relation.page(batch_page).per(batch_size) }
+
+    let(:expected_batch_files) do
+      paginated_rel.map(&:new_path)
+    end
+
+    it 'returns paginated diff files' do
+      expect(diff_files.size).to eq(3)
+    end
+
+    it 'returns a valid instance of a DiffCollection' do
+      expect(diff_files).to be_a(Gitlab::Git::DiffCollection)
+    end
+
+    context 'first page' do
+      it 'returns correct diff files' do
+        expect(diff_files.map(&:new_path)).to eq(expected_batch_files)
+      end
+    end
+
+    context 'another page' do
+      let(:batch_page) { 2 }
+
+      it 'returns correct diff files' do
+        expect(diff_files.map(&:new_path)).to eq(expected_batch_files)
+      end
+    end
+
+    context 'nil batch_page' do
+      let(:batch_page) { nil }
+
+      it 'returns correct diff files' do
+        expected_batch_files =
+          diff_files_relation.page(described_class::DEFAULT_BATCH_PAGE).per(batch_size).map(&:new_path)
+
+        expect(diff_files.map(&:new_path)).to eq(expected_batch_files)
+      end
+    end
+
+    context 'nil batch_size' do
+      let(:batch_size) { nil }
+
+      it 'returns correct diff files' do
+        expected_batch_files =
+          diff_files_relation.page(batch_page).per(described_class::DEFAULT_BATCH_SIZE).map(&:new_path)
+
+        expect(diff_files.map(&:new_path)).to eq(expected_batch_files)
+      end
+    end
+
+    context 'invalid page' do
+      let(:batch_page) { 999 }
+
+      it 'returns correct diff files' do
+        expect(diff_files.map(&:new_path)).to be_empty
+      end
+    end
+
+    context 'last page' do
+      it 'returns correct diff files' do
+        last_page = paginated_rel.total_pages
+        collection = described_class.new(diffable,
+                                         last_page,
+                                         batch_size,
+                                         diff_options: nil)
+
+        expected_batch_files = diff_files_relation.page(last_page).per(batch_size).map(&:new_path)
+
+        expect(collection.diff_files.map(&:new_path)).to eq(expected_batch_files)
+      end
+    end
+  end
+
+  it_behaves_like 'unfoldable diff' do
+    subject do
+      described_class.new(merge_request.merge_request_diff,
+                          batch_page,
+                          batch_size,
+                          diff_options: nil)
+    end
+  end
+
+  it_behaves_like 'diff statistics' do
+    let(:collection_default_args) do
+      { diff_options: {} }
+    end
+
+    let(:diffable) { merge_request.merge_request_diff }
+    let(:stub_path) { '.gitignore' }
+
+    subject do
+      described_class.new(merge_request.merge_request_diff,
+                          batch_page,
+                          batch_size,
+                          collection_default_args)
+    end
+  end
+end

spec/lib/gitlab/git/diff_collection_spec.rb

@@ -537,6 +537,70 @@ describe Gitlab::Git::DiffCollection, :seed_helper do
         end
       end
     end
+
+    context 'when offset_index is given' do
+      subject do
+        Gitlab::Git::DiffCollection.new(
+          iterator,
+          max_files: max_files,
+          max_lines: max_lines,
+          limits: limits,
+          offset_index: 2,
+          expanded: expanded
+        )
+      end
+
+      def diff(raw)
+        raw['diff']
+      end
+
+      let(:iterator) do
+        [
+          fake_diff(1, 1),
+          fake_diff(2, 2),
+          fake_diff(3, 3),
+          fake_diff(4, 4)
+        ]
+      end
+
+      it 'does not yield diffs before the offset' do
+        expect(subject.to_a.map(&:diff)).to eq(
+          [
+            diff(fake_diff(3, 3)),
+            diff(fake_diff(4, 4))
+          ]
+        )
+      end
+
+      context 'when go over safe limits on bytes' do
+        let(:iterator) do
+          [
+            fake_diff(1, 10), # 10
+            fake_diff(1, 10), # 20
+            fake_diff(1, 15), # 35
+            fake_diff(1, 20), # 55
+            fake_diff(1, 45), # 100 - limit hit
+            fake_diff(1, 45),
+            fake_diff(1, 20480),
+            fake_diff(1, 1)
+          ]
+        end
+
+        before do
+          stub_const('Gitlab::Git::DiffCollection::DEFAULT_LIMITS',
+                     { max_files: max_files, max_lines: 80 })
+        end
+
+        it 'considers size of diffs before the offset for prunning' do
+          expect(subject.to_a.map(&:diff)).to eq(
+            [
+              diff(fake_diff(1, 15)),
+              diff(fake_diff(1, 20))
+            ]
+          )
+        end
+      end
+    end
   end
 
   def fake_diff(line_length, line_count)

spec/serializers/paginated_diff_entity_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe PaginatedDiffEntity do
+  let(:user) { create(:user) }
+  let(:request) { double('request', current_user: user) }
+  let(:merge_request) { create(:merge_request, :with_diffs) }
+  let(:diff_batch) { merge_request.merge_request_diff.diffs_in_batch(2, 3, diff_options: nil) }
+  let(:options) do
+    {
+      request: request,
+      merge_request: merge_request,
+      pagination_data: diff_batch.pagination_data
+    }
+  end
+  let(:entity) { described_class.new(diff_batch, options) }
+
+  subject { entity.as_json }
+
+  it 'exposes diff_files' do
+    expect(subject[:diff_files]).to be_present
+  end
+
+  it 'exposes pagination data' do
+    expect(subject[:pagination]).to eq(
+      current_page: 2,
+      next_page: 3,
+      next_page_href: "/#{merge_request.project.full_path}/merge_requests/#{merge_request.iid}/diffs_batch.json?page=3",
+      total_pages: 7
+    )
+  end
+end

spec/support/shared_examples/diff_file_collections.rb

 # frozen_string_literal: true
 
 shared_examples 'diff statistics' do |test_include_stats_flag: true|
+  subject { described_class.new(diffable, collection_default_args) }
+
   def stub_stats_find_by_path(path, stats_mock)
     expect_next_instance_of(Gitlab::Git::DiffStatsCollection) do |collection|
       allow(collection).to receive(:find_by_path).and_call_original
@@ -10,8 +12,6 @@ shared_examples 'diff statistics' do |test_include_stats_flag: true|
 
   context 'when should request diff stats' do
     it 'Repository#diff_stats is called' do
-      subject = described_class.new(diffable, collection_default_args)
-
       expect(diffable.project.repository)
         .to receive(:diff_stats)
         .with(diffable.diff_refs.base_sha, diffable.diff_refs.head_sha)
@@ -21,8 +21,6 @@ shared_examples 'diff statistics' do |test_include_stats_flag: true|
     end
 
     it 'Gitlab::Diff::File is initialized with diff stats' do
-      subject = described_class.new(diffable, collection_default_args)
-
       stats_mock = double(Gitaly::DiffStats, path: '.gitignore', additions: 758, deletions: 120)
       stub_stats_find_by_path(stub_path, stats_mock)
 
@@ -37,8 +35,6 @@ shared_examples 'diff statistics' do |test_include_stats_flag: true|
     it 'Repository#diff_stats is not called' do
       collection_default_args[:diff_options][:include_stats] = false
 
-      subject = described_class.new(diffable, collection_default_args)
-
       expect(diffable.project.repository).not_to receive(:diff_stats)
 
       subject.diff_files