From b6b20a4d69e93c1726c97f82c3e533e5c28c854a Mon Sep 17 00:00:00 2001
From: Thong Kuah <tkuah@gitlab.com>
Date: Thu, 27 May 2021 11:44:22 +1200
Subject: [PATCH 01/10] Enable LB for Core by default

---
 lib/gitlab/database/load_balancing.rb         |  7 -----
 .../gitlab/database/load_balancing_spec.rb    | 29 -------------------
 2 files changed, 36 deletions(-)

diff --git a/lib/gitlab/database/load_balancing.rb b/lib/gitlab/database/load_balancing.rb
index ea4628ab757f54..88743cd2e75a5e 100644
--- a/lib/gitlab/database/load_balancing.rb
+++ b/lib/gitlab/database/load_balancing.rb
@@ -97,16 +97,9 @@ def self.enable?
       # posting the write location of the database if load balancing is
       # configured.
       def self.configured?
-        return false unless feature_available?
-
         hosts.any? || service_discovery_enabled?
       end
 
-      # Temporarily disabled for FOSS until move from EE to FOSS is complete
-      def self.feature_available?
-        Gitlab.ee? || Gitlab::Utils.to_boolean(ENV['ENABLE_LOAD_BALANCING_FOR_FOSS'], default: false)
-      end
-
       def self.start_service_discovery
         return unless service_discovery_enabled?
 
diff --git a/spec/lib/gitlab/database/load_balancing_spec.rb b/spec/lib/gitlab/database/load_balancing_spec.rb
index c3dcfa3eb4a542..b26431417d728e 100644
--- a/spec/lib/gitlab/database/load_balancing_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing_spec.rb
@@ -5,10 +5,6 @@
 RSpec.describe Gitlab::Database::LoadBalancing do
   include_context 'clear DB Load Balancing configuration'
 
-  before do
-    stub_env('ENABLE_LOAD_BALANCING_FOR_FOSS', 'true')
-  end
-
   describe '.proxy' do
     context 'when configured' do
       before do
@@ -182,31 +178,6 @@
         expect(described_class.enable?).to eq(true)
       end
     end
-
-    context 'FOSS' do
-      before do
-        allow(Gitlab).to receive(:ee?).and_return(false)
-
-        stub_env('ENABLE_LOAD_BALANCING_FOR_FOSS', 'false')
-      end
-
-      it 'is disabled' do
-        expect(described_class.enable?).to eq(false)
-      end
-    end
-
-    context 'EE' do
-      before do
-        allow(Gitlab).to receive(:ee?).and_return(true)
-      end
-
-      it 'is enabled' do
-        allow(described_class).to receive(:hosts).and_return(%w(foo))
-        allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(false)
-
-        expect(described_class.enable?).to eq(true)
-      end
-    end
   end
 
   describe '.configured?' do
-- 
GitLab


From f9a17e5555aacc1eab8fd0665476bc81ad157412 Mon Sep 17 00:00:00 2001
From: Thong Kuah <tkuah@gitlab.com>
Date: Thu, 27 May 2021 11:49:06 +1200
Subject: [PATCH 02/10] Move LB sticking for CI build to Core

---
 app/models/ci/build.rb          |  9 +++++++++
 ee/app/models/ee/ci/build.rb    |  8 --------
 ee/spec/models/ci/build_spec.rb | 14 --------------
 spec/models/ci/build_spec.rb    | 14 ++++++++++++++
 4 files changed, 23 insertions(+), 22 deletions(-)

diff --git a/app/models/ci/build.rb b/app/models/ci/build.rb
index 3c78b9f53d570a..e8fd039ed696f3 100644
--- a/app/models/ci/build.rb
+++ b/app/models/ci/build.rb
@@ -214,6 +214,8 @@ def persisted_environment
     before_save :ensure_token
     before_destroy { unscoped_project }
 
+    after_save :stick_build_if_status_changed
+
     after_create unless: :importing? do |build|
       run_after_commit { BuildHooksWorker.perform_async(build.id) }
     end
@@ -1082,6 +1084,13 @@ def run_status_commit_hooks!
 
     private
 
+    def stick_build_if_status_changed
+      return unless saved_change_to_status?
+      return unless running?
+
+      ::Gitlab::Database::LoadBalancing::Sticking.stick(:build, id)
+    end
+
     def status_commit_hooks
       @status_commit_hooks ||= []
     end
diff --git a/ee/app/models/ee/ci/build.rb b/ee/app/models/ee/ci/build.rb
index 7ed03cdd30f215..eea46ef67bd7e8 100644
--- a/ee/app/models/ee/ci/build.rb
+++ b/ee/app/models/ee/ci/build.rb
@@ -31,7 +31,6 @@ module Build
 
         has_many :security_scans, class_name: 'Security::Scan'
 
-        after_save :stick_build_if_status_changed
         after_commit :track_ci_secrets_management_usage, on: :create
         delegate :service_specification, to: :runner_session, allow_nil: true
 
@@ -62,13 +61,6 @@ def shared_runners_minutes_limit_enabled?
         project.shared_runners_minutes_limit_enabled? && runner&.minutes_cost_factor(project.visibility_level)&.positive?
       end
 
-      def stick_build_if_status_changed
-        return unless saved_change_to_status?
-        return unless running?
-
-        ::Gitlab::Database::LoadBalancing::Sticking.stick(:build, id)
-      end
-
       def log_geo_deleted_event
         # It is not needed to generate a Geo deleted event
         # since Legacy Artifacts are migrated to multi-build artifacts
diff --git a/ee/spec/models/ci/build_spec.rb b/ee/spec/models/ci/build_spec.rb
index d8c1b5e4d9cbc1..fa312aec2ba594 100644
--- a/ee/spec/models/ci/build_spec.rb
+++ b/ee/spec/models/ci/build_spec.rb
@@ -107,20 +107,6 @@
     end
   end
 
-  describe '#stick_build_if_status_changed' do
-    it 'sticks the build if the status changed' do
-      job = create(:ci_build, :pending)
-
-      allow(Gitlab::Database::LoadBalancing).to receive(:enable?)
-        .and_return(true)
-
-      expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:stick)
-        .with(:build, job.id)
-
-      job.update!(status: :running)
-    end
-  end
-
   describe '#variables' do
     subject { job.variables }
 
diff --git a/spec/models/ci/build_spec.rb b/spec/models/ci/build_spec.rb
index c2f49b1063983f..d5694020b5fe85 100644
--- a/spec/models/ci/build_spec.rb
+++ b/spec/models/ci/build_spec.rb
@@ -320,6 +320,20 @@
     end
   end
 
+  describe '#stick_build_if_status_changed' do
+    it 'sticks the build if the status changed' do
+      job = create(:ci_build, :pending)
+
+      allow(Gitlab::Database::LoadBalancing).to receive(:enable?)
+        .and_return(true)
+
+      expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:stick)
+        .with(:build, job.id)
+
+      job.update!(status: :running)
+    end
+  end
+
   describe '#enqueue' do
     let(:build) { create(:ci_build, :created) }
 
-- 
GitLab


From aafb0dd9df734e5eba861fffde24519b853c72c2 Mon Sep 17 00:00:00 2001
From: Thong Kuah <tkuah@gitlab.com>
Date: Thu, 27 May 2021 12:07:31 +1200
Subject: [PATCH 03/10] Move LB usage from runner to Core

---
 app/models/ci/runner.rb                       | 25 ++++++++++++---
 app/services/ci/register_job_service.rb       |  6 ++++
 ee/app/models/ee/ci/runner.rb                 | 32 ++++++-------------
 ee/app/services/ee/ci/register_job_service.rb | 10 ------
 ee/spec/models/ee/ci/runner_spec.rb           | 16 ----------
 spec/models/ci/runner_spec.rb                 | 16 ++++++++++
 6 files changed, 52 insertions(+), 53 deletions(-)

diff --git a/app/models/ci/runner.rb b/app/models/ci/runner.rb
index 144595ff7b0506..54cb34ef53b42e 100644
--- a/app/models/ci/runner.rb
+++ b/app/models/ci/runner.rb
@@ -338,6 +338,14 @@ def predefined_variables
     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 @@ def runner_queue_value_latest?(value)
     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)
diff --git a/app/services/ci/register_job_service.rb b/app/services/ci/register_job_service.rb
index f975457c34e1cc..04956c4813655e 100644
--- a/app/services/ci/register_job_service.rb
+++ b/app/services/ci/register_job_service.rb
@@ -128,6 +128,12 @@ def each_build(params, &blk)
     # 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
diff --git a/ee/app/models/ee/ci/runner.rb b/ee/app/models/ee/ci/runner.rb
index 0ec19d095f889e..ffd6a6203bdcd4 100644
--- a/ee/app/models/ee/ci/runner.rb
+++ b/ee/app/models/ee/ci/runner.rb
@@ -5,29 +5,17 @@ module Ci
     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)
+      def minutes_cost_factor(access_level)
+        return 0.0 unless instance_type?
 
-        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)
+        case access_level
+        when ::Gitlab::VisibilityLevel::PUBLIC
+          public_projects_minutes_cost_factor
+        when ::Gitlab::VisibilityLevel::PRIVATE, ::Gitlab::VisibilityLevel::INTERNAL
+          private_projects_minutes_cost_factor
+        else
+          raise ArgumentError, 'Invalid visibility level'
+        end
       end
 
       def visibility_levels_without_minutes_quota
diff --git a/ee/app/services/ee/ci/register_job_service.rb b/ee/app/services/ee/ci/register_job_service.rb
index d2a4acb2d17792..baf5ebd4a14033 100644
--- a/ee/app/services/ee/ci/register_job_service.rb
+++ b/ee/app/services/ee/ci/register_job_service.rb
@@ -28,16 +28,6 @@ def execute(params = {})
         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)
diff --git a/ee/spec/models/ee/ci/runner_spec.rb b/ee/spec/models/ee/ci/runner_spec.rb
index 4d04d9aabcf257..13cb9778430b96 100644
--- a/ee/spec/models/ee/ci/runner_spec.rb
+++ b/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) }
 
diff --git a/spec/models/ci/runner_spec.rb b/spec/models/ci/runner_spec.rb
index 491088a44a18eb..9ed9b10b1f38a9 100644
--- a/spec/models/ci/runner_spec.rb
+++ b/spec/models/ci/runner_spec.rb
@@ -365,6 +365,22 @@ def stub_redis_runner_contacted_at(value)
     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
 
-- 
GitLab


From 6ff510654f7f024c62811059b2cd6c5d54e7485a Mon Sep 17 00:00:00 2001
From: Thong Kuah <tkuah@gitlab.com>
Date: Thu, 27 May 2021 12:20:44 +1200
Subject: [PATCH 04/10] Move Project#mark_primary_write_location to Core

---
 app/models/project.rb          |  2 +-
 ee/app/models/ee/project.rb    |  5 -----
 ee/spec/models/project_spec.rb |  8 --------
 spec/models/project_spec.rb    | 10 ++++++++++
 4 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/app/models/project.rb b/app/models/project.rb
index 53054e5304ffbe..fe3c3375340ecf 100644
--- a/app/models/project.rb
+++ b/app/models/project.rb
@@ -2364,7 +2364,7 @@ def licensed_features
   end
 
   def mark_primary_write_location
-    # Overriden in EE
+    ::Gitlab::Database::LoadBalancing::Sticking.mark_primary_write_location(:project, self.id)
   end
 
   def toggle_ci_cd_settings!(settings_attribute)
diff --git a/ee/app/models/ee/project.rb b/ee/app/models/ee/project.rb
index 19c7fd1c69951f..8cf0c66dffd1f7 100644
--- a/ee/app/models/ee/project.rb
+++ b/ee/app/models/ee/project.rb
@@ -410,11 +410,6 @@ def github_external_pull_request_pipelines_available?
         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?
diff --git a/ee/spec/models/project_spec.rb b/ee/spec/models/project_spec.rb
index 254d2bd684d9ae..864b5380cd5b9a 100644
--- a/ee/spec/models/project_spec.rb
+++ b/ee/spec/models/project_spec.rb
@@ -2717,14 +2717,6 @@ def stub_default_url_options(host)
     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
-
   describe '#add_template_export_job' do
     it 'starts project template export job' do
       user = create(:user)
diff --git a/spec/models/project_spec.rb b/spec/models/project_spec.rb
index 6de1a246173375..886c381c475317 100644
--- a/spec/models/project_spec.rb
+++ b/spec/models/project_spec.rb
@@ -2854,6 +2854,16 @@ def subject
     end
   end
 
+  describe '#mark_primary_write_location' do
+    let(:project) { create(:project) }
+
+    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
+
   describe '#mark_stuck_remote_mirrors_as_failed!' do
     it 'fails stuck remote mirrors' do
       project = create(:project, :repository, :remote_mirror)
-- 
GitLab


From 672db1d38f251db95fb9f33797b51e542ee3499b Mon Sep 17 00:00:00 2001
From: Thong Kuah <tkuah@gitlab.com>
Date: Thu, 27 May 2021 12:29:00 +1200
Subject: [PATCH 05/10] Move non-stickiness in ProjectFeatureUsage to Core

---
 app/models/project_feature_usage.rb          | 14 +++---
 ee/app/models/ee/project_feature_usage.rb    | 15 ------
 ee/spec/models/project_feature_usage_spec.rb | 53 --------------------
 spec/models/project_feature_usage_spec.rb    | 50 ++++++++++++++++++
 4 files changed, 58 insertions(+), 74 deletions(-)
 delete mode 100644 ee/app/models/ee/project_feature_usage.rb
 delete mode 100644 ee/spec/models/project_feature_usage_spec.rb

diff --git a/app/models/project_feature_usage.rb b/app/models/project_feature_usage.rb
index d993db860c308d..520496c1527656 100644
--- a/app/models/project_feature_usage.rb
+++ b/app/models/project_feature_usage.rb
@@ -20,14 +20,16 @@ def jira_dvcs_integration_field(cloud: true)
   end
 
   def log_jira_dvcs_integration_usage(cloud: true)
-    integration_field = self.class.jira_dvcs_integration_field(cloud: cloud)
+    ::Gitlab::Database::LoadBalancing::Session.without_sticky_writes do
+      integration_field = self.class.jira_dvcs_integration_field(cloud: cloud)
 
-    # The feature usage is used only once later to query the feature usage in a
-    # long date range. Therefore, we just need to update the timestamp once per
-    # day
-    return if persisted? && updated_today?(integration_field)
+      # The feature usage is used only once later to query the feature usage in a
+      # long date range. Therefore, we just need to update the timestamp once per
+      # eay
+      break if persisted? && updated_today?(integration_field)
 
-    persist_jira_dvcs_usage(integration_field)
+      persist_jira_dvcs_usage(integration_field)
+    end
   end
 
   private
diff --git a/ee/app/models/ee/project_feature_usage.rb b/ee/app/models/ee/project_feature_usage.rb
deleted file mode 100644
index cd5f7e92b557b6..00000000000000
--- a/ee/app/models/ee/project_feature_usage.rb
+++ /dev/null
@@ -1,15 +0,0 @@
-# frozen_string_literal: true
-
-module EE
-  module ProjectFeatureUsage
-    extend ActiveSupport::Concern
-    extend ::Gitlab::Utils::Override
-
-    override :log_jira_dvcs_integration_usage
-    def log_jira_dvcs_integration_usage(**options)
-      ::Gitlab::Database::LoadBalancing::Session.without_sticky_writes do
-        super(**options)
-      end
-    end
-  end
-end
diff --git a/ee/spec/models/project_feature_usage_spec.rb b/ee/spec/models/project_feature_usage_spec.rb
deleted file mode 100644
index 9bc621f03e669e..00000000000000
--- a/ee/spec/models/project_feature_usage_spec.rb
+++ /dev/null
@@ -1,53 +0,0 @@
-# frozen_string_literal: true
-
-require 'spec_helper'
-
-RSpec.describe ProjectFeatureUsage, :request_store do
-  include_context 'clear DB Load Balancing configuration'
-
-  describe '#log_jira_dvcs_integration_usage' do
-    let!(:project) { create(:project) }
-
-    subject { project.feature_usage }
-
-    context 'database load balancing is configured' do
-      before do
-        # Do not pollute AR for other tests, but rather simulate effect of configure_proxy.
-        allow(ActiveRecord::Base.singleton_class).to receive(:prepend)
-        ::Gitlab::Database::LoadBalancing.configure_proxy
-        allow(ActiveRecord::Base).to receive(:connection).and_return(::Gitlab::Database::LoadBalancing.proxy)
-        ::Gitlab::Database::LoadBalancing::Session.clear_session
-      end
-
-      it 'logs Jira DVCS Cloud last sync' do
-        freeze_time do
-          subject.log_jira_dvcs_integration_usage
-
-          expect(subject.jira_dvcs_server_last_sync_at).to be_nil
-          expect(subject.jira_dvcs_cloud_last_sync_at).to be_like_time(Time.current)
-        end
-      end
-
-      it 'does not stick to primary' do
-        expect(::Gitlab::Database::LoadBalancing::Session.current).not_to be_performed_write
-        expect(::Gitlab::Database::LoadBalancing::Session.current).not_to be_using_primary
-
-        subject.log_jira_dvcs_integration_usage
-
-        expect(::Gitlab::Database::LoadBalancing::Session.current).to be_performed_write
-        expect(::Gitlab::Database::LoadBalancing::Session.current).not_to be_using_primary
-      end
-    end
-
-    context 'database load balancing is not cofigured' do
-      it 'logs Jira DVCS Cloud last sync' do
-        freeze_time do
-          subject.log_jira_dvcs_integration_usage
-
-          expect(subject.jira_dvcs_server_last_sync_at).to be_nil
-          expect(subject.jira_dvcs_cloud_last_sync_at).to be_like_time(Time.current)
-        end
-      end
-    end
-  end
-end
diff --git a/spec/models/project_feature_usage_spec.rb b/spec/models/project_feature_usage_spec.rb
index 4baa59535e4796..6ef407432b0382 100644
--- a/spec/models/project_feature_usage_spec.rb
+++ b/spec/models/project_feature_usage_spec.rb
@@ -126,4 +126,54 @@
       end
     end
   end
+
+  context 'ProjectFeatureUsage with DB Load Balancing', :request_store do
+    include_context 'clear DB Load Balancing configuration'
+
+    describe '#log_jira_dvcs_integration_usage' do
+      let!(:project) { create(:project) }
+
+      subject { project.feature_usage }
+
+      context 'database load balancing is configured' do
+        before do
+          # Do not pollute AR for other tests, but rather simulate effect of configure_proxy.
+          allow(ActiveRecord::Base.singleton_class).to receive(:prepend)
+          ::Gitlab::Database::LoadBalancing.configure_proxy
+          allow(ActiveRecord::Base).to receive(:connection).and_return(::Gitlab::Database::LoadBalancing.proxy)
+          ::Gitlab::Database::LoadBalancing::Session.clear_session
+        end
+
+        it 'logs Jira DVCS Cloud last sync' do
+          freeze_time do
+            subject.log_jira_dvcs_integration_usage
+
+            expect(subject.jira_dvcs_server_last_sync_at).to be_nil
+            expect(subject.jira_dvcs_cloud_last_sync_at).to be_like_time(Time.current)
+          end
+        end
+
+        it 'does not stick to primary' do
+          expect(::Gitlab::Database::LoadBalancing::Session.current).not_to be_performed_write
+          expect(::Gitlab::Database::LoadBalancing::Session.current).not_to be_using_primary
+
+          subject.log_jira_dvcs_integration_usage
+
+          expect(::Gitlab::Database::LoadBalancing::Session.current).to be_performed_write
+          expect(::Gitlab::Database::LoadBalancing::Session.current).not_to be_using_primary
+        end
+      end
+
+      context 'database load balancing is not cofigured' do
+        it 'logs Jira DVCS Cloud last sync' do
+          freeze_time do
+            subject.log_jira_dvcs_integration_usage
+
+            expect(subject.jira_dvcs_server_last_sync_at).to be_nil
+            expect(subject.jira_dvcs_cloud_last_sync_at).to be_like_time(Time.current)
+          end
+        end
+      end
+    end
+  end
 end
-- 
GitLab


From 6b9eec1e57d1c7b814ac2019ab615ed05a68889e Mon Sep 17 00:00:00 2001
From: Thong Kuah <tkuah@gitlab.com>
Date: Thu, 27 May 2021 12:48:39 +1200
Subject: [PATCH 06/10] Move caught-up check for Ci::RegisterJobService to Core

---
 app/services/ci/register_job_service.rb       | 18 +++++++++++-
 ee/app/services/ee/ci/register_job_service.rb | 18 ------------
 .../services/ci/register_job_service_spec.rb  | 28 -------------------
 spec/services/ci/register_job_service_spec.rb | 28 +++++++++++++++++++
 4 files changed, 45 insertions(+), 47 deletions(-)

diff --git a/app/services/ci/register_job_service.rb b/app/services/ci/register_job_service.rb
index 04956c4813655e..c97ab1ceaa8883 100644
--- a/app/services/ci/register_job_service.rb
+++ b/app/services/ci/register_job_service.rb
@@ -22,11 +22,27 @@ def initialize(runner)
     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
diff --git a/ee/app/services/ee/ci/register_job_service.rb b/ee/app/services/ee/ci/register_job_service.rb
index baf5ebd4a14033..cd1ffd1660c8f4 100644
--- a/ee/app/services/ee/ci/register_job_service.rb
+++ b/ee/app/services/ee/ci/register_job_service.rb
@@ -10,24 +10,6 @@ module RegisterJobService
       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)
diff --git a/ee/spec/services/ci/register_job_service_spec.rb b/ee/spec/services/ci/register_job_service_spec.rb
index 48d29b15065d11..caa6c3df64644b 100644
--- a/ee/spec/services/ci/register_job_service_spec.rb
+++ b/ee/spec/services/ci/register_job_service_spec.rb
@@ -10,34 +10,6 @@
   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 }
 
diff --git a/spec/services/ci/register_job_service_spec.rb b/spec/services/ci/register_job_service_spec.rb
index 067aed37468838..99b176a692cbe9 100644
--- a/spec/services/ci/register_job_service_spec.rb
+++ b/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
-- 
GitLab


From d9cac29093e6b7565bac93ef510e7d0d87bd3c86 Mon Sep 17 00:00:00 2001
From: Thong Kuah <tkuah@gitlab.com>
Date: Thu, 27 May 2021 13:41:44 +1200
Subject: [PATCH 07/10] Move non-sticky call in Users::ActivityService to Core

---
 app/services/users/activity_service.rb        |  4 +-
 ee/app/services/ee/users/activity_service.rb  | 16 ------
 .../services/users/activity_service_spec.rb   | 50 -------------------
 spec/services/users/activity_service_spec.rb  | 47 +++++++++++++++++
 4 files changed, 48 insertions(+), 69 deletions(-)
 delete mode 100644 ee/app/services/ee/users/activity_service.rb
 delete mode 100644 ee/spec/services/users/activity_service_spec.rb

diff --git a/app/services/users/activity_service.rb b/app/services/users/activity_service.rb
index c89a286cc8beac..20594bec28d52d 100644
--- a/app/services/users/activity_service.rb
+++ b/app/services/users/activity_service.rb
@@ -17,7 +17,7 @@ def initialize(author)
     def execute
       return unless @user
 
-      record_activity
+      ::Gitlab::Database::LoadBalancing::Session.without_sticky_writes { record_activity }
     end
 
     private
@@ -37,5 +37,3 @@ def record_activity
     end
   end
 end
-
-Users::ActivityService.prepend_mod
diff --git a/ee/app/services/ee/users/activity_service.rb b/ee/app/services/ee/users/activity_service.rb
deleted file mode 100644
index 77cc499cd6ac62..00000000000000
--- a/ee/app/services/ee/users/activity_service.rb
+++ /dev/null
@@ -1,16 +0,0 @@
-# frozen_string_literal: true
-
-module EE
-  module Users
-    module ActivityService
-      extend ::Gitlab::Utils::Override
-
-      private
-
-      override :record_activity
-      def record_activity
-        ::Gitlab::Database::LoadBalancing::Session.without_sticky_writes { super }
-      end
-    end
-  end
-end
diff --git a/ee/spec/services/users/activity_service_spec.rb b/ee/spec/services/users/activity_service_spec.rb
deleted file mode 100644
index f128b8afb6e320..00000000000000
--- a/ee/spec/services/users/activity_service_spec.rb
+++ /dev/null
@@ -1,50 +0,0 @@
-# frozen_string_literal: true
-
-require 'spec_helper'
-
-RSpec.describe Users::ActivityService, '#execute', :request_store, :redis, :clean_gitlab_redis_shared_state do
-  include_context 'clear DB Load Balancing configuration'
-
-  let(:user) { create(:user, last_activity_on: last_activity_on) }
-
-  context 'when last activity is in the past' do
-    let(:user) { create(:user, last_activity_on: Date.today - 1.week) }
-
-    context 'database load balancing is configured' do
-      before do
-        # Do not pollute AR for other tests, but rather simulate effect of configure_proxy.
-        allow(ActiveRecord::Base.singleton_class).to receive(:prepend)
-        ::Gitlab::Database::LoadBalancing.configure_proxy
-        allow(ActiveRecord::Base).to receive(:connection).and_return(::Gitlab::Database::LoadBalancing.proxy)
-      end
-
-      let(:service) do
-        service = described_class.new(user)
-
-        ::Gitlab::Database::LoadBalancing::Session.clear_session
-
-        service
-      end
-
-      it 'does not stick to primary' do
-        expect(::Gitlab::Database::LoadBalancing::Session.current).not_to be_performed_write
-
-        service.execute
-
-        expect(user.last_activity_on).to eq(Date.today)
-        expect(::Gitlab::Database::LoadBalancing::Session.current).to be_performed_write
-        expect(::Gitlab::Database::LoadBalancing::Session.current).not_to be_using_primary
-      end
-    end
-
-    context 'database load balancing is not configured' do
-      let(:service) { described_class.new(user) }
-
-      it 'updates user without error' do
-        service.execute
-
-        expect(user.last_activity_on).to eq(Date.today)
-      end
-    end
-  end
-end
diff --git a/spec/services/users/activity_service_spec.rb b/spec/services/users/activity_service_spec.rb
index 4bbf6a2bcb86fd..cfafa9eff454c0 100644
--- a/spec/services/users/activity_service_spec.rb
+++ b/spec/services/users/activity_service_spec.rb
@@ -84,4 +84,51 @@
       end
     end
   end
+
+  context 'with DB Load Balancing', :request_store, :redis, :clean_gitlab_redis_shared_state do
+    include_context 'clear DB Load Balancing configuration'
+
+    let(:user) { create(:user, last_activity_on: last_activity_on) }
+
+    context 'when last activity is in the past' do
+      let(:user) { create(:user, last_activity_on: Date.today - 1.week) }
+
+      context 'database load balancing is configured' do
+        before do
+          # Do not pollute AR for other tests, but rather simulate effect of configure_proxy.
+          allow(ActiveRecord::Base.singleton_class).to receive(:prepend)
+          ::Gitlab::Database::LoadBalancing.configure_proxy
+          allow(ActiveRecord::Base).to receive(:connection).and_return(::Gitlab::Database::LoadBalancing.proxy)
+        end
+
+        let(:service) do
+          service = described_class.new(user)
+
+          ::Gitlab::Database::LoadBalancing::Session.clear_session
+
+          service
+        end
+
+        it 'does not stick to primary' do
+          expect(::Gitlab::Database::LoadBalancing::Session.current).not_to be_performed_write
+
+          service.execute
+
+          expect(user.last_activity_on).to eq(Date.today)
+          expect(::Gitlab::Database::LoadBalancing::Session.current).to be_performed_write
+          expect(::Gitlab::Database::LoadBalancing::Session.current).not_to be_using_primary
+        end
+      end
+
+      context 'database load balancing is not configured' do
+        let(:service) { described_class.new(user) }
+
+        it 'updates user without error' do
+          service.execute
+
+          expect(user.last_activity_on).to eq(Date.today)
+        end
+      end
+    end
+  end
 end
-- 
GitLab


From a6c2b2cd2ea017d6cd3807b431fa6c032b43d774 Mon Sep 17 00:00:00 2001
From: Thong Kuah <tkuah@gitlab.com>
Date: Thu, 27 May 2021 21:14:20 +1200
Subject: [PATCH 08/10] Move LB calls for API helpers to Core

---
 ee/lib/ee/api/helpers.rb                      | 14 -----
 ee/lib/ee/api/helpers/runner.rb               | 24 --------
 ee/spec/lib/ee/api/helpers_spec.rb            | 38 ------------
 lib/api/helpers.rb                            |  5 ++
 lib/api/helpers/runner.rb                     | 10 ++++
 .../lib}/api/helpers/runner_spec.rb           |  2 +-
 spec/lib/api/helpers_spec.rb                  | 60 +++++++++++++++++++
 7 files changed, 76 insertions(+), 77 deletions(-)
 rename {ee/spec/lib/ee => spec/lib}/api/helpers/runner_spec.rb (97%)

diff --git a/ee/lib/ee/api/helpers.rb b/ee/lib/ee/api/helpers.rb
index 7d1df5e9f86883..4ad4ccb395701e 100644
--- a/ee/lib/ee/api/helpers.rb
+++ b/ee/lib/ee/api/helpers.rb
@@ -37,20 +37,6 @@ def check_gitlab_geo_request_ip!
         unauthorized! unless ::Gitlab::Geo.allowed_ip?(request.ip)
       end
 
-      override :current_user
-      def current_user
-        strong_memoize(:current_user) do
-          user = super
-
-          if user
-            ::Gitlab::Database::LoadBalancing::RackMiddleware
-              .stick_or_unstick(env, :user, user.id)
-          end
-
-          user
-        end
-      end
-
       def authorization_header_valid?
         return unless gitlab_geo_node_token?
 
diff --git a/ee/lib/ee/api/helpers/runner.rb b/ee/lib/ee/api/helpers/runner.rb
index d7c0ced22e2b9b..ecc3aac73dc65a 100644
--- a/ee/lib/ee/api/helpers/runner.rb
+++ b/ee/lib/ee/api/helpers/runner.rb
@@ -6,30 +6,6 @@ module Helpers
       module Runner
         extend ::Gitlab::Utils::Override
 
-        override :current_job
-        def current_job
-          id = params[:id]
-
-          if id
-            ::Gitlab::Database::LoadBalancing::RackMiddleware
-              .stick_or_unstick(env, :build, id)
-          end
-
-          super
-        end
-
-        override :current_runner
-        def current_runner
-          token = params[:token]
-
-          if token
-            ::Gitlab::Database::LoadBalancing::RackMiddleware
-              .stick_or_unstick(env, :runner, token)
-          end
-
-          super
-        end
-
         override :track_ci_minutes_usage!
         def track_ci_minutes_usage!(build, runner)
           ::Ci::Minutes::TrackLiveConsumptionService.new(build).execute
diff --git a/ee/spec/lib/ee/api/helpers_spec.rb b/ee/spec/lib/ee/api/helpers_spec.rb
index c53f61e9bd1424..5c84ae47094700 100644
--- a/ee/spec/lib/ee/api/helpers_spec.rb
+++ b/ee/spec/lib/ee/api/helpers_spec.rb
@@ -26,44 +26,6 @@ def app
     helper
   end
 
-  describe '#current_user' do
-    let(:user) { build(:user, id: 42) }
-
-    before do
-      allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
-    end
-
-    it 'handles sticking when a user could be found' do
-      allow_any_instance_of(API::Helpers).to receive(:initial_current_user).and_return(user)
-
-      expect(Gitlab::Database::LoadBalancing::RackMiddleware)
-        .to receive(:stick_or_unstick).with(any_args, :user, 42)
-
-      get 'user'
-
-      expect(Gitlab::Json.parse(last_response.body)).to eq({ 'id' => user.id })
-    end
-
-    it 'does not handle sticking if no user could be found' do
-      allow_any_instance_of(API::Helpers).to receive(:initial_current_user).and_return(nil)
-
-      expect(Gitlab::Database::LoadBalancing::RackMiddleware)
-        .not_to receive(:stick_or_unstick)
-
-      get 'user'
-
-      expect(Gitlab::Json.parse(last_response.body)).to eq({ 'found' => false })
-    end
-
-    it 'returns the user if one could be found' do
-      allow_any_instance_of(API::Helpers).to receive(:initial_current_user).and_return(user)
-
-      get 'user'
-
-      expect(Gitlab::Json.parse(last_response.body)).to eq({ 'id' => user.id })
-    end
-  end
-
   describe '#authenticate_by_gitlab_geo_node_token!' do
     let(:invalid_geo_auth_header) { "#{::Gitlab::Geo::BaseRequest::GITLAB_GEO_AUTH_TOKEN_TYPE}...Test" }
 
diff --git a/lib/api/helpers.rb b/lib/api/helpers.rb
index b719b046977526..2361e198603b6b 100644
--- a/lib/api/helpers.rb
+++ b/lib/api/helpers.rb
@@ -74,6 +74,11 @@ def current_user
 
       save_current_user_in_env(@current_user) if @current_user
 
+      if @current_user
+        ::Gitlab::Database::LoadBalancing::RackMiddleware
+          .stick_or_unstick(env, :user, @current_user.id)
+      end
+
       @current_user
     end
     # rubocop:enable Gitlab/ModuleWithInstanceVariables
diff --git a/lib/api/helpers/runner.rb b/lib/api/helpers/runner.rb
index e141abeef91ba4..65d66def148075 100644
--- a/lib/api/helpers/runner.rb
+++ b/lib/api/helpers/runner.rb
@@ -34,6 +34,11 @@ def get_runner_ip
       end
 
       def current_runner
+        if params[:token]
+          ::Gitlab::Database::LoadBalancing::RackMiddleware
+            .stick_or_unstick(env, :runner, params[:token])
+        end
+
         strong_memoize(:current_runner) do
           ::Ci::Runner.find_by_token(params[:token].to_s)
         end
@@ -65,6 +70,11 @@ def authenticate_job!(require_running: true)
       end
 
       def current_job
+        if params[:id]
+          ::Gitlab::Database::LoadBalancing::RackMiddleware
+            .stick_or_unstick(env, :build, params[:id])
+        end
+
         strong_memoize(:current_job) do
           ::Ci::Build.find_by_id(params[:id])
         end
diff --git a/ee/spec/lib/ee/api/helpers/runner_spec.rb b/spec/lib/api/helpers/runner_spec.rb
similarity index 97%
rename from ee/spec/lib/ee/api/helpers/runner_spec.rb
rename to spec/lib/api/helpers/runner_spec.rb
index a467b7c710fccf..e55c20b7ab62b5 100644
--- a/ee/spec/lib/ee/api/helpers/runner_spec.rb
+++ b/spec/lib/api/helpers/runner_spec.rb
@@ -2,7 +2,7 @@
 
 require 'spec_helper'
 
-RSpec.describe EE::API::Helpers::Runner do
+RSpec.describe API::Helpers::Runner do
   let(:helper) { Class.new { include API::Helpers::Runner }.new }
 
   before do
diff --git a/spec/lib/api/helpers_spec.rb b/spec/lib/api/helpers_spec.rb
index 87cd0d4388cdf4..6e48ee4c31596f 100644
--- a/spec/lib/api/helpers_spec.rb
+++ b/spec/lib/api/helpers_spec.rb
@@ -7,6 +7,66 @@
 
   subject { Class.new.include(described_class).new }
 
+  describe '#current_user' do
+    include Rack::Test::Methods
+
+    let(:user) { build(:user, id: 42) }
+
+    let(:helper) do
+      Class.new(Grape::API::Instance) do
+        helpers API::APIGuard::HelperMethods
+        helpers API::Helpers
+        format :json
+
+        get 'user' do
+          current_user ? { id: current_user.id } : { found: false }
+        end
+
+        get 'protected' do
+          authenticate_by_gitlab_geo_node_token!
+        end
+      end
+    end
+
+    def app
+      helper
+    end
+
+    before do
+      allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
+    end
+
+    it 'handles sticking when a user could be found' do
+      allow_any_instance_of(API::Helpers).to receive(:initial_current_user).and_return(user)
+
+      expect(Gitlab::Database::LoadBalancing::RackMiddleware)
+        .to receive(:stick_or_unstick).with(any_args, :user, 42)
+
+      get 'user'
+
+      expect(Gitlab::Json.parse(last_response.body)).to eq({ 'id' => user.id })
+    end
+
+    it 'does not handle sticking if no user could be found' do
+      allow_any_instance_of(API::Helpers).to receive(:initial_current_user).and_return(nil)
+
+      expect(Gitlab::Database::LoadBalancing::RackMiddleware)
+        .not_to receive(:stick_or_unstick)
+
+      get 'user'
+
+      expect(Gitlab::Json.parse(last_response.body)).to eq({ 'found' => false })
+    end
+
+    it 'returns the user if one could be found' do
+      allow_any_instance_of(API::Helpers).to receive(:initial_current_user).and_return(user)
+
+      get 'user'
+
+      expect(Gitlab::Json.parse(last_response.body)).to eq({ 'id' => user.id })
+    end
+  end
+
   describe '#find_project' do
     let(:project) { create(:project) }
 
-- 
GitLab


From 69c2b4b527457f67476a925cb8f97415cda767bc Mon Sep 17 00:00:00 2001
From: Thong Kuah <tkuah@gitlab.com>
Date: Thu, 27 May 2021 21:28:43 +1200
Subject: [PATCH 09/10] Move LB call in Auth to Core

---
 ee/lib/ee/gitlab/auth.rb        |  7 -------
 ee/spec/lib/gitlab/auth_spec.rb | 20 --------------------
 lib/gitlab/auth.rb              |  4 +++-
 spec/lib/gitlab/auth_spec.rb    | 22 ++++++++++++++++++++++
 4 files changed, 25 insertions(+), 28 deletions(-)

diff --git a/ee/lib/ee/gitlab/auth.rb b/ee/lib/ee/gitlab/auth.rb
index 06755033b10dbd..9e8a79e9057e4c 100644
--- a/ee/lib/ee/gitlab/auth.rb
+++ b/ee/lib/ee/gitlab/auth.rb
@@ -14,13 +14,6 @@ def find_with_user_password(login, password, increment_failed_attempts: false)
 
         super
       end
-
-      override :find_build_by_token
-      def find_build_by_token(token)
-        ::Gitlab::Database::LoadBalancing::Session.current.use_primary do
-          super
-        end
-      end
     end
   end
 end
diff --git a/ee/spec/lib/gitlab/auth_spec.rb b/ee/spec/lib/gitlab/auth_spec.rb
index 24bce22b87a0f3..427e3f02ac86d3 100644
--- a/ee/spec/lib/gitlab/auth_spec.rb
+++ b/ee/spec/lib/gitlab/auth_spec.rb
@@ -25,24 +25,4 @@
       expect( gl_auth.find_with_user_password(username, password) ).to eql user
     end
   end
-
-  describe '#build_access_token_check' do
-    subject { gl_auth.find_for_git_client('gitlab-ci-token', build.token, project: build.project, ip: '1.2.3.4') }
-
-    context 'for running build' do
-      let!(:build) { create(:ci_build, :running, user: user) }
-
-      it 'executes query using primary database' do
-        expect(Ci::Build).to receive(:find_by_token).with(build.token).and_wrap_original do |m, *args|
-          expect(::Gitlab::Database::LoadBalancing::Session.current.use_primary?).to eq(true)
-          m.call(*args)
-        end
-
-        expect(subject).to be_a(Gitlab::Auth::Result)
-        expect(subject.actor).to eq(user)
-        expect(subject.project).to eq(build.project)
-        expect(subject.type).to eq(:build)
-      end
-    end
-  end
 end
diff --git a/lib/gitlab/auth.rb b/lib/gitlab/auth.rb
index 4489fc9f3b2fa8..d2261148b213bf 100644
--- a/lib/gitlab/auth.rb
+++ b/lib/gitlab/auth.rb
@@ -371,7 +371,9 @@ def non_admin_available_scopes
       end
 
       def find_build_by_token(token)
-        ::Ci::AuthJobFinder.new(token: token).execute
+        ::Gitlab::Database::LoadBalancing::Session.current.use_primary do
+          ::Ci::AuthJobFinder.new(token: token).execute
+        end
       end
 
       def user_auth_attempt!(user, success:)
diff --git a/spec/lib/gitlab/auth_spec.rb b/spec/lib/gitlab/auth_spec.rb
index 7f06e66ad505b1..de61402b6d3f73 100644
--- a/spec/lib/gitlab/auth_spec.rb
+++ b/spec/lib/gitlab/auth_spec.rb
@@ -683,6 +683,28 @@ def operation
     end
   end
 
+  describe '#build_access_token_check' do
+    subject { gl_auth.find_for_git_client('gitlab-ci-token', build.token, project: build.project, ip: '1.2.3.4') }
+
+    let_it_be(:user) { create(:user) }
+
+    context 'for running build' do
+      let!(:build) { create(:ci_build, :running, user: user) }
+
+      it 'executes query using primary database' do
+        expect(Ci::Build).to receive(:find_by_token).with(build.token).and_wrap_original do |m, *args|
+          expect(::Gitlab::Database::LoadBalancing::Session.current.use_primary?).to eq(true)
+          m.call(*args)
+        end
+
+        expect(subject).to be_a(Gitlab::Auth::Result)
+        expect(subject.actor).to eq(user)
+        expect(subject.project).to eq(build.project)
+        expect(subject.type).to eq(:build)
+      end
+    end
+  end
+
   describe 'find_with_user_password' do
     let!(:user) do
       create(:user,
-- 
GitLab


From 52dc926bbe9926a1afae390fd8de3cbc3ef7f8ba Mon Sep 17 00:00:00 2001
From: Thong Kuah <tkuah@gitlab.com>
Date: Thu, 27 May 2021 21:41:18 +1200
Subject: [PATCH 10/10] Move LB call in Ci::Trace to Core

---
 ee/lib/ee/gitlab/ci/trace.rb | 35 -----------------------------------
 lib/gitlab/ci/trace.rb       | 23 ++++++++++++++++-------
 2 files changed, 16 insertions(+), 42 deletions(-)
 delete mode 100644 ee/lib/ee/gitlab/ci/trace.rb

diff --git a/ee/lib/ee/gitlab/ci/trace.rb b/ee/lib/ee/gitlab/ci/trace.rb
deleted file mode 100644
index 61b110ccd64cae..00000000000000
--- a/ee/lib/ee/gitlab/ci/trace.rb
+++ /dev/null
@@ -1,35 +0,0 @@
-# frozen_string_literal: true
-
-module EE
-  module Gitlab
-    module Ci
-      module Trace
-        extend ::Gitlab::Utils::Override
-
-        override :destroy_stream
-        def destroy_stream(build)
-          if consistent_archived_trace?(build)
-            ::Gitlab::Database::LoadBalancing::Sticking
-              .stick('ci/build/trace', build.id)
-          end
-
-          yield
-        end
-
-        override :read_trace_artifact
-        def read_trace_artifact(build)
-          if consistent_archived_trace?(build)
-            ::Gitlab::Database::LoadBalancing::Sticking
-              .unstick_or_continue_sticking('ci/build/trace', build.id)
-          end
-
-          yield
-        end
-
-        def consistent_archived_trace?(build)
-          ::Feature.enabled?(:gitlab_ci_archived_trace_consistent_reads, build.project, default_enabled: false)
-        end
-      end
-    end
-  end
-end
diff --git a/lib/gitlab/ci/trace.rb b/lib/gitlab/ci/trace.rb
index c4757edf74e034..19cc1ff91f7df0 100644
--- a/lib/gitlab/ci/trace.rb
+++ b/lib/gitlab/ci/trace.rb
@@ -296,25 +296,34 @@ def trace_artifact
         read_trace_artifact(job) { job.job_artifacts_trace }
       end
 
-      ##
-      # Overridden in EE
-      #
-      def destroy_stream(job)
+      def destroy_stream(build)
+        if consistent_archived_trace?(build)
+          ::Gitlab::Database::LoadBalancing::Sticking
+            .stick('ci/build/trace', build.id)
+        end
+
         yield
       end
 
       ##
       # Overriden in EE
       #
-      def read_trace_artifact(job)
+      def read_trace_artifact(build)
+        if consistent_archived_trace?(build)
+          ::Gitlab::Database::LoadBalancing::Sticking
+            .unstick_or_continue_sticking('ci/build/trace', build.id)
+        end
+
         yield
       end
 
+      def consistent_archived_trace?(build)
+        ::Feature.enabled?(:gitlab_ci_archived_trace_consistent_reads, build.project, default_enabled: false)
+      end
+
       def being_watched_cache_key
         "gitlab:ci:trace:#{job.id}:watched"
       end
     end
   end
 end
-
-::Gitlab::Ci::Trace.prepend_mod_with('Gitlab::Ci::Trace')
-- 
GitLab