Merge branch 'always-enable-load-balancer' into 'master'

Always enable the database load balancer

See merge request gitlab-org/gitlab!68042

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|
-    config.middleware.use(Gitlab::Database::LoadBalancing::RackMiddleware)
-  end
+Gitlab::Application.configure do |config|
+  config.middleware.use(Gitlab::Database::LoadBalancing::RackMiddleware)
+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.
-    Gitlab::Database::LoadBalancing.start_service_discovery
-  end
+# 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

config/initializers/zz_metrics.rb

@@ -29,12 +29,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.insert_before Gitlab::Database::LoadBalancing::RackMiddleware,
+      Gitlab::Metrics::RackMiddleware
 
     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,23 +13,21 @@ 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.
-          # 2. Gitaly: The UserMergeBranch RPC runs.
-          # 3. Gitaly (gitaly-ruby): This RPC calls the pre-receive hook.
-          # 4. Rails: This hook makes an API request to /api/v4/internal/allowed.
-          # 5. Rails: This API check does a SQL query for locked merge
-          #    requests with a matching SHA.
-          #
-          # Since steps 1 and 5 will happen on different database
-          # sessions, replication lag could erroneously cause step 5 to
-          # report no matching merge requests. To avoid this, we check
-          # the write location to ensure the replica can make this query.
-          track_session_metrics do
-            ::Gitlab::Database::LoadBalancing::Sticking.select_valid_host(:project, @project.id)
-          end
+        # When a user merges a merge request, the following sequence happens:
+        #
+        # 1. Sidekiq: MergeService runs and updates the merge request in a locked state.
+        # 2. Gitaly: The UserMergeBranch RPC runs.
+        # 3. Gitaly (gitaly-ruby): This RPC calls the pre-receive hook.
+        # 4. Rails: This hook makes an API request to /api/v4/internal/allowed.
+        # 5. Rails: This API check does a SQL query for locked merge
+        #    requests with a matching SHA.
+        #
+        # Since steps 1 and 5 will happen on different database
+        # sessions, replication lag could erroneously cause step 5 to
+        # report no matching merge requests. To avoid this, we check
+        # the write location to ensure the replica can make this query.
+        track_session_metrics do
+          ::Gitlab::Database::LoadBalancing::Sticking.select_valid_host(:project, @project.id)
         end
 
         # rubocop: disable CodeReuse/ActiveRecord

lib/gitlab/database/load_balancing.rb

@@ -4,34 +4,22 @@ module Gitlab
   module Database
     module LoadBalancing
       # The exceptions raised for connection errors.
-      CONNECTION_ERRORS = if defined?(PG)
-                            [
-                              PG::ConnectionBad,
-                              PG::ConnectionDoesNotExist,
-                              PG::ConnectionException,
-                              PG::ConnectionFailure,
-                              PG::UnableToSend,
-                              # During a failover this error may be raised when
-                              # writing to a primary.
-                              PG::ReadOnlySqlTransaction
-                            ].freeze
-                          else
-                            [].freeze
-                          end
-
-      ProxyNotConfiguredError = Class.new(StandardError)
+      CONNECTION_ERRORS = [
+        PG::ConnectionBad,
+        PG::ConnectionDoesNotExist,
+        PG::ConnectionException,
+        PG::ConnectionFailure,
+        PG::UnableToSend,
+        # During a failover this error may be raised when
+        # writing to a primary.
+        PG::ReadOnlySqlTransaction,
+        # This error is raised when we can't connect to the database in the
+        # first place (e.g. it's offline or the hostname is incorrect).
+        ActiveRecord::ConnectionNotEstablished
+      ].freeze
 
-      # 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 +27,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,17 +50,12 @@ 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)
-            .perform_service_discovery
-        end
+        ServiceDiscovery
+          .new(proxy.load_balancer, **configuration.service_discovery)
+          .perform_service_discovery
       end
 
       DB_ROLES = [

lib/gitlab/database/load_balancing/host.rb

@@ -9,19 +9,12 @@ module Gitlab
 
         delegate :connection, :release_connection, :enable_query_cache!, :disable_query_cache!, :query_cache_enabled, to: :pool
 
-        CONNECTION_ERRORS =
-          if defined?(PG)
-            [
-              ActionView::Template::Error,
-              ActiveRecord::StatementInvalid,
-              PG::Error
-            ].freeze
-          else
-            [
-              ActionView::Template::Error,
-              ActiveRecord::StatementInvalid
-            ].freeze
-          end
+        CONNECTION_ERRORS = [
+          ActionView::Template::Error,
+          ActiveRecord::StatementInvalid,
+          ActiveRecord::ConnectionNotEstablished,
+          PG::Error
+        ].freeze
 
         # host - The address of the database.
         # load_balancer - The LoadBalancer that manages this Host.

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,13 +47,11 @@ 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?
+          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
+          increment_db_role_counters(db_role, payload)
+          observe_db_role_duration(db_role, event)
         end
 
         def self.db_counter_payload
@@ -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,10 +53,7 @@ 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
+        @metrics[:sidekiq_load_balancing_count] = ::Gitlab::Metrics.counter(:sidekiq_load_balancing_count, 'Sidekiq jobs with load balancing')
       end
 
       def call(worker, job, queue)
@@ -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,10 +44,8 @@ module Peek
           count[item[:transaction]] += 1
         end
 
-        if ::Gitlab::Database::LoadBalancing.enable?
-          count[item[:db_role]] ||= 0
-          count[item[:db_role]] += 1
-        end
+        count[item[:db_role]] ||= 0
+        count[item[:db_role]] += 1
       end
 
       def setup_subscribers
@@ -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,39 +230,21 @@ 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
-              subscriber.process_action(event)
-
-              expect(log_data).to include(*db_load_balancing_logging_keys)
-            end
-          end
-
-          context 'when RequestStore is disabled' do
-            it 'does not include db counters for load balancing' do
-              subscriber.process_action(event)
+      context 'with db payload' do
+        context 'when RequestStore is enabled', :request_store do
+          it 'includes db counters for load balancing' do
+            subscriber.process_action(event)
 
-              expect(log_data).not_to include(*db_load_balancing_logging_keys)
-            end
+            expect(log_data).to include(*db_load_balancing_logging_keys)
           end
         end
-      end
 
-      context 'when load balancing is disabled' do
-        before do
-          allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false)
-        end
+        context 'when RequestStore is disabled' do
+          it 'does not include db counters for load balancing' do
+            subscriber.process_action(event)
 
-        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)
+            expect(log_data).not_to include(*db_load_balancing_logging_keys)
+          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/host_spec.rb

@@ -172,6 +172,14 @@ RSpec.describe Gitlab::Database::LoadBalancing::Host do
 
         expect(host).not_to be_online
       end
+
+      it 'returns false when ActiveRecord::ConnectionNotEstablished is raised' do
+        allow(host)
+          .to receive(:check_replica_status?)
+          .and_raise(ActiveRecord::ConnectionNotEstablished)
+
+        expect(host).not_to be_online
+      end
     end
   end
 

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

@@ -277,31 +277,26 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
   end
 
   describe '#connection_error?' do
-    before do
-      stub_const('Gitlab::Database::LoadBalancing::LoadBalancer::CONNECTION_ERRORS',
-                 [NotImplementedError])
-    end
-
     it 'returns true for a connection error' do
-      error = NotImplementedError.new
+      error = ActiveRecord::ConnectionNotEstablished.new
 
       expect(lb.connection_error?(error)).to eq(true)
     end
 
     it 'returns true for a wrapped connection error' do
-      wrapped = wrapped_exception(ActiveRecord::StatementInvalid, NotImplementedError)
+      wrapped = wrapped_exception(ActiveRecord::StatementInvalid, ActiveRecord::ConnectionNotEstablished)
 
       expect(lb.connection_error?(wrapped)).to eq(true)
     end
 
     it 'returns true for a wrapped connection error from a view' do
-      wrapped = wrapped_exception(ActionView::Template::Error, NotImplementedError)
+      wrapped = wrapped_exception(ActionView::Template::Error, ActiveRecord::ConnectionNotEstablished)
 
       expect(lb.connection_error?(wrapped)).to eq(true)
     end
 
     it 'returns true for deeply wrapped/nested errors' do
-      top = twice_wrapped_exception(ActionView::Template::Error, ActiveRecord::StatementInvalid, NotImplementedError)
+      top = twice_wrapped_exception(ActionView::Template::Error, ActiveRecord::StatementInvalid, ActiveRecord::ConnectionNotEstablished)
 
       expect(lb.connection_error?(top)).to eq(true)
     end

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,39 +8,24 @@ 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)
+    it 'does not stick if no write was performed' do
+      allow(Gitlab::Database::LoadBalancing::Session.current)
+        .to receive(:performed_write?)
+        .and_return(false)
 
-        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?)
-          .and_return(false)
+      expect(described_class).not_to receive(:stick)
 
-        expect(described_class).not_to receive(:stick)
-
-        described_class.stick_if_necessary(:user, 42)
-      end
+      described_class.stick_if_necessary(:user, 42)
+    end
 
-      it 'sticks to the primary if a write was performed' do
-        allow(Gitlab::Database::LoadBalancing::Session.current)
-          .to receive(:performed_write?)
-          .and_return(true)
+    it 'sticks to the primary if a write was performed' do
+      allow(Gitlab::Database::LoadBalancing::Session.current)
+        .to receive(:performed_write?)
+        .and_return(true)
 
-        expect(described_class).to receive(:stick).with(:user, 42)
+      expect(described_class).to receive(:stick).with(:user, 42)
 
-        described_class.stick_if_necessary(:user, 42)
-      end
+      described_class.stick_if_necessary(:user, 42)
     end
   end
 
@@ -155,35 +140,22 @@ 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!)
+    before do
+      lb = double(:lb, primary_write_location: 'foo')
 
-        described_class.bulk_stick(:user, ids)
-      end
+      allow(described_class).to receive(:load_balancer).and_return(lb)
     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)
+    it 'sticks an entity to the primary', :aggregate_failures do
+      ids.each do |id|
+        expect(described_class).to receive(:set_write_location_for)
+                                      .with(:user, id, 'foo')
       end
 
-      it 'sticks an entity to the primary', :aggregate_failures do
-        ids.each do |id|
-          expect(described_class).to receive(:set_write_location_for)
-                                       .with(:user, id, 'foo')
-        end
-
-        expect(Gitlab::Database::LoadBalancing::Session.current)
-          .to receive(:use_primary!)
+      expect(Gitlab::Database::LoadBalancing::Session.current)
+        .to receive(:use_primary!)
 
-        subject
-      end
+      subject
     end
   end
 
@@ -202,63 +174,15 @@ 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')
 
-      it 'updates the write location with the load balancer' do
-        lb = double(:lb, primary_write_location: 'foo')
-
-        allow(described_class).to receive(:load_balancer).and_return(lb)
-
-        expect(described_class).to receive(:set_write_location_for)
-          .with(:user, 42, 'foo')
-
-        described_class.mark_primary_write_location(:user, 42)
-      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
+      allow(described_class).to receive(:load_balancer).and_return(lb)
 
-      it 'updates the write location with the main ActiveRecord connection' do
-        expect(described_class).not_to receive(:set_write_location_for)
+      expect(described_class).to receive(:set_write_location_for)
+        .with(:user, 42, 'foo')
 
-        described_class.mark_primary_write_location(:user, 42)
-      end
+      described_class.mark_primary_write_location(:user, 42)
     end
   end
 

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
+    it 'returns the connection proxy' do
+      proxy = double(:connection_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
-
-    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,46 +26,53 @@ RSpec.describe Gitlab::Database::LoadBalancing do
     end
   end
 
-  describe '.enable?' 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([])
+  describe '.enable_replicas?' do
+    context 'when hosts are specified' do
+      before do
+        allow(described_class.configuration)
+          .to receive(:hosts)
+          .and_return(%w(foo))
+      end
 
-      expect(described_class.enable?).to eq(false)
-    end
+      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)
+      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)
-    end
+        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)
+      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
     end
 
-    it 'returns true when load balancing should be enabled' do
-      allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(false)
+    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)
 
-      expect(described_class.enable?).to eq(true)
-    end
+        allow(described_class.configuration)
+          .to receive(:service_discovery_enabled?)
+          .and_return(true)
 
-    it 'returns true when service discovery is enabled' do
-      allow(described_class.configuration).to receive(:hosts).and_return([])
-      allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(false)
+        expect(described_class.enable_replicas?).to eq(true)
+      end
+    end
 
-      allow(described_class.configuration)
-        .to receive(:service_discovery_enabled?)
-        .and_return(true)
+    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?).to eq(true)
+        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)
-      end
+      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,26 +158,15 @@ 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
+        it 'reads from replica' do
+          expect(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_replicas_for_read_queries).and_call_original
 
-            subject.execute
-          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
+          subject.execute
         end
       end
 

spec/lib/gitlab/instrumentation_helper_spec.rb

@@ -107,69 +107,44 @@ 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
-
-        expect(payload).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
-
-      context 'when replica caught up search was made' do
-        before do
-          Gitlab::SafeRequestStore[:caught_up_replica_pick_ok] = 2
-          Gitlab::SafeRequestStore[:caught_up_replica_pick_fail] = 1
-        end
+    it 'includes DB counts' do
+      subject
 
-        it 'includes related metrics' do
-          subject
+      expect(payload).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
 
-          expect(payload).to include(caught_up_replica_pick_ok: 2)
-          expect(payload).to include(caught_up_replica_pick_fail: 1)
-        end
+    context 'when replica caught up search was made' do
+      before do
+        Gitlab::SafeRequestStore[:caught_up_replica_pick_ok] = 2
+        Gitlab::SafeRequestStore[:caught_up_replica_pick_fail] = 1
       end
 
-      context 'when only a single counter was updated' do
-        before do
-          Gitlab::SafeRequestStore[:caught_up_replica_pick_ok] = 1
-          Gitlab::SafeRequestStore[:caught_up_replica_pick_fail] = nil
-        end
-
-        it 'includes only that counter into logging' do
-          subject
+      it 'includes related metrics' do
+        subject
 
-          expect(payload).to include(caught_up_replica_pick_ok: 1)
-          expect(payload).not_to include(:caught_up_replica_pick_fail)
-        end
+        expect(payload).to include(caught_up_replica_pick_ok: 2)
+        expect(payload).to include(caught_up_replica_pick_fail: 1)
       end
     end
 
-    context 'when load balancing is disabled' do
+    context 'when only a single counter was updated' do
       before do
-        allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false)
+        Gitlab::SafeRequestStore[:caught_up_replica_pick_ok] = 1
+        Gitlab::SafeRequestStore[:caught_up_replica_pick_fail] = nil
       end
 
-      it 'does not include DB counts' do
+      it 'includes only that counter into logging' 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)
+        expect(payload).to include(caught_up_replica_pick_ok: 1)
+        expect(payload).not_to include(:caught_up_replica_pick_fail)
       end
     end
   end

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,60 +250,30 @@ 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'
-
-          it 'increments load balancing counter with defined data consistency' do
-            process_job
-
-            expect(load_balancing_metric).to have_received(:increment).with(
-              a_hash_including(
-                data_consistency: :delayed,
-                load_balancing_strategy: 'replica'
-              ), 1)
-          end
-        end
-
-        context 'when worker does not declare data consistency' do
-          it 'increments load balancing counter with default data consistency' do
-            process_job
-
-            expect(load_balancing_metric).to have_received(:increment).with(
-              a_hash_including(
-                data_consistency: :always,
-                load_balancing_strategy: 'primary'
-              ), 1)
-          end
-        end
-      end
-    end
-
-    context 'when load_balancing is disabled' do
-      include_context 'worker declaring data consistency'
+    describe '#call' do
+      context 'when worker declares data consistency' 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)
+        it 'increments load balancing counter with defined data consistency' do
+          process_job
 
-          subject
+          expect(load_balancing_metric).to have_received(:increment).with(
+            a_hash_including(
+              data_consistency: :delayed,
+              load_balancing_strategy: 'replica'
+            ), 1)
         end
       end
 
-      describe '#call' do
-        it 'does not increment load balancing counter' do
+      context 'when worker does not declare data consistency' do
+        it 'increments load balancing counter with default data consistency' do
           process_job
 
-          expect(load_balancing_metric).not_to have_received(:increment)
+          expect(load_balancing_metric).to have_received(:increment).with(
+            a_hash_including(
+              data_consistency: :always,
+              load_balancing_strategy: 'primary'
+            ), 1)
         end
       end
     end

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,36 +117,16 @@ RSpec.describe Gitlab::SidekiqMiddleware do
         )
       end
 
-      context "when load balancing is enabled" do
-        let(:disabled_sidekiq_middlewares) do
-          [
-            Gitlab::SidekiqMiddleware::ServerMetrics,
-            Gitlab::SidekiqMiddleware::ArgumentsLogger,
-            Gitlab::SidekiqMiddleware::MemoryKiller
-          ]
-        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
+      let(:disabled_sidekiq_middlewares) do
+        [
+          Gitlab::SidekiqMiddleware::ServerMetrics,
+          Gitlab::SidekiqMiddleware::ArgumentsLogger,
+          Gitlab::SidekiqMiddleware::MemoryKiller
+        ]
       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,154 +53,82 @@ 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')
+
+    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)
+    allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).with(connection_unknown).and_return(nil)
   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
+  it 'includes db role data and db_config_name name' 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"
-          )
+    expect(subject.results).to match(
+      calls: 4,
+      summary: {
+        "Cached" => 1,
+        "In a transaction" => 1,
+        "Role: Primary" => 2,
+        "Role: Replica" => 1,
+        "Role: Unknown" => 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_role: 'Role: Primary',
+          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_role: 'Role: Replica',
+          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_role: 'Role: Primary',
+          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_role: 'Role: Unknown',
+          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
+  context 'when the GITLAB_MULTIPLE_DATABASE_METRICS env var is disabled' 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)
-      allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).with(connection_unknown).and_return(nil)
+      stub_env('GITLAB_MULTIPLE_DATABASE_METRICS', nil)
     end
 
-    it 'includes db role data and db_config_name name' 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,
-          "Role: Primary" => 2,
-          "Role: Replica" => 1,
-          "Role: Unknown" => 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_role: 'Role: Primary',
-            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_role: 'Role: Replica',
-            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_role: 'Role: Primary',
-            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_role: 'Role: Unknown',
-            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)
+    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
+      expect(subject.results[:details][0][:db_config_name]).to be_nil
     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