Merge branch 'sh-stick-primary-matching-mrs' into 'master'

Fix merge pre-receive errors when load balancing in use

See merge request gitlab-org/gitlab!42435

app/models/merge_request.rb

@@ -1302,6 +1302,14 @@ class MergeRequest < ApplicationRecord
     unlock_mr
   end
 
+  def update_and_mark_in_progress_merge_commit_sha(commit_id)
+    self.update(in_progress_merge_commit_sha: commit_id)
+    # Since another process checks for matching merge request, we need
+    # to make it possible to detect whether the query should go to the
+    # primary.
+    target_project.mark_primary_write_location
+  end
+
   def diverged_commits_count
     cache = Rails.cache.read(:"merge_request_#{id}_diverged_commits")
 

app/models/project.rb

@@ -2292,6 +2292,10 @@ class Project < ApplicationRecord
     []
   end
 
+  def mark_primary_write_location
+    # Overriden in EE
+  end
+
   def toggle_ci_cd_settings!(settings_attribute)
     ci_cd_settings.toggle!(settings_attribute)
   end

app/models/repository.rb

@@ -853,7 +853,7 @@ class Repository
   def merge(user, source_sha, merge_request, message)
     with_cache_hooks do
       raw_repository.merge(user, source_sha, merge_request.target_branch, message) do |commit_id|
-        merge_request.update(in_progress_merge_commit_sha: commit_id)
+        merge_request.update_and_mark_in_progress_merge_commit_sha(commit_id)
         nil # Return value does not matter.
       end
     end
@@ -873,7 +873,7 @@ class Repository
     their_commit_id = commit(source)&.id
     raise 'Invalid merge source' if their_commit_id.nil?
 
-    merge_request&.update(in_progress_merge_commit_sha: their_commit_id)
+    merge_request&.update_and_mark_in_progress_merge_commit_sha(their_commit_id)
 
     with_cache_hooks { raw.ff_merge(user, their_commit_id, target_branch) }
   end

app/services/merge_requests/ff_merge_service.rb

@@ -27,7 +27,7 @@ module MergeRequests
     rescue StandardError => e
       raise MergeError, "Something went wrong during merge: #{e.message}"
     ensure
-      merge_request.update(in_progress_merge_commit_sha: nil)
+      merge_request.update_and_mark_in_progress_merge_commit_sha(nil)
     end
   end
 end

app/services/merge_requests/merge_service.rb

@@ -84,7 +84,7 @@ module MergeRequests
 
       merge_request.update!(merge_commit_sha: commit_id)
     ensure
-      merge_request.update_column(:in_progress_merge_commit_sha, nil)
+      merge_request.update_and_mark_in_progress_merge_commit_sha(nil)
     end
 
     def try_merge

ee/app/models/ee/project.rb

@@ -350,6 +350,11 @@ module EE
         feature_available?(:github_project_service_integration)
     end
 
+    override :mark_primary_write_location
+    def mark_primary_write_location
+      ::Gitlab::Database::LoadBalancing::Sticking.mark_primary_write_location(:project, self.id)
+    end
+
     override :add_import_job
     def add_import_job
       return if gitlab_custom_project_template_import?

ee/changelogs/unreleased/sh-stick-primary-matching-mrs.yml

+---
+title: Fix merge pre-receive errors when load balancing in use
+merge_request: 42435
+author:
+type: fixed

ee/config/feature_flags/development/matching_merge_request_db_sync.yml

+---
+name: matching_merge_request_db_sync
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/42435
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/254488
+group: group::create
+type: development
+default_enabled: false
\ No newline at end of file

ee/lib/ee/gitlab/checks/matching_merge_request.rb

+# frozen_string_literal: true
+
+module EE
+  module Gitlab
+    module Checks
+      module MatchingMergeRequest
+        extend ::Gitlab::Utils::Override
+
+        override :match?
+        def match?
+          return super unless ::Feature.enabled?(:matching_merge_request_db_sync)
+          return super unless ::Gitlab::Database::LoadBalancing.enable?
+
+          # When a user merges a merge request, the following sequence happens:
+          #
+          # 1. Sidekiq: MergeService runs and updates the merge request in a locked state.
+          # 2. Gitaly: The UserMergeBranch RPC runs.
+          # 3. Gitaly (gitaly-ruby): This RPC calls the pre-receive hook.
+          # 4. Rails: This hook makes an API request to /api/v4/internal/allowed.
+          # 5. Rails: This API check does a SQL query for locked merge
+          #    requests with a matching SHA.
+          #
+          # Since steps 1 and 5 will happen on different database
+          # sessions, replication lag could erroneously cause step 5 to
+          # report no matching merge requests. To avoid this, we check
+          # the write location to ensure the replica can make this query.
+          ::Gitlab::Database::LoadBalancing::Sticking.unstick_or_continue_sticking(:project, @project.id) # rubocop:disable Gitlab/ModuleWithInstanceVariables
+          super
+        end
+      end
+    end
+  end
+end

ee/lib/gitlab/database/load_balancing/sticking.rb

@@ -49,11 +49,20 @@ module Gitlab
         def self.stick(namespace, id)
           return unless LoadBalancing.enable?
 
+          mark_primary_write_location_helper(namespace, id)
+          Session.current.use_primary!
+        end
+
+        def self.mark_primary_write_location(namespace, id)
+          return unless LoadBalancing.enable?
+
+          mark_primary_write_location_helper(namespace, id)
+        end
+
+        def self.mark_primary_write_location_helper(namespace, id)
           location = load_balancer.primary_write_location
 
           set_write_location_for(namespace, id, location)
-
-          Session.current.use_primary!
         end
 
         # Stops sticking to the primary.

ee/spec/lib/gitlab/checks/matching_merge_request_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Checks::MatchingMergeRequest do
+  describe '#match?' do
+    let_it_be(:newrev) { '012345678' }
+    let_it_be(:target_branch) { 'feature' }
+    let_it_be(:project) { create(:project, :repository) }
+    let_it_be(:locked_merge_request) do
+      create(:merge_request,
+        :locked,
+        source_project: project,
+        target_project: project,
+        target_branch: target_branch,
+        in_progress_merge_commit_sha: newrev)
+    end
+
+    subject { described_class.new(newrev, target_branch, project) }
+
+    context 'with load balancing disabled', :request_store, :redis do
+      before do
+        expect(::Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false)
+      end
+
+      it 'does not attempt to stick to primary' do
+        expect(::Gitlab::Database::LoadBalancing::Sticking).not_to receive(:unstick_or_continue_sticking)
+
+        expect(subject.match?).to be true
+      end
+    end
+
+    context 'with load balancing enabled', :request_store, :redis do
+      before do
+        allow(::Gitlab::Database::LoadBalancing::Sticking).to receive(:all_caught_up?).and_return(false)
+      end
+
+      context 'with feature flag disabled' do
+        before do
+          stub_feature_flags(matching_merge_request_db_sync: false)
+        end
+
+        it 'does not check load balancing state' do
+          expect(::Gitlab::Database::LoadBalancing).not_to receive(:enable?)
+          expect(::Gitlab::Database::LoadBalancing::Sticking).not_to receive(:unstick_or_continue_sticking)
+
+          expect(subject.match?).to be true
+        end
+      end
+
+      it 'sticks to the primary' do
+        session = ::Gitlab::Database::LoadBalancing::Session.current
+        expect(::Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
+        expect(::Gitlab::Database::LoadBalancing::Sticking).to receive(:unstick_or_continue_sticking).and_call_original
+
+        expect(subject.match?).to be true
+        expect(session.use_primary?).to be true
+      end
+    end
+  end
+end

ee/spec/lib/gitlab/database/load_balancing/sticking_spec.rb

@@ -156,6 +156,37 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
     end
   end
 
+  describe '.mark_primary_write_location' do
+    context 'when load balancing is disabled' do
+      before do
+        allow(Gitlab::Database::LoadBalancing).to receive(:enable?)
+          .and_return(false)
+      end
+
+      it 'does not attempt to write' do
+        expect(described_class).not_to receive(:set_write_location_for)
+
+        described_class.mark_primary_write_location(:user, 42)
+      end
+    end
+
+    context 'when load balancing is enabled' do
+      it 'updates the write location' do
+        allow(Gitlab::Database::LoadBalancing).to receive(:enable?)
+          .and_return(true)
+
+        lb = double(:lb, primary_write_location: 'foo')
+
+        allow(described_class).to receive(:load_balancer).and_return(lb)
+
+        expect(described_class).to receive(:set_write_location_for)
+          .with(:user, 42, 'foo')
+
+        described_class.mark_primary_write_location(:user, 42)
+      end
+    end
+  end
+
   describe '.unstick' do
     it 'removes the sticking data from Redis' do
       described_class.set_write_location_for(:user, 4, 'foo')

ee/spec/models/project_spec.rb

@@ -2584,4 +2584,12 @@ RSpec.describe Project do
       end
     end
   end
+
+  describe '#mark_primary_write_location' do
+    it 'marks the location with project ID' do
+      expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:mark_primary_write_location).with(:project, project.id)
+
+      project.mark_primary_write_location
+    end
+  end
 end

lib/gitlab/checks/matching_merge_request.rb

@@ -20,3 +20,5 @@ module Gitlab
     end
   end
 end
+
+Gitlab::Checks::MatchingMergeRequest.prepend_if_ee('EE::Gitlab::Checks::MatchingMergeRequest')

spec/lib/gitlab/checks/matching_merge_request_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Checks::MatchingMergeRequest do
+  describe '#match?' do
+    let_it_be(:newrev) { '012345678' }
+    let_it_be(:target_branch) { 'feature' }
+    let_it_be(:project) { create(:project, :repository) }
+    let_it_be(:locked_merge_request) do
+      create(:merge_request,
+        :locked,
+        source_project: project,
+        target_project: project,
+        target_branch: target_branch,
+        in_progress_merge_commit_sha: newrev)
+    end
+
+    subject { described_class.new(newrev, target_branch, project) }
+
+    it 'matches a merge request' do
+      expect(subject.match?).to be true
+    end
+
+    it 'does not match any merge request' do
+      matcher = described_class.new(newrev, 'test', project)
+
+      expect(matcher.match?).to be false
+    end
+  end
+end

spec/models/merge_request_spec.rb

@@ -4299,4 +4299,18 @@ RSpec.describe MergeRequest, factory_default: :keep do
       expect(merge_request.allows_reviewers?).to be(true)
     end
   end
+
+  describe '#update_and_mark_in_progress_merge_commit_sha' do
+    let(:ref) { subject.target_project.repository.commit.id }
+
+    before do
+      expect(subject.target_project).to receive(:mark_primary_write_location)
+    end
+
+    it 'updates commit ID' do
+      expect { subject.update_and_mark_in_progress_merge_commit_sha(ref) }
+        .to change { subject.in_progress_merge_commit_sha }
+        .from(nil).to(ref)
+    end
+  end
 end

spec/services/merge_requests/ff_merge_service_spec.rb

@@ -72,15 +72,22 @@ RSpec.describe MergeRequests::FfMergeService do
       end
 
       it 'does not update squash_commit_sha if it is not a squash' do
+        expect(merge_request).to receive(:update_and_mark_in_progress_merge_commit_sha).twice.and_call_original
+
         expect { execute_ff_merge }.not_to change { merge_request.squash_commit_sha }
+        expect(merge_request.in_progress_merge_commit_sha).to be_nil
       end
 
       it 'updates squash_commit_sha if it is a squash' do
+        expect(merge_request).to receive(:update_and_mark_in_progress_merge_commit_sha).twice.and_call_original
+
         merge_request.update!(squash: true)
 
         expect { execute_ff_merge }
           .to change { merge_request.squash_commit_sha }
           .from(nil)
+
+        expect(merge_request.in_progress_merge_commit_sha).to be_nil
       end
     end
 

spec/services/merge_requests/merge_service_spec.rb

@@ -22,6 +22,7 @@ RSpec.describe MergeRequests::MergeService do
     context 'valid params' do
       before do
         allow(service).to receive(:execute_hooks)
+        expect(merge_request).to receive(:update_and_mark_in_progress_merge_commit_sha).twice.and_call_original
 
         perform_enqueued_jobs do
           service.execute(merge_request)