Move LB call in overriden method for MatchingMergeRequest to Core

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
-
-        TOTAL_METRIC = :gitlab_merge_request_match_total
-        STALE_METRIC = :gitlab_merge_request_match_stale_secondary
-
-        # rubocop:disable Gitlab/ModuleWithInstanceVariables
-        override :match?
-        def match?
-          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.
-          track_session_metrics do
-            if ::Feature.enabled?(:load_balancing_atomic_replica, @project, default_enabled: :yaml)
-              ::Gitlab::Database::LoadBalancing::Sticking.select_valid_host(:project, @project.id)
-            else
-              ::Gitlab::Database::LoadBalancing::Sticking.unstick_or_continue_sticking(:project, @project.id)
-            end
-          end
-
-          super
-        end
-        # rubocop:disable Gitlab/ModuleWithInstanceVariables
-
-        private
-
-        def track_session_metrics
-          before = ::Gitlab::Database::LoadBalancing::Session.current.use_primary?
-
-          yield
-
-          after = ::Gitlab::Database::LoadBalancing::Session.current.use_primary?
-
-          increment_attempt_count
-
-          if !before && after
-            increment_stale_secondary_count
-          end
-        end
-
-        def increment_attempt_count
-          total_counter.increment
-        end
-
-        def increment_stale_secondary_count
-          stale_counter.increment
-        end
-
-        def total_counter
-          @total_counter ||= ::Gitlab::Metrics.counter(TOTAL_METRIC, 'Total number of merge request match attempts')
-        end
-
-        def stale_counter
-          @stale_counter ||= ::Gitlab::Metrics.counter(STALE_METRIC, 'Total number of merge request match attempts with lagging secondary')
-        end
-      end
-    end
-  end
-end

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
-
-    let(:total_counter) { subject.send(:total_counter) }
-    let(:stale_counter) { subject.send(:stale_counter) }
-
-    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?).at_least(:once).and_return(false)
-        expect(::Gitlab::Database::LoadBalancing::Sticking).not_to receive(:unstick_or_continue_sticking)
-        expect(::Gitlab::Database::LoadBalancing::Sticking).not_to receive(:select_valid_replicas)
-      end
-
-      it 'does not attempt to stick to primary' do
-        expect(subject.match?).to be true
-      end
-
-      it 'increments no counters' do
-        expect { subject.match? }
-          .to change { total_counter.get }.by(0)
-          .and change { stale_counter.get }.by(0)
-      end
-    end
-
-    context 'with load balancing enabled', :request_store, :redis do
-      let(:session) { ::Gitlab::Database::LoadBalancing::Session.current }
-      let(:all_caught_up) { true }
-
-      before do
-        expect(::Gitlab::Database::LoadBalancing).to receive(:enable?).at_least(:once).and_return(true)
-        allow(::Gitlab::Database::LoadBalancing::Sticking).to receive(:all_caught_up?).and_return(all_caught_up)
-      end
-
-      shared_examples 'secondary that has caught up to a primary' do
-        it 'continues to use the secondary' do
-          expect(session.use_primary?).to be false
-          expect(subject.match?).to be true
-        end
-
-        it 'only increments total counter' do
-          expect { subject.match? }
-            .to change { total_counter.get }.by(1)
-            .and change { stale_counter.get }.by(0)
-        end
-      end
-
-      shared_examples 'secondary that is lagging primary' do
-        it 'sticks to the primary' do
-          expect(subject.match?).to be true
-          expect(session.use_primary?).to be true
-        end
-
-        it 'increments both total and stale counters' do
-          expect { subject.match? }
-            .to change { total_counter.get }.by(1)
-            .and change { stale_counter.get }.by(1)
-        end
-      end
-
-      context 'with load_balancing_atomic_replica feature flag enabled' do
-        before do
-          stub_feature_flags(load_balancing_atomic_replica: true)
-
-          expect(::Gitlab::Database::LoadBalancing::Sticking).to receive(:select_valid_host).with(:project, project.id).and_call_original
-          allow(::Gitlab::Database::LoadBalancing::Sticking).to receive(:select_caught_up_replicas).with(:project, project.id).and_return(all_caught_up)
-        end
-
-        it_behaves_like 'secondary that has caught up to a primary'
-
-        context 'on secondary behind primary' do
-          let(:all_caught_up) { false }
-
-          it_behaves_like 'secondary that is lagging primary'
-        end
-      end
-
-      context 'with load_balancing_atomic_replica feature flag disabled' do
-        before do
-          stub_feature_flags(load_balancing_atomic_replica: false)
-
-          expect(::Gitlab::Database::LoadBalancing::Sticking).not_to receive(:select_valid_host)
-          expect(::Gitlab::Database::LoadBalancing::Sticking).to receive(:unstick_or_continue_sticking).and_call_original
-          allow(::Gitlab::Database::LoadBalancing::Sticking).to receive(:all_caught_up?).and_return(all_caught_up)
-        end
-
-        it_behaves_like 'secondary that has caught up to a primary'
-
-        context 'on secondary behind primary' do
-          let(:all_caught_up) { false }
-
-          it_behaves_like 'secondary that is lagging primary'
-        end
-      end
-    end
-  end
-end

lib/gitlab/checks/matching_merge_request.rb

@@ -3,22 +3,78 @@
 module Gitlab
   module Checks
     class MatchingMergeRequest
+      TOTAL_METRIC = :gitlab_merge_request_match_total
+      STALE_METRIC = :gitlab_merge_request_match_stale_secondary
+
       def initialize(newrev, branch_name, project)
         @newrev = newrev
         @branch_name = branch_name
         @project = project
       end
 
-      # rubocop: disable CodeReuse/ActiveRecord
       def match?
+        if ::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.
+          track_session_metrics do
+            if ::Feature.enabled?(:load_balancing_atomic_replica, @project, default_enabled: :yaml)
+              ::Gitlab::Database::LoadBalancing::Sticking.select_valid_host(:project, @project.id)
+            else
+              ::Gitlab::Database::LoadBalancing::Sticking.unstick_or_continue_sticking(:project, @project.id)
+            end
+          end
+        end
+
+        # rubocop: disable CodeReuse/ActiveRecord
         @project.merge_requests
           .with_state(:locked)
           .where(in_progress_merge_commit_sha: @newrev, target_branch: @branch_name)
           .exists?
+        # rubocop: enable CodeReuse/ActiveRecord
+      end
+
+      private
+
+      def track_session_metrics
+        before = ::Gitlab::Database::LoadBalancing::Session.current.use_primary?
+
+        yield
+
+        after = ::Gitlab::Database::LoadBalancing::Session.current.use_primary?
+
+        increment_attempt_count
+
+        if !before && after
+          increment_stale_secondary_count
+        end
+      end
+
+      def increment_attempt_count
+        total_counter.increment
+      end
+
+      def increment_stale_secondary_count
+        stale_counter.increment
+      end
+
+      def total_counter
+        @total_counter ||= ::Gitlab::Metrics.counter(TOTAL_METRIC, 'Total number of merge request match attempts')
+      end
+
+      def stale_counter
+        @stale_counter ||= ::Gitlab::Metrics.counter(STALE_METRIC, 'Total number of merge request match attempts with lagging secondary')
       end
-      # rubocop: enable CodeReuse/ActiveRecord
     end
   end
 end
-
-Gitlab::Checks::MatchingMergeRequest.prepend_mod_with('Gitlab::Checks::MatchingMergeRequest')

spec/lib/gitlab/checks/matching_merge_request_spec.rb

@@ -18,6 +18,9 @@ RSpec.describe Gitlab::Checks::MatchingMergeRequest do
 
     subject { described_class.new(newrev, target_branch, project) }
 
+    let(:total_counter) { subject.send(:total_counter) }
+    let(:stale_counter) { subject.send(:stale_counter) }
+
     it 'matches a merge request' do
       expect(subject.match?).to be true
     end
@@ -27,5 +30,94 @@ RSpec.describe Gitlab::Checks::MatchingMergeRequest do
 
       expect(matcher.match?).to be false
     end
+
+    context 'with load balancing disabled', :request_store, :redis do
+      before do
+        expect(::Gitlab::Database::LoadBalancing).to receive(:enable?).at_least(:once).and_return(false)
+        expect(::Gitlab::Database::LoadBalancing::Sticking).not_to receive(:unstick_or_continue_sticking)
+        expect(::Gitlab::Database::LoadBalancing::Sticking).not_to receive(:select_valid_replicas)
+      end
+
+      it 'does not attempt to stick to primary' do
+        expect(subject.match?).to be true
+      end
+
+      it 'increments no counters' do
+        expect { subject.match? }
+          .to change { total_counter.get }.by(0)
+          .and change { stale_counter.get }.by(0)
+      end
+    end
+
+    context 'with load balancing enabled', :request_store, :redis do
+      let(:session) { ::Gitlab::Database::LoadBalancing::Session.current }
+      let(:all_caught_up) { true }
+
+      before do
+        expect(::Gitlab::Database::LoadBalancing).to receive(:enable?).at_least(:once).and_return(true)
+        allow(::Gitlab::Database::LoadBalancing::Sticking).to receive(:all_caught_up?).and_return(all_caught_up)
+      end
+
+      shared_examples 'secondary that has caught up to a primary' do
+        it 'continues to use the secondary' do
+          expect(session.use_primary?).to be false
+          expect(subject.match?).to be true
+        end
+
+        it 'only increments total counter' do
+          expect { subject.match? }
+            .to change { total_counter.get }.by(1)
+            .and change { stale_counter.get }.by(0)
+        end
+      end
+
+      shared_examples 'secondary that is lagging primary' do
+        it 'sticks to the primary' do
+          expect(subject.match?).to be true
+          expect(session.use_primary?).to be true
+        end
+
+        it 'increments both total and stale counters' do
+          expect { subject.match? }
+            .to change { total_counter.get }.by(1)
+            .and change { stale_counter.get }.by(1)
+        end
+      end
+
+      context 'with load_balancing_atomic_replica feature flag enabled' do
+        before do
+          stub_feature_flags(load_balancing_atomic_replica: true)
+
+          expect(::Gitlab::Database::LoadBalancing::Sticking).to receive(:select_valid_host).with(:project, project.id).and_call_original
+          allow(::Gitlab::Database::LoadBalancing::Sticking).to receive(:select_caught_up_replicas).with(:project, project.id).and_return(all_caught_up)
+        end
+
+        it_behaves_like 'secondary that has caught up to a primary'
+
+        context 'on secondary behind primary' do
+          let(:all_caught_up) { false }
+
+          it_behaves_like 'secondary that is lagging primary'
+        end
+      end
+
+      context 'with load_balancing_atomic_replica feature flag disabled' do
+        before do
+          stub_feature_flags(load_balancing_atomic_replica: false)
+
+          expect(::Gitlab::Database::LoadBalancing::Sticking).not_to receive(:select_valid_host)
+          expect(::Gitlab::Database::LoadBalancing::Sticking).to receive(:unstick_or_continue_sticking).and_call_original
+          allow(::Gitlab::Database::LoadBalancing::Sticking).to receive(:all_caught_up?).and_return(all_caught_up)
+        end
+
+        it_behaves_like 'secondary that has caught up to a primary'
+
+        context 'on secondary behind primary' do
+          let(:all_caught_up) { false }
+
+          it_behaves_like 'secondary that is lagging primary'
+        end
+      end
+    end
   end
 end