Move LB usage from runner to Core

app/models/ci/runner.rb

@@ -338,6 +338,14 @@ module Ci
     end
 
     def tick_runner_queue
+      ##
+      # We only stick a runner to primary database to be able to detect the
+      # replication lag in `EE::Ci::RegisterJobService#execute`. The
+      # intention here is not to execute `Ci::RegisterJobService#execute` on
+      # the primary database.
+      #
+      ::Gitlab::Database::LoadBalancing::Sticking.stick(:runner, id)
+
       SecureRandom.hex.tap do |new_update|
         ::Gitlab::Workhorse.set_key_and_notify(runner_queue_key, new_update,
           expire: RUNNER_QUEUE_EXPIRY_TIME, overwrite: true)
@@ -355,13 +363,20 @@ module Ci
     end
 
     def heartbeat(values)
-      values = values&.slice(:version, :revision, :platform, :architecture, :ip_address, :config) || {}
-      values[:contacted_at] = Time.current
+      ##
+      # We can safely ignore writes performed by a runner heartbeat. We do
+      # not want to upgrade database connection proxy to use the primary
+      # database after heartbeat write happens.
+      #
+      ::Gitlab::Database::LoadBalancing::Session.without_sticky_writes do
+        values = values&.slice(:version, :revision, :platform, :architecture, :ip_address, :config) || {}
+        values[:contacted_at] = Time.current
 
-      cache_attributes(values)
+        cache_attributes(values)
 
-      # We save data without validation, it will always change due to `contacted_at`
-      self.update_columns(values) if persist_cached_data?
+        # We save data without validation, it will always change due to `contacted_at`
+        self.update_columns(values) if persist_cached_data?
+      end
     end
 
     def pick_build!(build)

app/services/ci/register_job_service.rb

@@ -128,6 +128,12 @@ module Ci
     # rubocop: enable CodeReuse/ActiveRecord
 
     def retrieve_queue(queue_query_proc)
+      ##
+      # We want to reset a load balancing session to discard the side
+      # effects of writes that could have happened prior to this moment.
+      #
+      ::Gitlab::Database::LoadBalancing::Session.clear_session
+
       @metrics.observe_queue_time(:retrieve, @runner.runner_type) do
         queue_query_proc.call
       end

ee/app/models/ee/ci/runner.rb

@@ -5,27 +5,6 @@ module EE
     module Runner
       extend ActiveSupport::Concern
 
-      def tick_runner_queue
-        ##
-        # We only stick a runner to primary database to be able to detect the
-        # replication lag in `EE::Ci::RegisterJobService#execute`. The
-        # intention here is not to execute `Ci::RegisterJobService#execute` on
-        # the primary database.
-        #
-        ::Gitlab::Database::LoadBalancing::Sticking.stick(:runner, id)
-
-        super
-      end
-
-      def heartbeat(values)
-        ##
-        # We can safely ignore writes performed by a runner heartbeat. We do
-        # not want to upgrade database connection proxy to use the primary
-        # database after heartbeat write happens.
-        #
-        ::Gitlab::Database::LoadBalancing::Session.without_sticky_writes { super }
-      end
-
       def minutes_cost_factor(visibility_level)
         ::Gitlab::Ci::Minutes::CostFactor.new(runner_matcher).for_visibility(visibility_level)
       end

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

@@ -28,16 +28,6 @@ module EE
         end
       end
 
-      def retrieve_queue(queue_query_proc)
-        ##
-        # We want to reset a load balancing session to discard the side
-        # effects of writes that could have happened prior to this moment.
-        #
-        ::Gitlab::Database::LoadBalancing::Session.clear_session
-
-        super
-      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/models/ee/ci/runner_spec.rb

@@ -3,22 +3,6 @@
 require 'spec_helper'
 
 RSpec.describe EE::Ci::Runner do
-  describe '#tick_runner_queue' do
-    it 'sticks the runner to the primary and calls the original method' do
-      runner = create(:ci_runner)
-
-      allow(Gitlab::Database::LoadBalancing).to receive(:enable?)
-        .and_return(true)
-
-      expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:stick)
-        .with(:runner, runner.id)
-
-      expect(Gitlab::Workhorse).to receive(:set_key_and_notify)
-
-      runner.tick_runner_queue
-    end
-  end
-
   describe '#minutes_cost_factor' do
     subject { runner.minutes_cost_factor(visibility_level) }
 

spec/models/ci/runner_spec.rb

@@ -365,6 +365,22 @@ RSpec.describe Ci::Runner do
     it { is_expected.to eq([@runner1])}
   end
 
+  describe '#tick_runner_queue' do
+    it 'sticks the runner to the primary and calls the original method' do
+      runner = create(:ci_runner)
+
+      allow(Gitlab::Database::LoadBalancing).to receive(:enable?)
+        .and_return(true)
+
+      expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:stick)
+        .with(:runner, runner.id)
+
+      expect(Gitlab::Workhorse).to receive(:set_key_and_notify)
+
+      runner.tick_runner_queue
+    end
+  end
+
   describe '#can_pick?' do
     using RSpec::Parameterized::TableSyntax