Use BatchLoader for merge_request_diff

- Implement a custom preloader for MergeRequestDiff

app/graphql/batch_loaders/merge_request_diff_summary_batch_loader.rb

+# frozen_string_literal: true
+
+module BatchLoaders
+  class MergeRequestDiffSummaryBatchLoader
+    NIL_STATS = { additions: 0, deletions: 0, file_count: 0 }.freeze
+
+    def self.load_for(merge_request)
+      BatchLoader::GraphQL.for(merge_request).batch(key: :diff_stats_summary) do |merge_requests, loader, args|
+        Preloaders::MergeRequestDiffPreloader.new(merge_requests).preload_all
+
+        merge_requests.each do |merge_request|
+          metrics = merge_request.metrics
+
+          summary = if metrics && metrics.added_lines && metrics.removed_lines
+                      { additions: metrics.added_lines, deletions: metrics.removed_lines, file_count: merge_request.merge_request_diff&.files_count || 0 }
+                    elsif merge_request.diff_stats.blank?
+                      NIL_STATS
+                    else
+                      merge_request.diff_stats.each_with_object(NIL_STATS.dup) do |status, summary|
+                        summary.merge!(additions: status.additions, deletions: status.deletions, file_count: 1) { |_, x, y| x + y }
+                      end
+                    end
+
+          loader.call(merge_request, summary)
+        end
+      end
+    end
+  end
+end

app/graphql/types/merge_request_type.rb

@@ -182,18 +182,7 @@ module Types
     end
 
     def diff_stats_summary
-      metrics = object.metrics
-
-      if metrics && metrics.added_lines && metrics.removed_lines
-        return { additions: metrics.added_lines, deletions: metrics.removed_lines, file_count: object.merge_request_diff&.files_count || 0 }
-      end
-
-      nil_stats = { additions: 0, deletions: 0, file_count: 0 }
-      return nil_stats unless object.diff_stats.present?
-
-      object.diff_stats.each_with_object(nil_stats) do |status, hash|
-        hash.merge!(additions: status.additions, deletions: status.deletions, file_count: 1) { |_, x, y| x + y }
-      end
+      BatchLoaders::MergeRequestDiffSummaryBatchLoader.load_for(object)
     end
 
     def commit_count

app/models/merge_request_diff.rb

@@ -106,6 +106,17 @@ class MergeRequestDiff < ApplicationRecord
     joins(merge_request: :metrics).where(condition)
   end
 
+  scope :latest_diff_for_merge_requests, -> (merge_requests) do
+    inner_select = MergeRequestDiff
+      .default_scoped
+      .distinct
+      .select("FIRST_VALUE(id) OVER (PARTITION BY merge_request_id ORDER BY created_at DESC) as id")
+      .where(merge_request: merge_requests)
+
+    joins("INNER JOIN (#{inner_select.to_sql}) latest_diffs ON latest_diffs.id = merge_request_diffs.id")
+      .includes(:merge_request_diff_commits)
+  end
+
   class << self
     def ids_for_external_storage_migration(limit:)
       return [] unless Gitlab.config.external_diffs.enabled
@@ -280,7 +291,13 @@ class MergeRequestDiff < ApplicationRecord
   end
 
   def commit_shas(limit: nil)
-    merge_request_diff_commits.limit(limit).pluck(:sha)
+    if association(:merge_request_diff_commits).loaded?
+      sorted_diff_commits = merge_request_diff_commits.sort_by { |diff_commit| [diff_commit.id, diff_commit.relative_order] }
+      sorted_diff_commits = sorted_diff_commits.take(limit) if limit
+      sorted_diff_commits.map(&:sha)
+    else
+      merge_request_diff_commits.limit(limit).pluck(:sha)
+    end
   end
 
   def includes_any_commits?(shas)

app/models/preloaders/merge_request_diff_preloader.rb

+# frozen_string_literal: true
+
+module Preloaders
+  # This class preloads the `merge_request_diff` association for the given merge request models.
+  #
+  # Usage:
+  #   merge_requests = MergeRequest.where(...)
+  #   Preloaders::MergeRequestDiffPreloader.new(merge_requests).preload_all
+  #   merge_requests.first.merge_request_diff # won't fire any query
+  class MergeRequestDiffPreloader
+    def initialize(merge_requests)
+      @merge_requests = merge_requests
+    end
+
+    def preload_all
+      merge_request_diffs = MergeRequestDiff.latest_diff_for_merge_requests(@merge_requests)
+      cache = merge_request_diffs.index_by { |diff| diff.merge_request_id }
+
+      @merge_requests.each do |merge_request|
+        merge_request_diff = cache[merge_request.id]
+
+        merge_request.association(:merge_request_diff).target = merge_request_diff
+        merge_request.association(:merge_request_diff).loaded!
+      end
+    end
+  end
+end

changelogs/unreleased/preload-merge-request-diffs-for-throughtput-table.yml

+---
+title: Optimize the loading of diffStats in merge request GraphQL API
+merge_request: 44752
+author:
+type: performance

spec/models/merge_request_diff_spec.rb

@@ -657,13 +657,32 @@ RSpec.describe MergeRequestDiff do
       expect(diff_with_commits.commit_shas).to all(match(/\h{40}/))
     end
 
-    context 'with limit attribute' do
+    shared_examples 'limited number of shas' do
       it 'returns limited number of shas' do
         expect(diff_with_commits.commit_shas(limit: 2).size).to eq(2)
         expect(diff_with_commits.commit_shas(limit: 100).size).to eq(29)
         expect(diff_with_commits.commit_shas.size).to eq(29)
       end
     end
+
+    context 'with limit attribute' do
+      it_behaves_like 'limited number of shas'
+    end
+
+    context 'with preloaded diff commits' do
+      before do
+        # preloads the merge_request_diff_commits association
+        diff_with_commits.merge_request_diff_commits.to_a
+      end
+
+      it_behaves_like 'limited number of shas'
+
+      it 'does not trigger any query' do
+        count = ActiveRecord::QueryRecorder.new { diff_with_commits.commit_shas(limit: 2) }.count
+
+        expect(count).to eq(0)
+      end
+    end
   end
 
   describe '#compare_with' do
@@ -876,4 +895,25 @@ RSpec.describe MergeRequestDiff do
       expect(subject.lines_count).to eq 189
     end
   end
+
+  describe '.latest_diff_for_merge_requests' do
+    let_it_be(:merge_request_1) { create(:merge_request_without_merge_request_diff) }
+    let_it_be(:merge_request_1_diff_1) { create(:merge_request_diff, merge_request: merge_request_1, created_at: 3.days.ago) }
+    let_it_be(:merge_request_1_diff_2) { create(:merge_request_diff, merge_request: merge_request_1, created_at: 1.day.ago) }
+
+    let_it_be(:merge_request_2) { create(:merge_request_without_merge_request_diff) }
+    let_it_be(:merge_request_2_diff_1) { create(:merge_request_diff, merge_request: merge_request_2, created_at: 3.days.ago) }
+
+    let_it_be(:merge_request_3) { create(:merge_request_without_merge_request_diff) }
+
+    subject { described_class.latest_diff_for_merge_requests([merge_request_1, merge_request_2]) }
+
+    it 'loads the latest merge_request_diff record for the given merge requests' do
+      expect(subject).to match_array([merge_request_1_diff_2, merge_request_2_diff_1])
+    end
+
+    it 'loads nothing if the merge request has no diff record' do
+      expect(described_class.latest_diff_for_merge_requests(merge_request_3)).to be_empty
+    end
+  end
 end

spec/models/preloaders/merge_request_diff_preloader_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Preloaders::MergeRequestDiffPreloader do
+  let_it_be(:merge_request_1) { create(:merge_request) }
+  let_it_be(:merge_request_2) { create(:merge_request) }
+  let_it_be(:merge_request_3) { create(:merge_request_without_merge_request_diff) }
+
+  let(:merge_requests) { [merge_request_1, merge_request_2, merge_request_3] }
+
+  def trigger(merge_requests)
+    Array(merge_requests).each(&:merge_request_diff)
+  end
+
+  def merge_requests_with_preloaded_diff
+    described_class.new(MergeRequest.where(id: merge_requests.map(&:id)).to_a).preload_all
+  end
+
+  it 'does not trigger N+1 queries' do
+    # warmup
+    trigger(merge_requests_with_preloaded_diff)
+
+    first_merge_request = merge_requests_with_preloaded_diff.first
+    clean_merge_requests = merge_requests_with_preloaded_diff
+
+    expect { trigger(clean_merge_requests) }.to issue_same_number_of_queries_as { trigger(first_merge_request) }
+  end
+end