Always enable the database load balancer

By always enabling the database load balancer we reduce the amount of
different code paths, make it easier for developers to interact with the
load balancing code (e.g. by adding metrics), and make it easier to add
support for multiple databases (e.g. by not having to check if the load
balancer is enabled or not).

As part of this commit we fixed various parts of the code that didn't
expect the load balancer to always be enabled (or at least when running
tests).

Changelog: added
Co-Authored-By: Dylan Griffith <dyl.griffith@gmail.com>

app/workers/authorized_project_update/user_refresh_from_replica_worker.rb

@@ -30,8 +30,6 @@ module AuthorizedProjectUpdate
     # does not allow us to deduplicate these jobs.
     # https://gitlab.com/gitlab-org/gitlab/-/issues/325291
     def use_replica_if_available(&block)
-      return yield unless ::Gitlab::Database::LoadBalancing.enable?
-
       ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries(&block)
     end
 

app/workers/container_expiration_policy_worker.rb

@@ -45,8 +45,6 @@ class ContainerExpirationPolicyWorker # rubocop:disable Scalability/IdempotentWo
   # not perfomed with a delay
   # https://gitlab.com/gitlab-org/gitlab/-/merge_requests/63635#note_603771207
   def use_replica_if_available(&blk)
-    return yield unless ::Gitlab::Database::LoadBalancing.enable?
-
     ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries(&blk)
   end
 

config/initializers/action_cable.rb

@@ -19,4 +19,4 @@ ActionCable::SubscriptionAdapter::Redis.redis_connector = lambda do |config|
 end
 
 Gitlab::ActionCable::RequestStoreCallbacks.install
-Gitlab::Database::LoadBalancing::ActionCableCallbacks.install if Gitlab::Database::LoadBalancing.enable?
+Gitlab::Database::LoadBalancing::ActionCableCallbacks.install

config/initializers/load_balancing.rb

@@ -2,27 +2,51 @@
 
 ActiveRecord::Base.singleton_class.attr_accessor :load_balancing_proxy
 
-if Gitlab::Database::LoadBalancing.enable?
-  Gitlab::Database.main.disable_prepared_statements
+Gitlab::Database.main.disable_prepared_statements
 
-  Gitlab::Application.configure do |config|
+Gitlab::Application.configure do |config|
   config.middleware.use(Gitlab::Database::LoadBalancing::RackMiddleware)
-  end
+end
+
+# This hijacks the "connection" method to ensure both
+# `ActiveRecord::Base.connection` and all models use the same load
+# balancing proxy.
+ActiveRecord::Base.singleton_class.prepend(Gitlab::Database::LoadBalancing::ActiveRecordProxy)
+
+# The load balancer needs to be configured immediately, and re-configured after
+# forking. This ensures queries that run before forking use the load balancer,
+# and queries running after a fork don't run into any errors when using dead
+# database connections.
+#
+# See https://gitlab.com/gitlab-org/gitlab/-/merge_requests/63485 for more
+# information.
+setup = proc do
+  lb = Gitlab::Database::LoadBalancing::LoadBalancer.new(
+    Gitlab::Database::LoadBalancing.configuration,
+    primary_only: !Gitlab::Database::LoadBalancing.enable_replicas?
+  )
 
-  # This hijacks the "connection" method to ensure both
-  # `ActiveRecord::Base.connection` and all models use the same load
-  # balancing proxy.
-  ActiveRecord::Base.singleton_class.prepend(Gitlab::Database::LoadBalancing::ActiveRecordProxy)
+  ActiveRecord::Base.load_balancing_proxy =
+    Gitlab::Database::LoadBalancing::ConnectionProxy.new(lb)
 
-  Gitlab::Database::LoadBalancing.configure_proxy
+  # Populate service discovery immediately if it is configured
+  Gitlab::Database::LoadBalancing.perform_service_discovery
+end
+
+setup.call
 
-  # This needs to be executed after fork of clustered processes
-  Gitlab::Cluster::LifecycleEvents.on_worker_start do
-    # For Host-based LB, we need to re-connect as Rails discards connections on fork
-    Gitlab::Database::LoadBalancing.configure_proxy
+# Database queries may be run before we fork, so we must set up the load
+# balancer as early as possible. When we do fork, we need to make sure all the
+# hosts are disconnected.
+Gitlab::Cluster::LifecycleEvents.on_before_fork do
+  # When forking, we don't want to wait until the connections aren't in use any
+  # more, as this could delay the boot cycle.
+  Gitlab::Database::LoadBalancing.proxy.load_balancer.disconnect!(timeout: 0)
+end
 
-    # Service discovery must be started after configuring the proxy, as service
-    # discovery depends on this.
+# Service discovery only needs to run in the worker processes, as the main one
+# won't be running many (if any) database queries.
+Gitlab::Cluster::LifecycleEvents.on_worker_start do
+  setup.call
   Gitlab::Database::LoadBalancing.start_service_discovery
-  end
 end

config/initializers/zz_metrics.rb

@@ -151,12 +151,8 @@ if Gitlab::Metrics.enabled? && !Rails.env.test? && !(Rails.env.development? && d
   Gitlab::Application.configure do |config|
     # We want to track certain metrics during the Load Balancing host resolving process.
     # Because of that, we need to have metrics code available earlier for Load Balancing.
-    if Gitlab::Database::LoadBalancing.enable?
     config.middleware.insert_before Gitlab::Database::LoadBalancing::RackMiddleware,
       Gitlab::Metrics::RackMiddleware
-    else
-      config.middleware.use(Gitlab::Metrics::RackMiddleware)
-    end
 
     config.middleware.use(Gitlab::Middleware::RailsQueueDuration)
     config.middleware.use(Gitlab::Metrics::ElasticsearchRackMiddleware)

ee/spec/lib/gitlab/geo_spec.rb

@@ -166,7 +166,7 @@ RSpec.describe Gitlab::Geo, :geo, :request_store do
   describe '.connected?' do
     context 'when there is a database issue' do
       it 'returns false when it cannot open an active database connection' do
-        allow(GeoNode.connection).to receive(:active?).and_return(false)
+        allow(GeoNode.retrieve_connection).to receive(:active?).and_return(false)
 
         expect(described_class.connected?).to be_falsey
       end

lib/gitlab/checks/matching_merge_request.rb

@@ -13,7 +13,6 @@ module Gitlab
       end
 
       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.
@@ -30,7 +29,6 @@ module Gitlab
         track_session_metrics do
           ::Gitlab::Database::LoadBalancing::Sticking.select_valid_host(:project, @project.id)
         end
-        end
 
         # rubocop: disable CodeReuse/ActiveRecord
         @project.merge_requests

lib/gitlab/database/load_balancing.rb

@@ -19,19 +19,8 @@ module Gitlab
                             [].freeze
                           end
 
-      ProxyNotConfiguredError = Class.new(StandardError)
-
-      # The connection proxy to use for load balancing (if enabled).
       def self.proxy
-        unless load_balancing_proxy = ActiveRecord::Base.load_balancing_proxy
-          Gitlab::ErrorTracking.track_exception(
-            ProxyNotConfiguredError.new(
-              "Attempting to access the database load balancing proxy, but it wasn't configured.\n" \
-              "Did you forget to call '#{self.name}.configure_proxy'?"
-            ))
-        end
-
-        load_balancing_proxy
+        ActiveRecord::Base.load_balancing_proxy
       end
 
       # Returns a Hash containing the load balancing configuration.
@@ -39,8 +28,11 @@ module Gitlab
         @configuration ||= Configuration.for_model(ActiveRecord::Base)
       end
 
-      # Returns true if load balancing is to be enabled.
-      def self.enable?
+      # Returns `true` if the use of load balancing replicas should be enabled.
+      #
+      # This is disabled for Rake tasks to ensure e.g. database migrations
+      # always produce consistent results.
+      def self.enable_replicas?
         return false if Gitlab::Runtime.rake?
 
         configured?
@@ -59,18 +51,13 @@ module Gitlab
           .start
       end
 
-      # Configures proxying of requests.
-      def self.configure_proxy
-        lb = LoadBalancer.new(configuration, primary_only: !enable?)
-        ActiveRecord::Base.load_balancing_proxy = ConnectionProxy.new(lb)
+      def self.perform_service_discovery
+        return unless configuration.service_discovery_enabled?
 
-        # Populate service discovery immediately if it is configured
-        if configuration.service_discovery_enabled?
         ServiceDiscovery
-            .new(lb, **configuration.service_discovery)
+          .new(proxy.load_balancer, **configuration.service_discovery)
           .perform_service_discovery
       end
-      end
 
       DB_ROLES = [
         ROLE_PRIMARY = :primary,

lib/gitlab/database/load_balancing/rack_middleware.rb

@@ -18,8 +18,6 @@ module Gitlab
         # namespace - The namespace to use for sticking.
         # id - The identifier to use for sticking.
         def self.stick_or_unstick(env, namespace, id)
-          return unless ::Gitlab::Database::LoadBalancing.enable?
-
           ::Gitlab::Database::LoadBalancing::Sticking.unstick_or_continue_sticking(namespace, id)
 
           env[STICK_OBJECT] ||= Set.new

lib/gitlab/database/load_balancing/sticking.rb

@@ -22,8 +22,6 @@ module Gitlab
 
         # Sticks to the primary if a write was performed.
         def self.stick_if_necessary(namespace, id)
-          return unless LoadBalancing.enable?
-
           stick(namespace, id) if Session.current.performed_write?
         end
 
@@ -75,15 +73,11 @@ module Gitlab
         # Starts sticking to the primary for the given namespace and id, using
         # the latest WAL pointer from the primary.
         def self.stick(namespace, id)
-          return unless LoadBalancing.enable?
-
           mark_primary_write_location(namespace, id)
           Session.current.use_primary!
         end
 
         def self.bulk_stick(namespace, ids)
-          return unless LoadBalancing.enable?
-
           with_primary_write_location do |location|
             ids.each do |id|
               set_write_location_for(namespace, id, location)
@@ -94,17 +88,7 @@ module Gitlab
         end
 
         def self.with_primary_write_location
-          return unless LoadBalancing.configured?
-
-          # Load balancing could be enabled for the Web application server,
-          # but it's not activated for Sidekiq. We should update Redis with
-          # the write location just in case load balancing is being used.
-          location =
-            if LoadBalancing.enable?
-              load_balancer.primary_write_location
-            else
-              Gitlab::Database.main.get_write_location(ActiveRecord::Base.connection)
-            end
+          location = load_balancer.primary_write_location
 
           return if location.blank?
 

lib/gitlab/import_export/json/streaming_serializer.rb

@@ -171,7 +171,6 @@ module Gitlab
 
         def read_from_replica_if_available(&block)
           return yield unless ::Feature.enabled?(:load_balancing_for_export_workers, type: :development, default_enabled: :yaml)
-          return yield unless ::Gitlab::Database::LoadBalancing.enable?
 
           ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries(&block)
         end

lib/gitlab/metrics/subscribers/active_record.rb

@@ -47,14 +47,12 @@ module Gitlab
             buckets SQL_DURATION_BUCKET
           end
 
-          if ::Gitlab::Database::LoadBalancing.enable?
           db_role = ::Gitlab::Database::LoadBalancing.db_role_for_connection(payload[:connection])
           return if db_role.blank?
 
           increment_db_role_counters(db_role, payload)
           observe_db_role_duration(db_role, event)
         end
-        end
 
         def self.db_counter_payload
           return {} unless Gitlab::SafeRequestStore.active?
@@ -64,7 +62,7 @@ module Gitlab
               payload[key] = Gitlab::SafeRequestStore[key].to_i
             end
 
-            if ::Gitlab::SafeRequestStore.active? && ::Gitlab::Database::LoadBalancing.enable?
+            if ::Gitlab::SafeRequestStore.active?
               load_balancing_metric_counter_keys.each do |counter|
                 payload[counter] = ::Gitlab::SafeRequestStore[counter].to_i
               end

lib/gitlab/metrics/subscribers/load_balancing.rb

@@ -10,7 +10,7 @@ module Gitlab
         LOG_COUNTERS = { true => :caught_up_replica_pick_ok, false => :caught_up_replica_pick_fail }.freeze
 
         def caught_up_replica_pick(event)
-          return unless Gitlab::SafeRequestStore.active? && ::Gitlab::Database::LoadBalancing.enable?
+          return unless Gitlab::SafeRequestStore.active?
 
           result = event.payload[:result]
           counter_name = counter(result)
@@ -20,13 +20,13 @@ module Gitlab
 
         # we want to update Prometheus counter after the controller/action are set
         def web_transaction_completed(_event)
-          return unless Gitlab::SafeRequestStore.active? && ::Gitlab::Database::LoadBalancing.enable?
+          return unless Gitlab::SafeRequestStore.active?
 
           LOG_COUNTERS.keys.each { |result| increment_prometheus_for_result_label(result) }
         end
 
         def self.load_balancing_payload
-          return {} unless Gitlab::SafeRequestStore.active? && ::Gitlab::Database::LoadBalancing.enable?
+          return {} unless Gitlab::SafeRequestStore.active?
 
           {}.tap do |payload|
             LOG_COUNTERS.values.each do |counter|

lib/gitlab/sidekiq_middleware.rb

@@ -39,7 +39,7 @@ module Gitlab
         # DuplicateJobs::Server should be placed  at the bottom, but before the SidekiqServerMiddleware,
         # so we can compare the latest WAL location against replica
         chain.add ::Gitlab::SidekiqMiddleware::DuplicateJobs::Server
-        chain.add ::Gitlab::Database::LoadBalancing::SidekiqServerMiddleware if load_balancing_enabled?
+        chain.add ::Gitlab::Database::LoadBalancing::SidekiqServerMiddleware
       end
     end
 
@@ -52,7 +52,7 @@ module Gitlab
         chain.add ::Labkit::Middleware::Sidekiq::Client
         # Sidekiq Client Middleware should be placed before DuplicateJobs::Client middleware,
         # so we can store WAL location before we deduplicate the job.
-        chain.add ::Gitlab::Database::LoadBalancing::SidekiqClientMiddleware if load_balancing_enabled?
+        chain.add ::Gitlab::Database::LoadBalancing::SidekiqClientMiddleware
         chain.add ::Gitlab::SidekiqMiddleware::DuplicateJobs::Client
         chain.add ::Gitlab::SidekiqStatus::ClientMiddleware
         chain.add ::Gitlab::SidekiqMiddleware::AdminMode::Client
@@ -61,10 +61,5 @@ module Gitlab
         chain.add ::Gitlab::SidekiqMiddleware::ClientMetrics
       end
     end
-
-    def self.load_balancing_enabled?
-      ::Gitlab::Database::LoadBalancing.enable?
-    end
-    private_class_method :load_balancing_enabled?
   end
 end

lib/gitlab/sidekiq_middleware/server_metrics.rb

@@ -53,11 +53,8 @@ module Gitlab
 
       def initialize
         @metrics = self.class.metrics
-
-        if ::Gitlab::Database::LoadBalancing.enable?
         @metrics[:sidekiq_load_balancing_count] = ::Gitlab::Metrics.counter(:sidekiq_load_balancing_count, 'Sidekiq jobs with load balancing')
       end
-      end
 
       def call(worker, job, queue)
         # This gives all the sidekiq worker threads a name, so we can recognize them
@@ -128,8 +125,6 @@ module Gitlab
       private
 
       def with_load_balancing_settings(job)
-        return unless ::Gitlab::Database::LoadBalancing.enable?
-
         keys = %w[load_balancing_strategy worker_data_consistency]
         return unless keys.all? { |k| job.key?(k) }
 

lib/peek/views/active_record.rb

@@ -44,11 +44,9 @@ module Peek
           count[item[:transaction]] += 1
         end
 
-        if ::Gitlab::Database::LoadBalancing.enable?
         count[item[:db_role]] ||= 0
         count[item[:db_role]] += 1
       end
-      end
 
       def setup_subscribers
         super
@@ -72,8 +70,6 @@ module Peek
       end
 
       def db_role(data)
-        return unless ::Gitlab::Database::LoadBalancing.enable?
-
         role = ::Gitlab::Database::LoadBalancing.db_role_for_connection(data[:connection]) ||
           ::Gitlab::Database::LoadBalancing::ROLE_UNKNOWN
 

spec/initializers/lograge_spec.rb

@@ -230,11 +230,6 @@ RSpec.describe 'lograge', type: :request do
         end
       end
 
-      context 'when load balancing is enabled' do
-        before do
-          allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
-        end
-
       context 'with db payload' do
         context 'when RequestStore is enabled', :request_store do
           it 'includes db counters for load balancing' do
@@ -253,18 +248,5 @@ RSpec.describe 'lograge', type: :request do
         end
       end
     end
-
-      context 'when load balancing is disabled' do
-        before do
-          allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false)
-        end
-
-        it 'does not include db counters for load balancing' do
-          subscriber.process_action(event)
-
-          expect(log_data).not_to include(*db_load_balancing_logging_keys)
-        end
-      end
-    end
   end
 end

spec/lib/api/helpers_spec.rb

@@ -32,10 +32,6 @@ RSpec.describe API::Helpers do
       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)
 

spec/lib/gitlab/checks/matching_merge_request_spec.rb

@@ -31,35 +31,22 @@ 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', :db_load_balancing do
+    context 'with load balancing enabled' do
       let(:session) { ::Gitlab::Database::LoadBalancing::Session.current }
       let(:all_caught_up) { true }
 
       before do
+        Gitlab::Database::LoadBalancing::Session.clear_session
         allow(::Gitlab::Database::LoadBalancing::Sticking).to receive(:all_caught_up?).and_return(all_caught_up)
 
         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
 
+      after do
+        Gitlab::Database::LoadBalancing::Session.clear_session
+      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

spec/lib/gitlab/database/bulk_update_spec.rb

@@ -122,7 +122,7 @@ RSpec.describe Gitlab::Database::BulkUpdate do
 
       stub_const('ActiveRecordBasePreparedStatementsInverted', klass)
 
-      c = ActiveRecord::Base.connection.instance_variable_get(:@config)
+      c = ActiveRecord::Base.retrieve_connection.instance_variable_get(:@config)
       inverted = c.merge(prepared_statements: !ActiveRecord::Base.connection.prepared_statements)
       ActiveRecordBasePreparedStatementsInverted.establish_connection(inverted)
 

spec/lib/gitlab/database/consistency_spec.rb

@@ -7,7 +7,7 @@ RSpec.describe Gitlab::Database::Consistency do
     Gitlab::Database::LoadBalancing::Session.current
   end
 
-  describe '.with_read_consistency' do
+  describe '.with_read_consistency', :db_load_balancing do
     it 'sticks to primary database' do
       expect(session).not_to be_using_primary
 

spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb

@@ -20,11 +20,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
   end
 
   describe '.stick_or_unstick' do
-    before do
-      allow(Gitlab::Database::LoadBalancing).to receive(:enable?)
-        .and_return(true)
-    end
-
     it 'sticks or unsticks a single object and updates the Rack environment' do
       expect(Gitlab::Database::LoadBalancing::Sticking)
         .to receive(:unstick_or_continue_sticking)

spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb

@@ -5,14 +5,13 @@ require 'spec_helper'
 RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
   let(:middleware) { described_class.new }
 
-  let(:load_balancer) { double.as_null_object }
+  let(:load_balancer) { Gitlab::Database::LoadBalancing.proxy.load_balancer }
   let(:worker_class) { 'TestDataConsistencyWorker' }
   let(:job) { { "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e" } }
 
   before do
     skip_feature_flags_yaml_validation
     skip_default_enabled_yaml_check
-    allow(::Gitlab::Database::LoadBalancing).to receive_message_chain(:proxy, :load_balancer).and_return(load_balancer)
   end
 
   after do

spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb

@@ -5,7 +5,7 @@ require 'spec_helper'
 RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware do
   let(:middleware) { described_class.new }
 
-  let(:load_balancer) { double.as_null_object }
+  let(:load_balancer) { Gitlab::Database::LoadBalancing.proxy.load_balancer }
 
   let(:worker) { worker_class.new }
   let(:job) { { "retry" => 3, "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", 'database_replica_location' => '0/D525E3A8' } }
@@ -13,7 +13,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware do
   before do
     skip_feature_flags_yaml_validation
     skip_default_enabled_yaml_check
-    allow(::Gitlab::Database::LoadBalancing).to receive_message_chain(:proxy, :load_balancer).and_return(load_balancer)
 
     replication_lag!(false)
   end

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

@@ -8,20 +8,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
   end
 
   describe '.stick_if_necessary' do
-    context 'when sticking is disabled' do
-      it 'does not perform any sticking' do
-        expect(described_class).not_to receive(:stick)
-
-        described_class.stick_if_necessary(:user, 42)
-      end
-    end
-
-    context 'when sticking is enabled' do
-      before do
-        allow(Gitlab::Database::LoadBalancing).to receive(:enable?)
-          .and_return(true)
-      end
-
     it 'does not stick if no write was performed' do
       allow(Gitlab::Database::LoadBalancing::Session.current)
         .to receive(:performed_write?)
@@ -42,7 +28,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
       described_class.stick_if_necessary(:user, 42)
     end
   end
-  end
 
   describe '.all_caught_up?' do
     let(:lb) { double(:lb) }
@@ -155,19 +140,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
   end
 
   RSpec.shared_examples 'sticking' do
-    context 'when sticking is disabled' do
-      it 'does not perform any sticking', :aggregate_failures do
-        expect(described_class).not_to receive(:set_write_location_for)
-        expect(Gitlab::Database::LoadBalancing::Session.current).not_to receive(:use_primary!)
-
-        described_class.bulk_stick(:user, ids)
-      end
-    end
-
-    context 'when sticking is enabled' do
     before do
-        allow(Gitlab::Database::LoadBalancing).to receive(:configured?).and_return(true)
-
       lb = double(:lb, primary_write_location: 'foo')
 
       allow(described_class).to receive(:load_balancer).and_return(lb)
@@ -185,7 +158,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
       subject
     end
   end
-  end
 
   describe '.stick' do
     it_behaves_like 'sticking' do
@@ -202,12 +174,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
   end
 
   describe '.mark_primary_write_location' do
-    context 'when enabled' do
-      before do
-        allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
-        allow(Gitlab::Database::LoadBalancing).to receive(:configured?).and_return(true)
-      end
-
     it 'updates the write location with the load balancer' do
       lb = double(:lb, primary_write_location: 'foo')
 
@@ -220,48 +186,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
     end
   end
 
-    context 'when load balancing is configured but not enabled' do
-      before do
-        allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false)
-        allow(Gitlab::Database::LoadBalancing).to receive(:configured?).and_return(true)
-      end
-
-      it 'updates the write location with the main ActiveRecord connection' do
-        allow(described_class).to receive(:load_balancer).and_return(nil)
-        expect(ActiveRecord::Base).to receive(:connection).and_call_original
-        expect(described_class).to receive(:set_write_location_for)
-          .with(:user, 42, anything)
-
-        described_class.mark_primary_write_location(:user, 42)
-      end
-
-      context 'when write location is nil' do
-        before do
-          allow(Gitlab::Database.main).to receive(:get_write_location).and_return(nil)
-        end
-
-        it 'does not update the write location' do
-          expect(described_class).not_to receive(:set_write_location_for)
-
-          described_class.mark_primary_write_location(:user, 42)
-        end
-      end
-    end
-
-    context 'when load balancing is disabled' do
-      before do
-        allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false)
-        allow(Gitlab::Database::LoadBalancing).to receive(:configured?).and_return(false)
-      end
-
-      it 'updates the write location with the main ActiveRecord connection' do
-        expect(described_class).not_to receive(:set_write_location_for)
-
-        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')

spec/lib/gitlab/database/load_balancing_spec.rb

@@ -4,38 +4,14 @@ require 'spec_helper'
 
 RSpec.describe Gitlab::Database::LoadBalancing do
   describe '.proxy' do
-    before do
-      @previous_proxy = ActiveRecord::Base.load_balancing_proxy
-
-      ActiveRecord::Base.load_balancing_proxy = connection_proxy
-    end
-
-    after do
-      ActiveRecord::Base.load_balancing_proxy = @previous_proxy
-    end
-
-    context 'when configured' do
-      let(:connection_proxy) { double(:connection_proxy) }
-
     it 'returns the connection proxy' do
-        expect(subject.proxy).to eq(connection_proxy)
-      end
-    end
+      proxy = double(:connection_proxy)
 
-    context 'when not configured' do
-      let(:connection_proxy) { nil }
-
-      it 'returns nil' do
-        expect(subject.proxy).to be_nil
-      end
-
-      it 'tracks an error to sentry' do
-        expect(Gitlab::ErrorTracking).to receive(:track_exception).with(
-          an_instance_of(subject::ProxyNotConfiguredError)
-        )
+      allow(ActiveRecord::Base)
+        .to receive(:load_balancing_proxy)
+        .and_return(proxy)
 
-        subject.proxy
-      end
+      expect(described_class.proxy).to eq(proxy)
     end
   end
 
@@ -50,38 +26,33 @@ RSpec.describe Gitlab::Database::LoadBalancing do
     end
   end
 
-  describe '.enable?' do
+  describe '.enable_replicas?' do
+    context 'when hosts are specified' do
       before do
         allow(described_class.configuration)
           .to receive(:hosts)
           .and_return(%w(foo))
       end
 
-    it 'returns false when no hosts are specified' do
-      allow(described_class.configuration).to receive(:hosts).and_return([])
-
-      expect(described_class.enable?).to eq(false)
+      it 'returns true' do
+        expect(described_class.enable_replicas?).to eq(true)
       end
 
       it 'returns true when Sidekiq is being used' do
         allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(true)
 
-      expect(described_class.enable?).to eq(true)
+        expect(described_class.enable_replicas?).to eq(true)
       end
 
       it 'returns false when running inside a Rake task' do
         allow(Gitlab::Runtime).to receive(:rake?).and_return(true)
 
-      expect(described_class.enable?).to eq(false)
+        expect(described_class.enable_replicas?).to eq(false)
       end
-
-    it 'returns true when load balancing should be enabled' do
-      allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(false)
-
-      expect(described_class.enable?).to eq(true)
     end
 
-    it 'returns true when service discovery is enabled' do
+    context 'when no hosts are specified but service discovery is enabled' do
+      it 'returns true' do
         allow(described_class.configuration).to receive(:hosts).and_return([])
         allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(false)
 
@@ -89,7 +60,19 @@ RSpec.describe Gitlab::Database::LoadBalancing do
           .to receive(:service_discovery_enabled?)
           .and_return(true)
 
-      expect(described_class.enable?).to eq(true)
+        expect(described_class.enable_replicas?).to eq(true)
+      end
+    end
+
+    context 'when no hosts are specified and service discovery is disabled' do
+      it 'returns false' do
+        allow(described_class.configuration).to receive(:hosts).and_return([])
+        allow(described_class.configuration)
+          .to receive(:service_discovery_enabled?)
+          .and_return(false)
+
+        expect(described_class.enable_replicas?).to eq(false)
+      end
     end
   end
 
@@ -121,41 +104,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
     end
   end
 
-  describe '.configure_proxy' do
-    before do
-      allow(ActiveRecord::Base).to receive(:load_balancing_proxy=)
-    end
-
-    it 'configures the connection proxy' do
-      described_class.configure_proxy
-
-      expect(ActiveRecord::Base).to have_received(:load_balancing_proxy=)
-        .with(Gitlab::Database::LoadBalancing::ConnectionProxy)
-    end
-
-    context 'when service discovery is enabled' do
-      it 'runs initial service discovery when configuring the connection proxy' do
-        discover = instance_spy(Gitlab::Database::LoadBalancing::ServiceDiscovery)
-
-        allow(described_class.configuration)
-          .to receive(:service_discovery)
-          .and_return({ record: 'foo' })
-
-        expect(Gitlab::Database::LoadBalancing::ServiceDiscovery)
-          .to receive(:new)
-          .with(
-            an_instance_of(Gitlab::Database::LoadBalancing::LoadBalancer),
-            an_instance_of(Hash)
-          )
-          .and_return(discover)
-
-        expect(discover).to receive(:perform_service_discovery)
-
-        described_class.configure_proxy
-      end
-    end
-  end
-
   describe '.start_service_discovery' do
     it 'does not start if service discovery is disabled' do
       expect(Gitlab::Database::LoadBalancing::ServiceDiscovery)
@@ -191,15 +139,41 @@ RSpec.describe Gitlab::Database::LoadBalancing do
     end
   end
 
-  describe '.db_role_for_connection' do
-    context 'when the load balancing is not configured' do
-      let(:connection) { ActiveRecord::Base.connection }
+  describe '.perform_service_discovery' do
+    it 'does nothing if service discovery is disabled' do
+      expect(Gitlab::Database::LoadBalancing::ServiceDiscovery)
+        .not_to receive(:new)
 
-      it 'returns primary' do
-        expect(described_class.db_role_for_connection(connection)).to eq(:primary)
+      described_class.perform_service_discovery
+    end
+
+    it 'performs service discovery when enabled' do
+      allow(described_class.configuration)
+        .to receive(:service_discovery_enabled?)
+        .and_return(true)
+
+      sv = instance_spy(Gitlab::Database::LoadBalancing::ServiceDiscovery)
+      cfg = Gitlab::Database::LoadBalancing::Configuration
+        .new(ActiveRecord::Base)
+      lb = Gitlab::Database::LoadBalancing::LoadBalancer.new(cfg)
+      proxy = Gitlab::Database::LoadBalancing::ConnectionProxy.new(lb)
+
+      allow(described_class)
+        .to receive(:proxy)
+        .and_return(proxy)
+
+      expect(Gitlab::Database::LoadBalancing::ServiceDiscovery)
+        .to receive(:new)
+        .with(lb, cfg.service_discovery)
+        .and_return(sv)
+
+      expect(sv).to receive(:perform_service_discovery)
+
+      described_class.perform_service_discovery
     end
   end
 
+  describe '.db_role_for_connection' do
     context 'when the NullPool is used for connection' do
       let(:pool) { ActiveRecord::ConnectionAdapters::NullPool.new }
       let(:connection) { double(:connection, pool: pool) }
@@ -274,10 +248,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
       end
     end
 
-    before do
-      model.singleton_class.prepend ::Gitlab::Database::LoadBalancing::ActiveRecordProxy
-    end
-
     where(:queries, :include_transaction, :expected_results) do
       [
         # Read methods

spec/lib/gitlab/database/with_lock_retries_spec.rb

@@ -7,7 +7,7 @@ RSpec.describe Gitlab::Database::WithLockRetries do
   let(:logger) { Gitlab::Database::WithLockRetries::NULL_LOGGER }
   let(:subject) { described_class.new(env: env, logger: logger, allow_savepoints: allow_savepoints, timing_configuration: timing_configuration) }
   let(:allow_savepoints) { true }
-  let(:connection) { ActiveRecord::Base.connection }
+  let(:connection) { ActiveRecord::Base.retrieve_connection }
 
   let(:timing_configuration) do
     [

spec/lib/gitlab/database_spec.rb

@@ -185,10 +185,23 @@ RSpec.describe Gitlab::Database do
 
   describe '.db_config_name' do
     it 'returns the db_config name for the connection' do
-      connection = ActiveRecord::Base.connection
+      model = ActiveRecord::Base
 
-      expect(described_class.db_config_name(connection)).to be_a(String)
-      expect(described_class.db_config_name(connection)).to eq(connection.pool.db_config.name)
+      # This is a ConnectionProxy
+      expect(described_class.db_config_name(model.connection))
+        .to eq('unknown')
+
+      # This is an actual connection
+      expect(described_class.db_config_name(model.retrieve_connection))
+        .to eq('main')
+    end
+
+    context 'when replicas are configured', :db_load_balancing do
+      it 'returns the name for a replica' do
+        replica = ActiveRecord::Base.connection.load_balancer.host
+
+        expect(described_class.db_config_name(replica)).to eq('main_replica')
+      end
     end
   end
 
@@ -279,7 +292,7 @@ RSpec.describe Gitlab::Database do
         expect(event).not_to be_nil
         expect(event.duration).to be > 0.0
         expect(event.payload).to a_hash_including(
-          connection: be_a(ActiveRecord::ConnectionAdapters::AbstractAdapter)
+          connection: be_a(Gitlab::Database::LoadBalancing::ConnectionProxy)
         )
       end
     end
@@ -296,7 +309,7 @@ RSpec.describe Gitlab::Database do
         expect(event).not_to be_nil
         expect(event.duration).to be > 0.0
         expect(event.payload).to a_hash_including(
-          connection: be_a(ActiveRecord::ConnectionAdapters::AbstractAdapter)
+          connection: be_a(Gitlab::Database::LoadBalancing::ConnectionProxy)
         )
       end
     end
@@ -319,7 +332,7 @@ RSpec.describe Gitlab::Database do
           expect(event).not_to be_nil
           expect(event.duration).to be > 0.0
           expect(event.payload).to a_hash_including(
-            connection: be_a(ActiveRecord::ConnectionAdapters::AbstractAdapter)
+            connection: be_a(Gitlab::Database::LoadBalancing::ConnectionProxy)
           )
         end
       end
@@ -340,7 +353,7 @@ RSpec.describe Gitlab::Database do
         expect(event).not_to be_nil
         expect(event.duration).to be > 0.0
         expect(event.payload).to a_hash_including(
-          connection: be_a(ActiveRecord::ConnectionAdapters::AbstractAdapter)
+          connection: be_a(Gitlab::Database::LoadBalancing::ConnectionProxy)
         )
       end
     end

spec/lib/gitlab/import_export/json/streaming_serializer_spec.rb

@@ -158,12 +158,11 @@ RSpec.describe Gitlab::ImportExport::Json::StreamingSerializer do
     end
 
     describe 'load balancing' do
-      context 'when feature flag load_balancing_for_export_workers is enabled' do
+      context 'when feature flag load_balancing_for_export_workers is enabled', :db_load_balancing do
         before do
           stub_feature_flags(load_balancing_for_export_workers: true)
         end
 
-        context 'when enabled', :db_load_balancing do
         it 'reads from replica' do
           expect(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_replicas_for_read_queries).and_call_original
 
@@ -171,16 +170,6 @@ RSpec.describe Gitlab::ImportExport::Json::StreamingSerializer do
         end
       end
 
-        context 'when disabled' do
-          it 'reads from primary' do
-            allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false)
-            expect(Gitlab::Database::LoadBalancing::Session.current).not_to receive(:use_replicas_for_read_queries)
-
-            subject.execute
-          end
-        end
-      end
-
       context 'when feature flag load_balancing_for_export_workers is disabled' do
         it 'reads from primary' do
           stub_feature_flags(load_balancing_for_export_workers: false)

spec/lib/gitlab/instrumentation_helper_spec.rb

@@ -107,11 +107,6 @@ RSpec.describe Gitlab::InstrumentationHelper do
       end
     end
 
-    context 'when load balancing is enabled' do
-      before do
-        allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
-      end
-
     it 'includes DB counts' do
       subject
 
@@ -154,26 +149,6 @@ RSpec.describe Gitlab::InstrumentationHelper do
     end
   end
 
-    context 'when load balancing is disabled' do
-      before do
-        allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false)
-      end
-
-      it 'does not include DB counts' do
-        subject
-
-        expect(payload).not_to include(db_replica_count: 0,
-                                   db_replica_cached_count: 0,
-                                   db_primary_count: 0,
-                                   db_primary_cached_count: 0,
-                                   db_primary_wal_count: 0,
-                                   db_replica_wal_count: 0,
-                                   db_primary_wal_cached_count: 0,
-                                   db_replica_wal_cached_count: 0)
-      end
-    end
-  end
-
   describe '.queue_duration_for_job' do
     where(:enqueued_at, :created_at, :time_now, :expected_duration) do
       "2019-06-01T00:00:00.000+0000" | nil                            | "2019-06-01T02:00:00.000+0000" | 2.hours.to_f

spec/lib/gitlab/metrics/subscribers/active_record_spec.rb

@@ -7,7 +7,7 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
 
   let(:env) { {} }
   let(:subscriber) { described_class.new }
-  let(:connection) { ActiveRecord::Base.connection }
+  let(:connection) { ActiveRecord::Base.retrieve_connection }
   let(:db_config_name) { ::Gitlab::Database.db_config_name(connection) }
 
   describe '#transaction' do
@@ -135,7 +135,7 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
         end
 
         it_behaves_like 'record ActiveRecord metrics'
-        it_behaves_like 'store ActiveRecord info in RequestStore'
+        it_behaves_like 'store ActiveRecord info in RequestStore', :primary
       end
     end
 
@@ -195,11 +195,7 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
       with_them do
         let(:payload) { { name: name, sql: sql(sql_query, comments: comments), connection: connection } }
 
-        before do
-          allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
-        end
-
-        context 'query using a connection to a replica' do
+        context 'query using a connection to a replica', :db_load_balancing do
           before do
             allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).and_return(:replica)
           end

spec/lib/gitlab/metrics/subscribers/load_balancing_spec.rb

@@ -5,10 +5,6 @@ require 'spec_helper'
 RSpec.describe Gitlab::Metrics::Subscribers::LoadBalancing, :request_store do
   let(:subscriber) { described_class.new }
 
-  before do
-    allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
-  end
-
   describe '#caught_up_replica_pick' do
     shared_examples 'having payload result value' do |result, counter_name|
       subject { subscriber.caught_up_replica_pick(event) }

spec/lib/gitlab/sidekiq_logging/structured_logger_spec.rb

@@ -243,8 +243,22 @@ RSpec.describe Gitlab::SidekiqLogging::StructuredLogger do
         expected_end_payload.merge(
           'db_duration_s' => a_value >= 0.1,
           'db_count' => a_value >= 1,
-          'db_cached_count' => 0,
-          'db_write_count' => 0
+          "db_replica_#{db_config_name}_count" => 0,
+          'db_replica_duration_s' => a_value >= 0,
+          'db_primary_count' => a_value >= 1,
+          "db_primary_#{db_config_name}_count" => a_value >= 1,
+          'db_primary_duration_s' => a_value > 0,
+          "db_primary_#{db_config_name}_duration_s" => a_value > 0
+        )
+      end
+
+      let(:end_payload) do
+        start_payload.merge(db_payload_defaults).merge(
+          'message' => 'TestWorker JID-da883554ee4fe414012f5f42: done: 0.0 sec',
+          'job_status' => 'done',
+          'duration_s' => 0.0,
+          'completed_at' => timestamp.to_f,
+          'cpu_s' => 1.111112
         )
       end
 
@@ -274,59 +288,9 @@ RSpec.describe Gitlab::SidekiqLogging::StructuredLogger do
         end
       end
 
-      context 'when load balancing is disabled' do
-        before do
-          allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false)
-        end
-
-        let(:expected_end_payload_with_db) do
-          expected_end_payload.merge(
-            'db_duration_s' => a_value >= 0.1,
-            'db_count' => a_value >= 1,
-            'db_cached_count' => 0,
-            'db_write_count' => 0
-          )
-        end
-
-        include_examples 'performs database queries'
-      end
-
       context 'when load balancing is enabled', :db_load_balancing do
-        let(:db_config_name) { ::Gitlab::Database.db_config_name(ApplicationRecord.connection) }
-
-        let(:expected_db_payload_defaults) do
-          metrics =
-            ::Gitlab::Metrics::Subscribers::ActiveRecord.load_balancing_metric_counter_keys +
-            ::Gitlab::Metrics::Subscribers::ActiveRecord.load_balancing_metric_duration_keys +
-            ::Gitlab::Metrics::Subscribers::ActiveRecord.db_counter_keys +
-            [:db_duration_s]
-
-          metrics.each_with_object({}) do |key, result|
-            result[key.to_s] = 0
-          end
-        end
-
-        let(:expected_end_payload_with_db) do
-          expected_end_payload.merge(expected_db_payload_defaults).merge(
-            'db_duration_s' => a_value >= 0.1,
-            'db_count' => a_value >= 1,
-            "db_replica_#{db_config_name}_count" => 0,
-            'db_replica_duration_s' => a_value >= 0,
-            'db_primary_count' => a_value >= 1,
-            "db_primary_#{db_config_name}_count" => a_value >= 1,
-            'db_primary_duration_s' => a_value > 0,
-            "db_primary_#{db_config_name}_duration_s" => a_value > 0
-          )
-        end
-
-        let(:end_payload) do
-          start_payload.merge(expected_db_payload_defaults).merge(
-            'message' => 'TestWorker JID-da883554ee4fe414012f5f42: done: 0.0 sec',
-            'job_status' => 'done',
-            'duration_s' => 0.0,
-            'completed_at' => timestamp.to_f,
-            'cpu_s' => 1.111112
-          )
+        let(:db_config_name) do
+          ::Gitlab::Database.db_config_name(ApplicationRecord.retrieve_connection)
         end
 
         include_examples 'performs database queries'

spec/lib/gitlab/sidekiq_middleware/server_metrics_spec.rb

@@ -250,11 +250,6 @@ RSpec.describe Gitlab::SidekiqMiddleware::ServerMetrics do
       end
     end
 
-    context 'when load_balancing is enabled' do
-      before do
-        allow(::Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
-      end
-
     describe '#call' do
       context 'when worker declares data consistency' do
         include_context 'worker declaring data consistency'
@@ -284,31 +279,6 @@ RSpec.describe Gitlab::SidekiqMiddleware::ServerMetrics do
     end
   end
 
-    context 'when load_balancing is disabled' do
-      include_context 'worker declaring data consistency'
-
-      before do
-        allow(::Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false)
-      end
-
-      describe '#initialize' do
-        it 'does not set load_balancing metrics' do
-          expect(Gitlab::Metrics).not_to receive(:counter).with(:sidekiq_load_balancing_count, anything)
-
-          subject
-        end
-      end
-
-      describe '#call' do
-        it 'does not increment load balancing counter' do
-          process_job
-
-          expect(load_balancing_metric).not_to have_received(:increment)
-        end
-      end
-    end
-  end
-
   context 'feature attribution' do
     let(:test_worker) do
       category = worker_category

spec/lib/gitlab/sidekiq_middleware_spec.rb

@@ -28,9 +28,8 @@ RSpec.describe Gitlab::SidekiqMiddleware do
     stub_const('TestWorker', worker_class)
   end
 
-  shared_examples "a middleware chain" do |load_balancing_enabled|
+  shared_examples "a middleware chain" do
     before do
-      allow(::Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(load_balancing_enabled)
       configurator.call(chain)
     end
 
@@ -45,10 +44,10 @@ RSpec.describe Gitlab::SidekiqMiddleware do
     end
   end
 
-  shared_examples "a middleware chain for mailer" do |load_balancing_enabled|
+  shared_examples "a middleware chain for mailer" do
     let(:worker_class) { ActiveJob::QueueAdapters::SidekiqAdapter::JobWrapper }
 
-    it_behaves_like "a middleware chain", load_balancing_enabled
+    it_behaves_like "a middleware chain"
   end
 
   describe '.server_configurator' do
@@ -105,25 +104,8 @@ RSpec.describe Gitlab::SidekiqMiddleware do
     end
 
     context "all optional middlewares on" do
-      context "when load balancing is enabled" do
-        before do
-          allow(::Gitlab::Database::LoadBalancing).to receive_message_chain(:proxy, :load_balancer, :release_host)
-        end
-
-        it_behaves_like "a middleware chain", true
-        it_behaves_like "a middleware chain for mailer", true
-      end
-
-      context "when load balancing is disabled" do
-        let(:disabled_sidekiq_middlewares) do
-          [
-            Gitlab::Database::LoadBalancing::SidekiqServerMiddleware
-          ]
-        end
-
-        it_behaves_like "a middleware chain", false
-        it_behaves_like "a middleware chain for mailer", false
-      end
+      it_behaves_like "a middleware chain"
+      it_behaves_like "a middleware chain for mailer"
     end
 
     context "all optional middlewares off" do
@@ -135,7 +117,6 @@ RSpec.describe Gitlab::SidekiqMiddleware do
         )
       end
 
-      context "when load balancing is enabled" do
       let(:disabled_sidekiq_middlewares) do
         [
           Gitlab::SidekiqMiddleware::ServerMetrics,
@@ -144,27 +125,8 @@ RSpec.describe Gitlab::SidekiqMiddleware do
         ]
       end
 
-        before do
-          allow(::Gitlab::Database::LoadBalancing).to receive_message_chain(:proxy, :load_balancer, :release_host)
-        end
-
-        it_behaves_like "a middleware chain", true
-        it_behaves_like "a middleware chain for mailer", true
-      end
-
-      context "when load balancing is disabled" do
-        let(:disabled_sidekiq_middlewares) do
-          [
-            Gitlab::SidekiqMiddleware::ServerMetrics,
-            Gitlab::SidekiqMiddleware::ArgumentsLogger,
-            Gitlab::SidekiqMiddleware::MemoryKiller,
-            Gitlab::Database::LoadBalancing::SidekiqServerMiddleware
-          ]
-        end
-
-        it_behaves_like "a middleware chain", false
-        it_behaves_like "a middleware chain for mailer", false
-      end
+      it_behaves_like "a middleware chain"
+      it_behaves_like "a middleware chain for mailer"
     end
   end
 
@@ -186,30 +148,7 @@ RSpec.describe Gitlab::SidekiqMiddleware do
       ]
     end
 
-    context "when load balancing is disabled" do
-      let(:disabled_sidekiq_middlewares) do
-        [
-          Gitlab::Database::LoadBalancing::SidekiqClientMiddleware
-        ]
-      end
-
-      it_behaves_like "a middleware chain", false
-      it_behaves_like "a middleware chain for mailer", false
-
-      # Sidekiq documentation states that the worker class could be a string
-      # or a class reference. We should test for both
-      context "worker_class as string value" do
-        let(:worker_args) { [worker_class.to_s, { 'args' => job_args }, queue, redis_pool] }
-        let(:middleware_expected_args) { [worker_class.to_s, hash_including({ 'args' => job_args }), queue, redis_pool] }
-
-        it_behaves_like "a middleware chain", false
-        it_behaves_like "a middleware chain for mailer", false
-      end
-    end
-
-    context "when load balancing is enabled" do
-      it_behaves_like "a middleware chain", true
-      it_behaves_like "a middleware chain for mailer", true
-    end
+    it_behaves_like "a middleware chain"
+    it_behaves_like "a middleware chain for mailer"
   end
 end

spec/lib/peek/views/active_record_spec.rb

@@ -53,78 +53,7 @@ RSpec.describe Peek::Views::ActiveRecord, :request_store do
     allow(connection_primary_2).to receive(:transaction_open?).and_return(true)
     allow(connection_unknown).to receive(:transaction_open?).and_return(false)
     allow(::Gitlab::Database).to receive(:db_config_name).and_return('the_db_config_name')
-  end
 
-  context 'when database load balancing is not enabled' do
-    it 'subscribes and store data into peek views' do
-      Timecop.freeze(2021, 2, 23, 10, 0) do
-        ActiveSupport::Notifications.publish('sql.active_record', Time.current, Time.current + 1.second, '1', event_1)
-        ActiveSupport::Notifications.publish('sql.active_record', Time.current, Time.current + 2.seconds, '2', event_2)
-        ActiveSupport::Notifications.publish('sql.active_record', Time.current, Time.current + 3.seconds, '3', event_3)
-        ActiveSupport::Notifications.publish('sql.active_record', Time.current, Time.current + 4.seconds, '4', event_4)
-      end
-
-      expect(subject.results).to match(
-        calls: 4,
-        summary: {
-          "Cached" => 1,
-          "In a transaction" => 1
-        },
-        duration: '10000.00ms',
-        warnings: ["active-record duration: 10000.0 over 3000"],
-        details: contain_exactly(
-          a_hash_including(
-            start: be_a(Time),
-            cached: '',
-            transaction: '',
-            duration: 1000.0,
-            sql: 'SELECT * FROM users WHERE id = 10',
-            db_config_name: "Config name: the_db_config_name"
-          ),
-          a_hash_including(
-            start: be_a(Time),
-            cached: 'Cached',
-            transaction: '',
-            duration: 2000.0,
-            sql: 'SELECT * FROM users WHERE id = 10',
-            db_config_name: "Config name: the_db_config_name"
-          ),
-          a_hash_including(
-            start: be_a(Time),
-            cached: '',
-            transaction: 'In a transaction',
-            duration: 3000.0,
-            sql: 'UPDATE users SET admin = true WHERE id = 10',
-            db_config_name: "Config name: the_db_config_name"
-          ),
-          a_hash_including(
-            start: be_a(Time),
-            cached: '',
-            transaction: '',
-            duration: 4000.0,
-            sql: 'SELECT VERSION()',
-            db_config_name: "Config name: the_db_config_name"
-          )
-        )
-      )
-    end
-
-    context 'when the GITLAB_MULTIPLE_DATABASE_METRICS env var is disabled' do
-      before do
-        stub_env('GITLAB_MULTIPLE_DATABASE_METRICS', nil)
-      end
-
-      it 'does not include db_config_name field' do
-        ActiveSupport::Notifications.publish('sql.active_record', Time.current, Time.current + 1.second, '1', event_1)
-
-        expect(subject.results[:details][0][:db_config_name]).to be_nil
-      end
-    end
-  end
-
-  context 'when database load balancing is enabled' do
-    before do
-      allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
     allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).with(connection_replica).and_return(:replica)
     allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).with(connection_primary_1).and_return(:primary)
     allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).with(connection_primary_2).and_return(:primary)
@@ -202,5 +131,4 @@ RSpec.describe Peek::Views::ActiveRecord, :request_store do
       expect(subject.results[:details][0][:db_config_name]).to be_nil
     end
   end
-  end
 end

spec/models/application_record_spec.rb

@@ -194,7 +194,7 @@ RSpec.describe ApplicationRecord do
     end
 
     context 'with database load balancing' do
-      let(:session) { double(:session) }
+      let(:session) { Gitlab::Database::LoadBalancing::Session.new }
 
       before do
         allow(::Gitlab::Database::LoadBalancing::Session).to receive(:current).and_return(session)

spec/models/ci/pipeline_spec.rb

@@ -2816,8 +2816,9 @@ RSpec.describe Ci::Pipeline, :mailer, factory_default: :keep do
 
         extra_update_queries = 4 # transition ... => :canceled, queue pop
         extra_generic_commit_status_validation_queries = 2 # name_uniqueness_across_types
+        extra_load_balancer_queries = 3
 
-        expect(control2.count).to eq(control1.count + extra_update_queries + extra_generic_commit_status_validation_queries)
+        expect(control2.count).to eq(control1.count + extra_update_queries + extra_generic_commit_status_validation_queries + extra_load_balancer_queries)
       end
     end
 

spec/models/ci/runner_spec.rb

@@ -383,9 +383,6 @@ RSpec.describe Ci::Runner 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)
 

spec/services/user_project_access_changed_service_spec.rb

@@ -47,8 +47,6 @@ RSpec.describe UserProjectAccessChangedService do
     let(:service) { UserProjectAccessChangedService.new([1, 2]) }
 
     before do
-      allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
-
       expect(AuthorizedProjectsWorker).to receive(:bulk_perform_and_wait)
                                             .with([[1], [2]])
                                             .and_return(10)

spec/support/database_load_balancing.rb

@@ -2,8 +2,6 @@
 
 RSpec.configure do |config|
   config.before(:each, :db_load_balancing) do
-    allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
-
     config = Gitlab::Database::LoadBalancing::Configuration
       .new(ActiveRecord::Base, [Gitlab::Database.main.config['host']])
     lb = ::Gitlab::Database::LoadBalancing::LoadBalancer.new(config)

spec/support/shared_contexts/lib/gitlab/sidekiq_logging/structured_logger_shared_context.rb

@@ -39,17 +39,25 @@ RSpec.shared_context 'structured_logger' do
     )
   end
 
+  let(:db_payload_defaults) do
+    metrics =
+      ::Gitlab::Metrics::Subscribers::ActiveRecord.load_balancing_metric_counter_keys +
+      ::Gitlab::Metrics::Subscribers::ActiveRecord.load_balancing_metric_duration_keys +
+      ::Gitlab::Metrics::Subscribers::ActiveRecord.db_counter_keys +
+      [:db_duration_s]
+
+    metrics.each_with_object({}) do |key, result|
+      result[key.to_s] = 0
+    end
+  end
+
   let(:end_payload) do
-    start_payload.merge(
+    start_payload.merge(db_payload_defaults).merge(
       'message' => 'TestWorker JID-da883554ee4fe414012f5f42: done: 0.0 sec',
       'job_status' => 'done',
       'duration_s' => 0.0,
       'completed_at' => timestamp.to_f,
-      'cpu_s' => 1.111112,
-      'db_duration_s' => 0.0,
-      'db_cached_count' => 0,
-      'db_count' => 0,
-      'db_write_count' => 0
+      'cpu_s' => 1.111112
     )
   end
 

spec/support/shared_contexts/lib/gitlab/sidekiq_middleware/server_metrics_shared_context.rb

@@ -15,6 +15,7 @@ RSpec.shared_context 'server metrics with mocked prometheus' do
   let(:redis_seconds_metric) { double('redis seconds metric') }
   let(:elasticsearch_seconds_metric) { double('elasticsearch seconds metric') }
   let(:elasticsearch_requests_total) { double('elasticsearch calls total metric') }
+  let(:load_balancing_metric) { double('load balancing metric') }
 
   before do
     allow(Gitlab::Metrics).to receive(:histogram).and_call_original
@@ -31,6 +32,7 @@ RSpec.shared_context 'server metrics with mocked prometheus' do
     allow(Gitlab::Metrics).to receive(:counter).with(:sidekiq_jobs_retried_total, anything).and_return(retried_total_metric)
     allow(Gitlab::Metrics).to receive(:counter).with(:sidekiq_redis_requests_total, anything).and_return(redis_requests_total)
     allow(Gitlab::Metrics).to receive(:counter).with(:sidekiq_elasticsearch_requests_total, anything).and_return(elasticsearch_requests_total)
+    allow(Gitlab::Metrics).to receive(:counter).with(:sidekiq_load_balancing_count, anything).and_return(load_balancing_metric)
     allow(Gitlab::Metrics).to receive(:gauge).with(:sidekiq_running_jobs, anything, {}, :all).and_return(running_jobs_metric)
     allow(Gitlab::Metrics).to receive(:gauge).with(:sidekiq_concurrency, anything, {}, :all).and_return(concurrency_metric)
 

spec/support/shared_examples/metrics/active_record_subscriber_shared_examples.rb

@@ -4,14 +4,20 @@ RSpec.shared_examples 'store ActiveRecord info in RequestStore' do |db_role|
   let(:db_config_name) { ::Gitlab::Database.db_config_names.first }
 
   let(:expected_payload_defaults) do
+    result = {}
     metrics =
       ::Gitlab::Metrics::Subscribers::ActiveRecord.load_balancing_metric_counter_keys +
-      ::Gitlab::Metrics::Subscribers::ActiveRecord.load_balancing_metric_duration_keys +
       ::Gitlab::Metrics::Subscribers::ActiveRecord.db_counter_keys
 
-    metrics.each_with_object({}) do |key, result|
+    metrics.each do |key|
       result[key] = 0
     end
+
+    ::Gitlab::Metrics::Subscribers::ActiveRecord.load_balancing_metric_duration_keys.each do |key|
+      result[key] = 0.0
+    end
+
+    result
   end
 
   def transform_hash(hash, another_hash)
@@ -36,8 +42,8 @@ RSpec.shared_examples 'store ActiveRecord info in RequestStore' do |db_role|
                        "db_primary_#{db_config_name}_cached_count": record_cached_query ? 1 : 0,
                        db_primary_count: record_query ? 1 : 0,
                        "db_primary_#{db_config_name}_count": record_query ? 1 : 0,
-                       db_primary_duration_s: record_query ? 0.002 : 0,
-                       "db_primary_#{db_config_name}_duration_s": record_query ? 0.002 : 0,
+                       db_primary_duration_s: record_query ? 0.002 : 0.0,
+                       "db_primary_#{db_config_name}_duration_s": record_query ? 0.002 : 0.0,
                        db_primary_wal_count: record_wal_query ? 1 : 0,
                        "db_primary_#{db_config_name}_wal_count": record_wal_query ? 1 : 0,
                        db_primary_wal_cached_count: record_wal_query && record_cached_query ? 1 : 0,
@@ -52,19 +58,29 @@ RSpec.shared_examples 'store ActiveRecord info in RequestStore' do |db_role|
                        "db_replica_#{db_config_name}_cached_count": record_cached_query ? 1 : 0,
                        db_replica_count: record_query ? 1 : 0,
                        "db_replica_#{db_config_name}_count": record_query ? 1 : 0,
-                       db_replica_duration_s: record_query ? 0.002 : 0,
-                       "db_replica_#{db_config_name}_duration_s": record_query ? 0.002 : 0,
+                       db_replica_duration_s: record_query ? 0.002 : 0.0,
+                       "db_replica_#{db_config_name}_duration_s": record_query ? 0.002 : 0.0,
                        db_replica_wal_count: record_wal_query ? 1 : 0,
                        "db_replica_#{db_config_name}_wal_count": record_wal_query ? 1 : 0,
                        db_replica_wal_cached_count: record_wal_query && record_cached_query ? 1 : 0,
                        "db_replica_#{db_config_name}_wal_cached_count": record_wal_query && record_cached_query ? 1 : 0
                      })
                    else
-                     {
+                     transform_hash(expected_payload_defaults, {
                        db_count: record_query ? 1 : 0,
                        db_write_count: record_write_query ? 1 : 0,
-                       db_cached_count: record_cached_query ? 1 : 0
-                     }
+                       db_cached_count: record_cached_query ? 1 : 0,
+                       db_primary_cached_count: 0,
+                       "db_primary_#{db_config_name}_cached_count": 0,
+                       db_primary_count: 0,
+                       "db_primary_#{db_config_name}_count": 0,
+                       db_primary_duration_s: 0.0,
+                       "db_primary_#{db_config_name}_duration_s": 0.0,
+                       db_primary_wal_count: 0,
+                       "db_primary_#{db_config_name}_wal_count": 0,
+                       db_primary_wal_cached_count: 0,
+                       "db_primary_#{db_config_name}_wal_cached_count": 0
+                     })
                    end
 
         expect(described_class.db_counter_payload).to eq(expected)
@@ -89,7 +105,7 @@ RSpec.shared_examples 'store ActiveRecord info in RequestStore' do |db_role|
 end
 
 RSpec.shared_examples 'record ActiveRecord metrics in a metrics transaction' do |db_role|
-  let(:db_config_name) { ::Gitlab::Database.db_config_name(ApplicationRecord.connection) }
+  let(:db_config_name) { ::Gitlab::Database.db_config_name(ApplicationRecord.retrieve_connection) }
 
   it 'increments only db counters' do
     if record_query