Merge branch '31290-mr-diffs-batch-load' into 'master'

Introduce diffs_batch JSON endpoint for paginated diffs

Closes #31290

See merge request gitlab-org/gitlab!17651

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