Move caught-up check for Ci::RegisterJobService to Core

app/services/ci/register_job_service.rb

@@ -22,11 +22,27 @@ module Ci
     end
 
     def execute(params = {})
+      db_all_caught_up = ::Gitlab::Database::LoadBalancing::Sticking.all_caught_up?(:runner, runner.id)
+
       @metrics.increment_queue_operation(:queue_attempt)
 
-      @metrics.observe_queue_time(:process, @runner.runner_type) do
+      result = @metrics.observe_queue_time(:process, @runner.runner_type) do
         process_queue(params)
       end
+
+      # Since we execute this query against replica it might lead to false-positive
+      # We might receive the positive response: "hi, we don't have any more builds for you".
+      # This might not be true. If our DB replica is not up-to date with when runner event was generated
+      # we might still have some CI builds to be picked. Instead we should say to runner:
+      # "Hi, we don't have any more builds now,  but not everything is right anyway, so try again".
+      # Runner will retry, but again, against replica, and again will check if replication lag did catch-up.
+      if !db_all_caught_up && !result.build
+        metrics.increment_queue_operation(:queue_replication_lag)
+
+        ::Ci::RegisterJobService::Result.new(nil, false) # rubocop:disable Cop/AvoidReturnFromBlocks
+      else
+        result
+      end
     end
 
     private

ee/app/services/ee/ci/register_job_service.rb

@@ -10,24 +10,6 @@ module EE
       extend ActiveSupport::Concern
       extend ::Gitlab::Utils::Override
 
-      def execute(params = {})
-        db_all_caught_up = ::Gitlab::Database::LoadBalancing::Sticking.all_caught_up?(:runner, runner.id)
-
-        super.tap do |result|
-          # Since we execute this query against replica it might lead to false-positive
-          # We might receive the positive response: "hi, we don't have any more builds for you".
-          # This might not be true. If our DB replica is not up-to date with when runner event was generated
-          # we might still have some CI builds to be picked. Instead we should say to runner:
-          # "Hi, we don't have any more builds now,  but not everything is right anyway, so try again".
-          # Runner will retry, but again, against replica, and again will check if replication lag did catch-up.
-          if !db_all_caught_up && !result.build
-            metrics.increment_queue_operation(:queue_replication_lag)
-
-            return ::Ci::RegisterJobService::Result.new(nil, false) # rubocop:disable Cop/AvoidReturnFromBlocks
-          end
-        end
-      end
-
       def builds_for_shared_runner
         # if disaster recovery is enabled, we disable quota
         if ::Feature.enabled?(:ci_queueing_disaster_recovery, runner, type: :ops, default_enabled: :yaml)

ee/spec/services/ci/register_job_service_spec.rb

@@ -10,34 +10,6 @@ RSpec.describe Ci::RegisterJobService do
   let!(:pending_build) { create :ci_build, pipeline: pipeline }
 
   describe '#execute' do
-    context 'checks database loadbalancing stickiness' do
-      subject { described_class.new(shared_runner).execute }
-
-      before do
-        project.update!(shared_runners_enabled: false)
-      end
-
-      it 'result is valid if replica did caught-up' do
-        allow(Gitlab::Database::LoadBalancing).to receive(:enable?)
-          .and_return(true)
-
-        expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:all_caught_up?)
-          .with(:runner, shared_runner.id) { true }
-
-        expect(subject).to be_valid
-      end
-
-      it 'result is invalid if replica did not caught-up' do
-        allow(Gitlab::Database::LoadBalancing).to receive(:enable?)
-          .and_return(true)
-
-        expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:all_caught_up?)
-          .with(:runner, shared_runner.id) { false }
-
-        expect(subject).not_to be_valid
-      end
-    end
-
     context 'shared runners minutes limit' do
       subject { described_class.new(shared_runner).execute.build }
 

spec/services/ci/register_job_service_spec.rb

@@ -14,6 +14,34 @@ module Ci
     let!(:pending_job) { create(:ci_build, pipeline: pipeline) }
 
     describe '#execute' do
+      context 'checks database loadbalancing stickiness' do
+        subject { described_class.new(shared_runner).execute }
+
+        before do
+          project.update!(shared_runners_enabled: false)
+        end
+
+        it 'result is valid if replica did caught-up' do
+          allow(Gitlab::Database::LoadBalancing).to receive(:enable?)
+            .and_return(true)
+
+          expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:all_caught_up?)
+            .with(:runner, shared_runner.id) { true }
+
+          expect(subject).to be_valid
+        end
+
+        it 'result is invalid if replica did not caught-up' do
+          allow(Gitlab::Database::LoadBalancing).to receive(:enable?)
+            .and_return(true)
+
+          expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:all_caught_up?)
+            .with(:runner, shared_runner.id) { false }
+
+          expect(subject).not_to be_valid
+        end
+      end
+
       shared_examples 'handles runner assignment' do
         context 'runner follow tag list' do
           it "picks build with the same tag" do