Avoid idling in transaction when fetching source for merge requests

Previously when a merge request was created, the `after_create` Rails
callback would fire the `FetchSourceBranch` Gitaly RPC. However, this
call is made inside a database transaction, which may have an idle in
transaction timeout.

On GitLab.com, this timeout is set to 30 seconds, so any
`FetchSourceBranch` calls taking longer than 30 seconds causes the
database connection to be terminated quietly. When Rails attempts to
make another database write (e.g. writing the metrics), the query
fails with a 500 error after the RPC completed, usually because a
subsequent query ran into an error.

Instead of failing after the idle in transaction, this commit will
cause the query to be killed if the total Web request time expires (60
seconds or so). While this isn't ideal, this is better than the
alternative because we should never hold open a database transaction
for a long Gitaly call. Gitaly team has a few optimizations for
`FetchSourceBranch` in the works, but allowing this RPC to run for up
to 60 seconds during a Web request is better than the status quo.

Relates to https://gitlab.com/gitlab-org/gitlab/-/issues/336657

Changelog: fixed

app/models/merge_request.rb

@@ -1016,8 +1016,23 @@ class MergeRequest < ApplicationRecord
     merge_request_diff.persisted? || create_merge_request_diff
   end
 
-  def create_merge_request_diff
+  def eager_fetch_ref!
+    return unless valid?
+
+    # has_internal_id normally attempts to allocate the iid in the
+    # before_create hook, but we need the iid to be available before
+    # that to fetch the ref into the target project.
+    track_target_project_iid!
+    ensure_target_project_iid!
+
     fetch_ref!
+    # Prevent the after_create hook from fetching the source branch again
+    # Drop this field after rollout in https://gitlab.com/gitlab-org/gitlab/-/issues/353044.
+    @skip_fetch_ref = true
+  end
+
+  def create_merge_request_diff
+    fetch_ref! unless skip_fetch_ref
 
     # n+1: https://gitlab.com/gitlab-org/gitlab/-/issues/19377
     Gitlab::GitalyClient.allow_n_plus_1_calls do
@@ -1950,6 +1965,8 @@ class MergeRequest < ApplicationRecord
 
   private
 
+  attr_accessor :skip_fetch_ref
+
   def set_draft_status
     self.draft = draft?
   end

app/services/merge_requests/create_service.rb

@@ -31,6 +31,16 @@ module MergeRequests
 
     private
 
+    def before_create(merge_request)
+      # If the fetching of the source branch occurs in an ActiveRecord
+      # callback (e.g. after_create), a database transaction will be
+      # open while the Gitaly RPC waits. To avoid an idle in transaction
+      # timeout, we do this before we attempt to save the merge request.
+      if Feature.enabled?(:merge_request_eager_fetch_ref, @project, default_enabled: :yaml)
+        merge_request.eager_fetch_ref!
+      end
+    end
+
     def set_projects!
       # @project is used to determine whether the user can set the merge request's
       # assignee, milestone and labels. Whether they can depends on their

config/feature_flags/development/merge_request_eager_fetch_ref.yml

+---
+name: merge_request_eager_fetch_ref
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/80876
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/353044
+milestone: '14.9'
+type: development
+group: group::code review
+default_enabled: false

spec/models/merge_request_spec.rb

@@ -4249,6 +4249,29 @@ RSpec.describe MergeRequest, factory_default: :keep do
     end
   end
 
+  describe '#eager_fetch_ref!' do
+    let(:project) { create(:project, :repository) }
+
+    # We use build instead of create to test that an IID is allocated
+    subject { build(:merge_request, source_project: project) }
+
+    it 'fetches the ref correctly' do
+      expect(subject.iid).to be_nil
+
+      expect { subject.eager_fetch_ref! }.to change { subject.iid.to_i }.by(1)
+
+      expect(subject.target_project.repository.ref_exists?(subject.ref_path)).to be_truthy
+    end
+
+    it 'only fetches the ref once after saved' do
+      expect(subject.target_project.repository).to receive(:fetch_source_branch!).once.and_call_original
+
+      subject.save!
+
+      expect(subject.target_project.repository.ref_exists?(subject.ref_path)).to be_truthy
+    end
+  end
+
   describe 'removing a merge request' do
     it 'refreshes the number of open merge requests of the target project' do
       project = subject.target_project

spec/services/merge_requests/create_service_spec.rb

@@ -454,7 +454,7 @@ RSpec.describe MergeRequests::CreateService, :clean_gitlab_redis_shared_state do
       end
     end
 
-    context 'when source and target projects are different' do
+    shared_examples 'when source and target projects are different' do |eager_fetch_ref_enabled|
       let(:target_project) { fork_project(project, nil, repository: true) }
 
       let(:opts) do
@@ -497,9 +497,18 @@ RSpec.describe MergeRequests::CreateService, :clean_gitlab_redis_shared_state do
         end
 
         it 'creates the merge request', :sidekiq_might_not_need_inline do
+          expect_next_instance_of(MergeRequest) do |instance|
+            if eager_fetch_ref_enabled
+              expect(instance).to receive(:eager_fetch_ref!).and_call_original
+            else
+              expect(instance).not_to receive(:eager_fetch_ref!)
+            end
+          end
+
           merge_request = described_class.new(project: project, current_user: user, params: opts).execute
 
           expect(merge_request).to be_persisted
+          expect(merge_request.iid).to be > 0
         end
 
         it 'does not create the merge request when the target project is archived' do
@@ -511,6 +520,18 @@ RSpec.describe MergeRequests::CreateService, :clean_gitlab_redis_shared_state do
       end
     end
 
+    context 'when merge_request_eager_fetch_ref is enabled' do
+      it_behaves_like 'when source and target projects are different', true
+    end
+
+    context 'when merge_request_eager_fetch_ref is disabled' do
+      before do
+        stub_feature_flags(merge_request_eager_fetch_ref: false)
+      end
+
+      it_behaves_like 'when source and target projects are different', false
+    end
+
     context 'when user sets source project id' do
       let(:another_project) { create(:project) }