Create merge head diff on mergeability check

We create the merge head ref when we check for mergeability. Since
we are getting the data for merge head diff off that, we do it at
the same time.

app/controllers/projects/merge_requests/diffs_controller.rb

@@ -122,10 +122,7 @@ class Projects::MergeRequests::DiffsController < Projects::MergeRequests::Applic
       end
     end
 
-    if render_merge_ref_head_diff?
-      return CompareService.new(@project, @merge_request.merge_ref_head.sha)
-        .execute(@project, @merge_request.target_branch)
-    end
+    return @merge_request.merge_head_diff if render_merge_ref_head_diff?
 
     if @start_sha
       @merge_request_diff.compare_with(@start_sha)

app/models/merge_request.rb

@@ -479,13 +479,17 @@ class MergeRequest < ApplicationRecord
   # This is used after project import, to reset the IDs to the correct
   # values. It is not intended to be called without having already scoped the
   # relation.
+  #
+  # Only set `regular` merge request diffs as latest so `merge_head` diff
+  # won't be considered as `MergeRequest#merge_request_diff`.
   def self.set_latest_merge_request_diff_ids!
-    update = '
+    update = "
       latest_merge_request_diff_id = (
         SELECT MAX(id)
         FROM merge_request_diffs
         WHERE merge_requests.id = merge_request_diffs.merge_request_id
-      )'.squish
+        AND merge_request_diffs.diff_type = #{MergeRequestDiff.diff_types[:regular]}
+      )".squish
 
     self.each_batch do |batch|
       batch.update_all(update)
@@ -924,7 +928,7 @@ class MergeRequest < ApplicationRecord
   def create_merge_request_diff
     fetch_ref!
 
-    # n+1: https://gitlab.com/gitlab-org/gitlab-foss/issues/37435
+    # n+1: https://gitlab.com/gitlab-org/gitlab/-/issues/19377
     Gitlab::GitalyClient.allow_n_plus_1_calls do
       merge_request_diffs.create!
       reload_merge_request_diff
@@ -998,7 +1002,7 @@ class MergeRequest < ApplicationRecord
   # rubocop: enable CodeReuse/ServiceClass
 
   def diffable_merge_ref?
-    open? && merge_ref_head.present? && (Feature.enabled?(:display_merge_conflicts_in_diff, project) || can_be_merged?)
+    open? && merge_head_diff.present? && (Feature.enabled?(:display_merge_conflicts_in_diff, project) || can_be_merged?)
   end
 
   # Returns boolean indicating the merge_status should be rechecked in order to

app/models/merge_request_diff.rb

@@ -78,6 +78,7 @@ class MergeRequestDiff < ApplicationRecord
 
     join_condition = merge_requests[:id].eq(mr_diffs[:merge_request_id])
       .and(mr_diffs[:id].not_eq(merge_requests[:latest_merge_request_diff_id]))
+      .and(mr_diffs[:diff_type].eq(diff_types[:regular]))
 
     arel_join = mr_diffs.join(merge_requests).on(join_condition)
     joins(arel_join.join_sources)
@@ -202,6 +203,10 @@ class MergeRequestDiff < ApplicationRecord
   end
 
   def set_as_latest_diff
+    # Don't set merge_head diff as latest so it won't get considered as the
+    # MergeRequest#merge_request_diff.
+    return if merge_head?
+
     MergeRequest
       .where('id = ? AND COALESCE(latest_merge_request_diff_id, 0) < ?', self.merge_request_id, self.id)
       .update_all(latest_merge_request_diff_id: self.id)
@@ -209,8 +214,16 @@ class MergeRequestDiff < ApplicationRecord
 
   def ensure_commit_shas
     self.start_commit_sha ||= merge_request.target_branch_sha
-    self.head_commit_sha  ||= merge_request.source_branch_sha
-    self.base_commit_sha  ||= find_base_sha
+
+    if merge_head? && merge_request.merge_ref_head.present?
+      diff_refs = merge_request.merge_ref_head.diff_refs
+
+      self.head_commit_sha  ||= diff_refs.head_sha
+      self.base_commit_sha  ||= diff_refs.base_sha
+    else
+      self.head_commit_sha  ||= merge_request.source_branch_sha
+      self.base_commit_sha  ||= find_base_sha
+    end
   end
 
   # Override head_commit_sha to keep compatibility with merge request diff

app/services/merge_requests/mergeability_check_service.rb

@@ -114,6 +114,7 @@ module MergeRequests
 
       merge_to_ref_success = merge_to_ref
 
+      reload_merge_head_diff
       update_diff_discussion_positions! if merge_to_ref_success
 
       if merge_to_ref_success && can_git_merge?
@@ -123,6 +124,10 @@ module MergeRequests
       end
     end
 
+    def reload_merge_head_diff
+      MergeRequests::ReloadMergeHeadDiffService.new(merge_request).execute
+    end
+
     def update_diff_discussion_positions!
       Discussions::CaptureDiffNotePositionsService.new(merge_request).execute
     end
@@ -153,6 +158,7 @@ module MergeRequests
     def merge_to_ref
       params = { allow_conflicts: Feature.enabled?(:display_merge_conflicts_in_diff, project) }
       result = MergeRequests::MergeToRefService.new(project, merge_request.author, params).execute(merge_request)
+
       result[:status] == :success
     end
 

app/services/merge_requests/reload_merge_head_diff_service.rb

+# frozen_string_literal: true
+
+module MergeRequests
+  class ReloadMergeHeadDiffService
+    include BaseServiceUtility
+
+    def initialize(merge_request)
+      @merge_request = merge_request
+    end
+
+    def execute
+      return error("default_merge_ref_for_diffs feature flag is disabled") unless enabled?
+      return error("Merge request has no merge ref head.") unless merge_request.merge_ref_head.present?
+
+      error_msg = recreate_merge_head_diff
+
+      return error(error_msg) if error_msg
+
+      success
+    end
+
+    private
+
+    attr_reader :merge_request
+
+    def enabled?
+      Feature.enabled?(:default_merge_ref_for_diffs, merge_request.project)
+    end
+
+    def recreate_merge_head_diff
+      merge_request.merge_head_diff&.destroy!
+
+      # n+1: https://gitlab.com/gitlab-org/gitlab/-/issues/19377
+      Gitlab::GitalyClient.allow_n_plus_1_calls do
+        merge_request.create_merge_head_diff!
+      end
+
+      # Reset the merge request so it won't load the merge head diff as the
+      # MergeRequest#merge_request_diff.
+      merge_request.reset
+
+      nil
+    rescue StandardError => e
+      message = "Failed to recreate merge head diff: #{e.message}"
+
+      Gitlab::AppLogger.error(message: message, merge_request_id: merge_request.id)
+      message
+    end
+  end
+end

spec/controllers/projects/merge_requests/diffs_controller_spec.rb

@@ -226,11 +226,7 @@ RSpec.describe Projects::MergeRequests::DiffsController do
         let(:diffable_merge_ref) { true }
 
         it 'compares diffs with the head' do
-          MergeRequests::MergeToRefService.new(project, merge_request.author).execute(merge_request)
-
-          expect(CompareService).to receive(:new).with(
-            project, merge_request.merge_ref_head.sha
-          ).and_call_original
+          create(:merge_request_diff, :merge_head, merge_request: merge_request)
 
           go(diff_head: true)
 
@@ -242,8 +238,6 @@ RSpec.describe Projects::MergeRequests::DiffsController do
         let(:diffable_merge_ref) { false }
 
         it 'compares diffs with the base' do
-          expect(CompareService).not_to receive(:new)
-
           go(diff_head: true)
 
           expect(response).to have_gitlab_http_status(:ok)

spec/models/merge_request_diff_spec.rb

@@ -55,6 +55,26 @@ RSpec.describe MergeRequestDiff do
     it { expect(subject.head_commit_sha).to eq('b83d6e391c22777fca1ed3012fce84f633d7fed0') }
     it { expect(subject.base_commit_sha).to eq('ae73cb07c9eeaf35924a10f713b364d32b2dd34f') }
     it { expect(subject.start_commit_sha).to eq('0b4bc9a49b562e85de7cc9e834518ea6828729b9') }
+
+    context 'when diff_type is merge_head' do
+      let_it_be(:merge_request) { create(:merge_request) }
+
+      let_it_be(:merge_head) do
+        MergeRequests::MergeToRefService
+          .new(merge_request.project, merge_request.author)
+          .execute(merge_request)
+
+        merge_request.create_merge_head_diff
+      end
+
+      it { expect(merge_head).to be_valid }
+      it { expect(merge_head).to be_persisted }
+      it { expect(merge_head.commits.count).to eq(30) }
+      it { expect(merge_head.diffs.count).to eq(20) }
+      it { expect(merge_head.head_commit_sha).to eq(merge_request.merge_ref_head.diff_refs.head_sha) }
+      it { expect(merge_head.base_commit_sha).to eq(merge_request.merge_ref_head.diff_refs.base_sha) }
+      it { expect(merge_head.start_commit_sha).to eq(merge_request.target_branch_sha) }
+    end
   end
 
   describe '.by_commit_sha' do
@@ -83,6 +103,7 @@ RSpec.describe MergeRequestDiff do
     let_it_be(:merge_request) { create(:merge_request) }
     let_it_be(:outdated) { merge_request.merge_request_diff }
     let_it_be(:latest) { merge_request.create_merge_request_diff }
+    let_it_be(:merge_head) { merge_request.create_merge_head_diff }
 
     let_it_be(:closed_mr) { create(:merge_request, :closed_last_month) }
     let(:closed) { closed_mr.merge_request_diff }
@@ -123,14 +144,14 @@ RSpec.describe MergeRequestDiff do
         stub_external_diffs_setting(enabled: true)
       end
 
-      it { is_expected.to contain_exactly(outdated.id, latest.id, closed.id, merged.id, closed_recently.id, merged_recently.id) }
+      it { is_expected.to contain_exactly(outdated.id, latest.id, closed.id, merged.id, closed_recently.id, merged_recently.id, merge_head.id) }
 
       it 'ignores diffs with 0 files' do
         MergeRequestDiffFile.where(merge_request_diff_id: [closed_recently.id, merged_recently.id]).delete_all
         closed_recently.update!(files_count: 0)
         merged_recently.update!(files_count: 0)
 
-        is_expected.to contain_exactly(outdated.id, latest.id, closed.id, merged.id)
+        is_expected.to contain_exactly(outdated.id, latest.id, closed.id, merged.id, merge_head.id)
       end
     end
 

spec/models/merge_request_spec.rb

@@ -491,6 +491,7 @@ RSpec.describe MergeRequest, factory_default: :keep do
 
       create(:merge_request, params).tap do |mr|
         diffs.times { mr.merge_request_diffs.create }
+        mr.create_merge_head_diff
       end
     end
 
@@ -4379,37 +4380,41 @@ RSpec.describe MergeRequest, factory_default: :keep do
   end
 
   describe '#diffable_merge_ref?' do
+    let(:merge_request) { create(:merge_request) }
+
     context 'merge request can be merged' do
-      context 'merge_to_ref is not calculated' do
+      context 'merge_head diff is not created' do
         it 'returns true' do
-          expect(subject.diffable_merge_ref?).to eq(false)
+          expect(merge_request.diffable_merge_ref?).to eq(false)
         end
       end
 
-      context 'merge_to_ref is calculated' do
+      context 'merge_head diff is created' do
         before do
-          MergeRequests::MergeToRefService.new(subject.project, subject.author).execute(subject)
+          create(:merge_request_diff, :merge_head, merge_request: merge_request)
         end
 
         it 'returns true' do
-          expect(subject.diffable_merge_ref?).to eq(true)
+          expect(merge_request.diffable_merge_ref?).to eq(true)
         end
 
         context 'merge request is merged' do
-          subject { build_stubbed(:merge_request, :merged, project: project) }
+          before do
+            merge_request.mark_as_merged!
+          end
 
           it 'returns false' do
-            expect(subject.diffable_merge_ref?).to eq(false)
+            expect(merge_request.diffable_merge_ref?).to eq(false)
           end
         end
 
         context 'merge request cannot be merged' do
           before do
-            subject.mark_as_unchecked!
+            merge_request.mark_as_unchecked!
           end
 
           it 'returns false' do
-            expect(subject.diffable_merge_ref?).to eq(true)
+            expect(merge_request.diffable_merge_ref?).to eq(true)
           end
 
           context 'display_merge_conflicts_in_diff is disabled' do
@@ -4418,7 +4423,7 @@ RSpec.describe MergeRequest, factory_default: :keep do
             end
 
             it 'returns false' do
-              expect(subject.diffable_merge_ref?).to eq(false)
+              expect(merge_request.diffable_merge_ref?).to eq(false)
             end
           end
         end

spec/services/merge_requests/delete_non_latest_diffs_service_spec.rb

@@ -12,6 +12,8 @@ RSpec.describe MergeRequests::DeleteNonLatestDiffsService, :clean_gitlab_redis_s
       stub_const("#{described_class.name}::BATCH_SIZE", 2)
 
       3.times { merge_request.create_merge_request_diff }
+      merge_request.create_merge_head_diff
+      merge_request.reset
     end
 
     it 'schedules non-latest merge request diffs removal' do

spec/services/merge_requests/mergeability_check_service_spec.rb

@@ -33,6 +33,14 @@ RSpec.describe MergeRequests::MergeabilityCheckService, :clean_gitlab_redis_shar
       expect(merge_request.merge_status).to eq('can_be_merged')
     end
 
+    it 'reloads merge head diff' do
+      expect_next_instance_of(MergeRequests::ReloadMergeHeadDiffService) do |service|
+        expect(service).to receive(:execute)
+      end
+
+      subject
+    end
+
     it 'update diff discussion positions' do
       expect_next_instance_of(Discussions::CaptureDiffNotePositionsService) do |service|
         expect(service).to receive(:execute)
@@ -142,7 +150,11 @@ RSpec.describe MergeRequests::MergeabilityCheckService, :clean_gitlab_redis_shar
       end
 
       it 'resets one merge request upon execution' do
-        expect_any_instance_of(MergeRequest).to receive(:reset).once
+        expect_next_instance_of(MergeRequests::ReloadMergeHeadDiffService) do |svc|
+          expect(svc).to receive(:execute).and_return(status: :success)
+        end
+
+        expect_any_instance_of(MergeRequest).to receive(:reset).once.and_call_original
 
         execute_within_threads(amount: 2)
       end
@@ -266,6 +278,14 @@ RSpec.describe MergeRequests::MergeabilityCheckService, :clean_gitlab_redis_shar
 
       it_behaves_like 'unmergeable merge request'
 
+      it 'reloads merge head diff' do
+        expect_next_instance_of(MergeRequests::ReloadMergeHeadDiffService) do |service|
+          expect(service).to receive(:execute)
+        end
+
+        subject
+      end
+
       it 'returns ServiceResponse.error' do
         result = subject
 
@@ -329,6 +349,12 @@ RSpec.describe MergeRequests::MergeabilityCheckService, :clean_gitlab_redis_shar
 
           subject
         end
+
+        it 'does not reload merge head diff' do
+          expect(MergeRequests::ReloadMergeHeadDiffService).not_to receive(:new)
+
+          subject
+        end
       end
     end
 

spec/services/merge_requests/reload_merge_head_diff_service_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe MergeRequests::ReloadMergeHeadDiffService do
+  let(:merge_request) { create(:merge_request) }
+
+  subject { described_class.new(merge_request).execute }
+
+  describe '#execute' do
+    before do
+      MergeRequests::MergeToRefService
+        .new(merge_request.project, merge_request.author)
+        .execute(merge_request)
+    end
+
+    it 'creates a merge head diff' do
+      expect(subject[:status]).to eq(:success)
+      expect(merge_request.reload.merge_head_diff).to be_present
+    end
+
+    context 'when merge ref head is not present' do
+      before do
+        allow(merge_request).to receive(:merge_ref_head).and_return(nil)
+      end
+
+      it 'returns error' do
+        expect(subject[:status]).to eq(:error)
+      end
+    end
+
+    context 'when failed to create merge head diff' do
+      before do
+        allow(merge_request).to receive(:create_merge_head_diff!).and_raise('fail')
+      end
+
+      it 'returns error' do
+        expect(subject[:status]).to eq(:error)
+      end
+    end
+
+    context 'when there is existing merge head diff' do
+      let!(:existing_merge_head_diff) { create(:merge_request_diff, :merge_head, merge_request: merge_request) }
+
+      it 'recreates merge head diff' do
+        expect(subject[:status]).to eq(:success)
+        expect(merge_request.reload.merge_head_diff).not_to eq(existing_merge_head_diff)
+      end
+    end
+
+    context 'when default_merge_ref_for_diffs feature flag is disabled' do
+      before do
+        stub_feature_flags(default_merge_ref_for_diffs: false)
+      end
+
+      it 'returns error' do
+        expect(subject[:status]).to eq(:error)
+      end
+    end
+  end
+end