Merge branch 'load-balancer-multiple-databases' into 'master'

Support load balancing of multiple database clusters

See merge request gitlab-org/gitlab!67773

app/models/ci/build.rb

@@ -1109,7 +1109,7 @@ module Ci
       return unless saved_change_to_status?
       return unless running?
 
-      ::Gitlab::Database::LoadBalancing::Sticking.stick(:build, id)
+      self.class.sticking.stick(:build, id)
     end
 
     def status_commit_hooks

app/models/ci/ci_database_record.rb

@@ -12,13 +12,6 @@ module Ci
 
     if Gitlab::Database.has_config?(:ci)
       connects_to database: { writing: :ci, reading: :ci }
-
-      # TODO: Load Balancing messes with `CiDatabaseRecord`
-      # returning wrong connection. To be removed once merged:
-      # https://gitlab.com/gitlab-org/gitlab/-/merge_requests/67773
-      def self.connection
-        retrieve_connection
-      end
     end
   end
 end

app/models/ci/runner.rb

@@ -348,7 +348,7 @@ module Ci
       # intention here is not to execute `Ci::RegisterJobService#execute` on
       # the primary database.
       #
-      ::Gitlab::Database::LoadBalancing::Sticking.stick(:runner, id)
+      ::Ci::Runner.sticking.stick(:runner, id)
 
       SecureRandom.hex.tap do |new_update|
         ::Gitlab::Workhorse.set_key_and_notify(runner_queue_key, new_update,

app/models/project.rb

@@ -2400,7 +2400,7 @@ class Project < ApplicationRecord
   end
 
   def mark_primary_write_location
-    ::Gitlab::Database::LoadBalancing::Sticking.mark_primary_write_location(:project, self.id)
+    self.class.sticking.mark_primary_write_location(:project, self.id)
   end
 
   def toggle_ci_cd_settings!(settings_attribute)

app/services/ci/register_job_service.rb

@@ -22,7 +22,8 @@ module Ci
     end
 
     def execute(params = {})
-      db_all_caught_up = ::Gitlab::Database::LoadBalancing::Sticking.all_caught_up?(:runner, runner.id)
+      db_all_caught_up =
+        ::Ci::Runner.sticking.all_caught_up?(:runner, runner.id)
 
       @metrics.increment_queue_operation(:queue_attempt)
 

app/services/user_project_access_changed_service.rb

@@ -30,7 +30,7 @@ class UserProjectAccessChangedService
         end
       end
 
-    ::Gitlab::Database::LoadBalancing::Sticking.bulk_stick(:user, @user_ids)
+    ::User.sticking.bulk_stick(:user, @user_ids)
 
     result
   end

config/initializers/load_balancing.rb

 # frozen_string_literal: true
 
-ActiveRecord::Base.singleton_class.attr_accessor :load_balancing_proxy
-
-Gitlab::Database.main.disable_prepared_statements
-
 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?
-  )
-
-  ActiveRecord::Base.load_balancing_proxy =
-    Gitlab::Database::LoadBalancing::ConnectionProxy.new(lb)
-
-  # Populate service discovery immediately if it is configured
-  Gitlab::Database::LoadBalancing.perform_service_discovery
-end
-
-setup.call
-
-# 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 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
+Gitlab::Database::LoadBalancing.base_models.each do |model|
+  # 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.
+  Gitlab::Database::LoadBalancing::Setup.new(model).setup
+
+  # 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.
+    model.connection.load_balancer.disconnect!(timeout: 0)
+  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
+    Gitlab::Database::LoadBalancing::Setup
+      .new(model, start_service_discovery: true)
+      .setup
+  end
 end

lib/api/ci/helpers/runner.rb

@@ -42,8 +42,7 @@ module API
           token = params[:token]
 
           if token
-            ::Gitlab::Database::LoadBalancing::RackMiddleware
-              .stick_or_unstick(env, :runner, token)
+            ::Ci::Runner.sticking.stick_or_unstick_request(env, :runner, token)
           end
 
           strong_memoize(:current_runner) do
@@ -80,8 +79,9 @@ module API
           id = params[:id]
 
           if id
-            ::Gitlab::Database::LoadBalancing::RackMiddleware
-              .stick_or_unstick(env, :build, id)
+            ::Ci::Build
+              .sticking
+              .stick_or_unstick_request(env, :build, id)
           end
 
           strong_memoize(:current_job) do

lib/api/helpers.rb

@@ -75,8 +75,9 @@ module API
       save_current_user_in_env(@current_user) if @current_user
 
       if @current_user
-        ::Gitlab::Database::LoadBalancing::RackMiddleware
-          .stick_or_unstick(env, :user, @current_user.id)
+        ::ApplicationRecord
+          .sticking
+          .stick_or_unstick_request(env, :user, @current_user.id)
       end
 
       @current_user

lib/gitlab/checks/matching_merge_request.rb

@@ -27,7 +27,7 @@ module Gitlab
         # 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)
+          ::ApplicationRecord.sticking.select_valid_host(:project, @project.id)
         end
 
         # rubocop: disable CodeReuse/ActiveRecord

lib/gitlab/ci/trace.rb

@@ -286,7 +286,8 @@ module Gitlab
 
       def destroy_stream(build)
         if consistent_archived_trace?(build)
-          ::Gitlab::Database::LoadBalancing::Sticking
+          ::Ci::Build
+            .sticking
             .stick(LOAD_BALANCING_STICKING_NAMESPACE, build.id)
         end
 
@@ -295,7 +296,8 @@ module Gitlab
 
       def read_trace_artifact(build)
         if consistent_archived_trace?(build)
-          ::Gitlab::Database::LoadBalancing::Sticking
+          ::Ci::Build
+            .sticking
             .unstick_or_continue_sticking(LOAD_BALANCING_STICKING_NAMESPACE, build.id)
         end
 

lib/gitlab/database.rb

@@ -53,7 +53,12 @@ module Gitlab
 
     def self.database_base_models
       @database_base_models ||= {
-        main: ::ApplicationRecord,
+        # Note that we use ActiveRecord::Base here and not ApplicationRecord.
+        # This is deliberate, as we also use these classes to apply load
+        # balancing to, and the load balancer must be enabled for _all_ models
+        # that inher from ActiveRecord::Base; not just our own models that
+        # inherit from ApplicationRecord.
+        main: ::ActiveRecord::Base,
         ci: ::Ci::CiDatabaseRecord.connection_class? ? ::Ci::CiDatabaseRecord : nil
       }.compact.freeze
     end

lib/gitlab/database/load_balancing.rb

@@ -18,44 +18,20 @@ module Gitlab
         ActiveRecord::ConnectionNotEstablished
       ].freeze
 
-      def self.proxy
-        ActiveRecord::Base.load_balancing_proxy
+      def self.base_models
+        @base_models ||= ::Gitlab::Database.database_base_models.values.freeze
       end
 
-      # Returns a Hash containing the load balancing configuration.
-      def self.configuration
-        @configuration ||= Configuration.for_model(ActiveRecord::Base)
-      end
-
-      # 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?
+      def self.each_load_balancer
+        return to_enum(__method__) unless block_given?
 
-        configured?
-      end
-
-      def self.configured?
-        configuration.load_balancing_enabled? ||
-          configuration.service_discovery_enabled?
-      end
-
-      def self.start_service_discovery
-        return unless configuration.service_discovery_enabled?
-
-        ServiceDiscovery
-          .new(proxy.load_balancer, **configuration.service_discovery)
-          .start
+        base_models.each do |model|
+          yield model.connection.load_balancer
+        end
       end
 
-      def self.perform_service_discovery
-        return unless configuration.service_discovery_enabled?
-
-        ServiceDiscovery
-          .new(proxy.load_balancer, **configuration.service_discovery)
-          .perform_service_discovery
+      def self.release_hosts
+        each_load_balancer(&:release_host)
       end
 
       DB_ROLES = [

lib/gitlab/database/load_balancing/action_cable_callbacks.rb

@@ -16,7 +16,7 @@ module Gitlab
 
             inner.call
           ensure
-            ::Gitlab::Database::LoadBalancing.proxy.load_balancer.release_host
+            ::Gitlab::Database::LoadBalancing.release_hosts
             ::Gitlab::Database::LoadBalancing::Session.clear_session
           end
         end

lib/gitlab/database/load_balancing/active_record_proxy.rb

-# frozen_string_literal: true
-
-module Gitlab
-  module Database
-    module LoadBalancing
-      # Module injected into ActiveRecord::Base to allow hijacking of the
-      # "connection" method.
-      module ActiveRecordProxy
-        def connection
-          ::Gitlab::Database::LoadBalancing.proxy || super
-        end
-      end
-    end
-  end
-end

lib/gitlab/database/load_balancing/configuration.rb

@@ -72,7 +72,14 @@ module Gitlab
             Database.default_pool_size
         end
 
+        # 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 load_balancing_enabled?
+          return false if Gitlab::Runtime.rake?
+
           hosts.any? || service_discovery_enabled?
         end
 

lib/gitlab/database/load_balancing/load_balancer.rb

@@ -12,22 +12,22 @@ module Gitlab
 
         REPLICA_SUFFIX = '_replica'
 
-        attr_reader :host_list, :configuration
+        attr_reader :name, :host_list, :configuration
 
         # configuration - An instance of `LoadBalancing::Configuration` that
         #                 contains the configuration details (such as the hosts)
         #                 for this load balancer.
-        # primary_only - If set, the replicas are ignored and the primary is
-        #                always used.
-        def initialize(configuration, primary_only: false)
+        def initialize(configuration)
           @configuration = configuration
-          @primary_only = primary_only
+          @primary_only = !configuration.load_balancing_enabled?
           @host_list =
-            if primary_only
+            if @primary_only
               HostList.new([PrimaryHost.new(self)])
             else
               HostList.new(configuration.hosts.map { |addr| Host.new(addr, self) })
             end
+
+          @name = @configuration.model.connection_db_config.name.to_sym
         end
 
         def primary_only?

lib/gitlab/database/load_balancing/rack_middleware.rb

@@ -9,21 +9,6 @@ module Gitlab
       class RackMiddleware
         STICK_OBJECT = 'load_balancing.stick_object'
 
-        # Unsticks or continues sticking the current request.
-        #
-        # This method also updates the Rack environment so #call can later
-        # determine if we still need to stick or not.
-        #
-        # env - The Rack environment.
-        # namespace - The namespace to use for sticking.
-        # id - The identifier to use for sticking.
-        def self.stick_or_unstick(env, namespace, id)
-          ::Gitlab::Database::LoadBalancing::Sticking.unstick_or_continue_sticking(namespace, id)
-
-          env[STICK_OBJECT] ||= Set.new
-          env[STICK_OBJECT] << [namespace, id]
-        end
-
         def initialize(app)
           @app = app
         end
@@ -51,41 +36,46 @@ module Gitlab
         # Typically this code will only be reachable for Rails requests as
         # Grape data is not yet available at this point.
         def unstick_or_continue_sticking(env)
-          namespaces_and_ids = sticking_namespaces_and_ids(env)
+          namespaces_and_ids = sticking_namespaces(env)
 
-          namespaces_and_ids.each do |namespace, id|
-            ::Gitlab::Database::LoadBalancing::Sticking.unstick_or_continue_sticking(namespace, id)
+          namespaces_and_ids.each do |(model, namespace, id)|
+            model.sticking.unstick_or_continue_sticking(namespace, id)
           end
         end
 
         # Determine if we need to stick after handling a request.
         def stick_if_necessary(env)
-          namespaces_and_ids = sticking_namespaces_and_ids(env)
+          namespaces_and_ids = sticking_namespaces(env)
 
-          namespaces_and_ids.each do |namespace, id|
-            ::Gitlab::Database::LoadBalancing::Sticking.stick_if_necessary(namespace, id)
+          namespaces_and_ids.each do |model, namespace, id|
+            model.sticking.stick_if_necessary(namespace, id)
           end
         end
 
         def clear
-          load_balancer.release_host
+          ::Gitlab::Database::LoadBalancing.release_hosts
           ::Gitlab::Database::LoadBalancing::Session.clear_session
         end
 
-        def load_balancer
-          ::Gitlab::Database::LoadBalancing.proxy.load_balancer
-        end
-
         # Determines the sticking namespace and identifier based on the Rack
         # environment.
         #
         # For Rails requests this uses warden, but Grape and others have to
         # manually set the right environment variable.
-        def sticking_namespaces_and_ids(env)
+        def sticking_namespaces(env)
           warden = env['warden']
 
           if warden && warden.user
-            [[:user, warden.user.id]]
+            # When sticking per user, _only_ sticking the main connection could
+            # result in the application trying to read data from a different
+            # connection, while that data isn't available yet.
+            #
+            # To prevent this from happening, we scope sticking to all the
+            # models that support load balancing. In the future (if we
+            # determined this to be OK) we may be able to relax this.
+            LoadBalancing.base_models.map do |model|
+              [model, :user, warden.user.id]
+            end
           elsif env[STICK_OBJECT].present?
             env[STICK_OBJECT].to_a
           else

lib/gitlab/database/load_balancing/setup.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module LoadBalancing
+      # Class for setting up load balancing of a specific model.
+      class Setup
+        attr_reader :configuration
+
+        def initialize(model, start_service_discovery: false)
+          @model = model
+          @configuration = Configuration.for_model(model)
+          @start_service_discovery = start_service_discovery
+        end
+
+        def setup
+          disable_prepared_statements
+          setup_load_balancer
+          setup_service_discovery
+        end
+
+        def disable_prepared_statements
+          db_config_object = @model.connection_db_config
+          config =
+            db_config_object.configuration_hash.merge(prepared_statements: false)
+
+          hash_config = ActiveRecord::DatabaseConfigurations::HashConfig.new(
+            db_config_object.env_name,
+            db_config_object.name,
+            config
+          )
+
+          @model.establish_connection(hash_config)
+        end
+
+        def setup_load_balancer
+          lb = LoadBalancer.new(configuration)
+
+          # We just use a simple `class_attribute` here so we don't need to
+          # inject any modules and/or expose unnecessary methods.
+          @model.class_attribute(:connection)
+          @model.class_attribute(:sticking)
+
+          @model.connection = ConnectionProxy.new(lb)
+          @model.sticking = Sticking.new(lb)
+        end
+
+        def setup_service_discovery
+          return unless configuration.service_discovery_enabled?
+
+          lb = @model.connection.load_balancer
+          sv = ServiceDiscovery.new(lb, **configuration.service_discovery)
+
+          sv.perform_service_discovery
+
+          sv.start if @start_service_discovery
+        end
+      end
+    end
+  end
+end

lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb

@@ -30,26 +30,23 @@ module Gitlab
         end
 
         def set_data_consistency_locations!(job)
-          # Once we add support for multiple databases to our load balancer, we would use something like this:
-          #   job['wal_locations'] = Gitlab::Database.databases.transform_values do |connection|
-          #      connection.load_balancer.primary_write_location
-          #   end
-          #
-          job['wal_locations'] = { ::Gitlab::Database::MAIN_DATABASE_NAME.to_sym => wal_location } if wal_location
-        end
+          locations = {}
 
-        def wal_location
-          strong_memoize(:wal_location) do
-            if ::Gitlab::Database::LoadBalancing::Session.current.use_primary?
-              load_balancer.primary_write_location
-            else
-              load_balancer.host.database_replica_location
+          ::Gitlab::Database::LoadBalancing.each_load_balancer do |lb|
+            if (location = wal_location_for(lb))
+              locations[lb.name] = location
             end
           end
+
+          job['wal_locations'] = locations
         end
 
-        def load_balancer
-          ::Gitlab::Database::LoadBalancing.proxy.load_balancer
+        def wal_location_for(load_balancer)
+          if ::Gitlab::Database::LoadBalancing::Session.current.use_primary?
+            load_balancer.primary_write_location
+          else
+            load_balancer.host.database_replica_location
+          end
         end
       end
     end

lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb

@@ -29,7 +29,7 @@ module Gitlab
         private
 
         def clear
-          release_hosts
+          LoadBalancing.release_hosts
           Session.clear_session
         end
 
@@ -44,7 +44,7 @@ module Gitlab
 
           return :primary_no_wal unless wal_locations
 
-          if all_databases_has_replica_caught_up?(wal_locations)
+          if databases_in_sync?(wal_locations)
             # Happy case: we can read from a replica.
             retried_before?(worker_class, job) ? :replica_retried : :replica
           elsif can_retry?(worker_class, job)
@@ -89,27 +89,18 @@ module Gitlab
           job['retry_count'].nil?
         end
 
-        def all_databases_has_replica_caught_up?(wal_locations)
-          wal_locations.all? do |_config_name, location|
-            # Once we add support for multiple databases to our load balancer, we would use something like this:
-            # Gitlab::Database.databases[config_name].load_balancer.select_up_to_date_host(location)
-            load_balancer.select_up_to_date_host(location)
+        def databases_in_sync?(wal_locations)
+          LoadBalancing.each_load_balancer.all? do |lb|
+            if (location = wal_locations[lb.name])
+              lb.select_up_to_date_host(location)
+            else
+              # If there's no entry for a load balancer it means the Sidekiq
+              # job doesn't care for it. In this case we'll treat the load
+              # balancer as being in sync.
+              true
+            end
           end
         end
-
-        def release_hosts
-          # Once we add support for multiple databases to our load balancer, we would use something like this:
-          # connection.load_balancer.primary_write_location
-          #
-          # Gitlab::Database.databases.values.each do |connection|
-          #   connection.load_balancer.release_host
-          # end
-          load_balancer.release_host
-        end
-
-        def load_balancer
-          LoadBalancing.proxy.load_balancer
-        end
       end
     end
   end

lib/gitlab/database/load_balancing/sticking.rb

@@ -5,34 +5,47 @@ module Gitlab
     module LoadBalancing
       # Module used for handling sticking connections to a primary, if
       # necessary.
-      #
-      # ## Examples
-      #
-      # Sticking a user to the primary:
-      #
-      #     Sticking.stick_if_necessary(:user, current_user.id)
-      #
-      # To unstick if possible, or continue using the primary otherwise:
-      #
-      #     Sticking.unstick_or_continue_sticking(:user, current_user.id)
-      module Sticking
+      class Sticking
         # The number of seconds after which a session should stop reading from
         # the primary.
         EXPIRATION = 30
 
+        def initialize(load_balancer)
+          @load_balancer = load_balancer
+          @model = load_balancer.configuration.model
+        end
+
+        # Unsticks or continues sticking the current request.
+        #
+        # This method also updates the Rack environment so #call can later
+        # determine if we still need to stick or not.
+        #
+        # env - The Rack environment.
+        # namespace - The namespace to use for sticking.
+        # id - The identifier to use for sticking.
+        # model - The ActiveRecord model to scope sticking to.
+        def stick_or_unstick_request(env, namespace, id)
+          unstick_or_continue_sticking(namespace, id)
+
+          env[RackMiddleware::STICK_OBJECT] ||= Set.new
+          env[RackMiddleware::STICK_OBJECT] << [@model, namespace, id]
+        end
+
         # Sticks to the primary if a write was performed.
-        def self.stick_if_necessary(namespace, id)
+        def stick_if_necessary(namespace, id)
           stick(namespace, id) if Session.current.performed_write?
         end
 
-        # Checks if we are caught-up with all the work
-        def self.all_caught_up?(namespace, id)
+        def all_caught_up?(namespace, id)
           location = last_write_location_for(namespace, id)
 
           return true unless location
 
-          load_balancer.select_up_to_date_host(location).tap do |found|
-            ActiveSupport::Notifications.instrument('caught_up_replica_pick.load_balancing', { result: found } )
+          @load_balancer.select_up_to_date_host(location).tap do |found|
+            ActiveSupport::Notifications.instrument(
+              'caught_up_replica_pick.load_balancing',
+              { result: found }
+            )
 
             unstick(namespace, id) if found
           end
@@ -43,7 +56,7 @@ module Gitlab
         # in another thread.
         #
         # Returns true if one host was selected.
-        def self.select_caught_up_replicas(namespace, id)
+        def select_caught_up_replicas(namespace, id)
           location = last_write_location_for(namespace, id)
 
           # Unlike all_caught_up?, we return false if no write location exists.
@@ -51,33 +64,36 @@ module Gitlab
           # write location. If no such location exists, err on the side of caution.
           return false unless location
 
-          load_balancer.select_up_to_date_host(location).tap do |selected|
+          @load_balancer.select_up_to_date_host(location).tap do |selected|
             unstick(namespace, id) if selected
           end
         end
 
         # Sticks to the primary if necessary, otherwise unsticks an object (if
         # it was previously stuck to the primary).
-        def self.unstick_or_continue_sticking(namespace, id)
-          Session.current.use_primary! unless all_caught_up?(namespace, id)
+        def unstick_or_continue_sticking(namespace, id)
+          return if all_caught_up?(namespace, id)
+
+          Session.current.use_primary!
         end
 
         # Select a replica that has caught up with the primary. If one has not been
         # found, stick to the primary.
-        def self.select_valid_host(namespace, id)
-          replica_selected = select_caught_up_replicas(namespace, id)
+        def select_valid_host(namespace, id)
+          replica_selected =
+            select_caught_up_replicas(namespace, id)
 
           Session.current.use_primary! unless replica_selected
         end
 
         # Starts sticking to the primary for the given namespace and id, using
         # the latest WAL pointer from the primary.
-        def self.stick(namespace, id)
+        def stick(namespace, id)
           mark_primary_write_location(namespace, id)
           Session.current.use_primary!
         end
 
-        def self.bulk_stick(namespace, ids)
+        def bulk_stick(namespace, ids)
           with_primary_write_location do |location|
             ids.each do |id|
               set_write_location_for(namespace, id, location)
@@ -87,45 +103,49 @@ module Gitlab
           Session.current.use_primary!
         end
 
-        def self.with_primary_write_location
-          location = load_balancer.primary_write_location
+        def with_primary_write_location
+          location = @load_balancer.primary_write_location
 
           return if location.blank?
 
           yield(location)
         end
 
-        def self.mark_primary_write_location(namespace, id)
+        def mark_primary_write_location(namespace, id)
           with_primary_write_location do |location|
             set_write_location_for(namespace, id, location)
           end
         end
 
-        # Stops sticking to the primary.
-        def self.unstick(namespace, id)
+        def unstick(namespace, id)
           Gitlab::Redis::SharedState.with do |redis|
             redis.del(redis_key_for(namespace, id))
+            redis.del(old_redis_key_for(namespace, id))
           end
         end
 
-        def self.set_write_location_for(namespace, id, location)
+        def set_write_location_for(namespace, id, location)
           Gitlab::Redis::SharedState.with do |redis|
             redis.set(redis_key_for(namespace, id), location, ex: EXPIRATION)
+            redis.set(old_redis_key_for(namespace, id), location, ex: EXPIRATION)
           end
         end
 
-        def self.last_write_location_for(namespace, id)
+        def last_write_location_for(namespace, id)
           Gitlab::Redis::SharedState.with do |redis|
-            redis.get(redis_key_for(namespace, id))
+            redis.get(redis_key_for(namespace, id)) ||
+              redis.get(old_redis_key_for(namespace, id))
           end
         end
 
-        def self.redis_key_for(namespace, id)
-          "database-load-balancing/write-location/#{namespace}/#{id}"
+        def redis_key_for(namespace, id)
+          name = @load_balancer.name
+
+          "database-load-balancing/write-location/#{name}/#{namespace}/#{id}"
         end
 
-        def self.load_balancer
-          LoadBalancing.proxy.load_balancer
+        def old_redis_key_for(namespace, id)
+          "database-load-balancing/write-location/#{namespace}/#{id}"
         end
       end
     end

spec/lib/api/ci/helpers/runner_spec.rb

@@ -15,8 +15,8 @@ RSpec.describe API::Ci::Helpers::Runner do
     it 'handles sticking of a build when a build ID is specified' do
       allow(helper).to receive(:params).and_return(id: build.id)
 
-      expect(Gitlab::Database::LoadBalancing::RackMiddleware)
-        .to receive(:stick_or_unstick)
+      expect(ApplicationRecord.sticking)
+        .to receive(:stick_or_unstick_request)
         .with({}, :build, build.id)
 
       helper.current_job
@@ -25,8 +25,8 @@ RSpec.describe API::Ci::Helpers::Runner do
     it 'does not handle sticking if no build ID was specified' do
       allow(helper).to receive(:params).and_return({})
 
-      expect(Gitlab::Database::LoadBalancing::RackMiddleware)
-        .not_to receive(:stick_or_unstick)
+      expect(ApplicationRecord.sticking)
+        .not_to receive(:stick_or_unstick_request)
 
       helper.current_job
     end
@@ -44,8 +44,8 @@ RSpec.describe API::Ci::Helpers::Runner do
     it 'handles sticking of a runner if a token is specified' do
       allow(helper).to receive(:params).and_return(token: runner.token)
 
-      expect(Gitlab::Database::LoadBalancing::RackMiddleware)
-        .to receive(:stick_or_unstick)
+      expect(ApplicationRecord.sticking)
+        .to receive(:stick_or_unstick_request)
         .with({}, :runner, runner.token)
 
       helper.current_runner
@@ -54,8 +54,8 @@ RSpec.describe API::Ci::Helpers::Runner do
     it 'does not handle sticking if no token was specified' do
       allow(helper).to receive(:params).and_return({})
 
-      expect(Gitlab::Database::LoadBalancing::RackMiddleware)
-        .not_to receive(:stick_or_unstick)
+      expect(ApplicationRecord.sticking)
+        .not_to receive(:stick_or_unstick_request)
 
       helper.current_runner
     end

spec/lib/api/helpers_spec.rb

@@ -35,8 +35,8 @@ RSpec.describe API::Helpers do
     it 'handles sticking when a user could be found' do
       allow_any_instance_of(API::Helpers).to receive(:initial_current_user).and_return(user)
 
-      expect(Gitlab::Database::LoadBalancing::RackMiddleware)
-        .to receive(:stick_or_unstick).with(any_args, :user, 42)
+      expect(ApplicationRecord.sticking)
+        .to receive(:stick_or_unstick_request).with(any_args, :user, 42)
 
       get 'user'
 
@@ -46,8 +46,8 @@ RSpec.describe API::Helpers do
     it 'does not handle sticking if no user could be found' do
       allow_any_instance_of(API::Helpers).to receive(:initial_current_user).and_return(nil)
 
-      expect(Gitlab::Database::LoadBalancing::RackMiddleware)
-        .not_to receive(:stick_or_unstick)
+      expect(ApplicationRecord.sticking)
+        .not_to receive(:stick_or_unstick_request)
 
       get 'user'
 

spec/lib/gitlab/checks/matching_merge_request_spec.rb

@@ -37,10 +37,20 @@ RSpec.describe Gitlab::Checks::MatchingMergeRequest do
 
       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)
+        allow(::ApplicationRecord.sticking)
+          .to receive(:all_caught_up?)
+          .and_return(all_caught_up)
+
+        expect(::ApplicationRecord.sticking)
+          .to receive(:select_valid_host)
+          .with(:project, project.id)
+          .and_call_original
+
+        allow(::ApplicationRecord.sticking)
+          .to receive(:select_caught_up_replicas)
+          .with(:project, project.id)
+          .and_return(all_caught_up)
       end
 
       after do

spec/lib/gitlab/database/consistency_spec.rb

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

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

@@ -5,7 +5,7 @@ require 'spec_helper'
 RSpec.describe Gitlab::Database::LoadBalancing::ActionCableCallbacks, :request_store do
   describe '.wrapper' do
     it 'uses primary and then releases the connection and clears the session' do
-      expect(Gitlab::Database::LoadBalancing).to receive_message_chain(:proxy, :load_balancer, :release_host)
+      expect(Gitlab::Database::LoadBalancing).to receive(:release_hosts)
       expect(Gitlab::Database::LoadBalancing::Session).to receive(:clear_session)
 
       described_class.wrapper.call(
@@ -18,7 +18,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ActionCableCallbacks, :request_s
 
     context 'with an exception' do
       it 'releases the connection and clears the session' do
-        expect(Gitlab::Database::LoadBalancing).to receive_message_chain(:proxy, :load_balancer, :release_host)
+        expect(Gitlab::Database::LoadBalancing).to receive(:release_hosts)
         expect(Gitlab::Database::LoadBalancing::Session).to receive(:clear_session)
 
         expect do

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

-# frozen_string_literal: true
-
-require 'spec_helper'
-
-RSpec.describe Gitlab::Database::LoadBalancing::ActiveRecordProxy do
-  describe '#connection' do
-    it 'returns a connection proxy' do
-      dummy = Class.new do
-        include Gitlab::Database::LoadBalancing::ActiveRecordProxy
-      end
-
-      proxy = double(:proxy)
-
-      expect(Gitlab::Database::LoadBalancing).to receive(:proxy)
-        .and_return(proxy)
-
-      expect(dummy.new.connection).to eq(proxy)
-    end
-
-    it 'returns a connection when no proxy is present' do
-      allow(Gitlab::Database::LoadBalancing).to receive(:proxy).and_return(nil)
-
-      expect(ActiveRecord::Base.connection)
-        .to eq(ActiveRecord::Base.retrieve_connection)
-    end
-  end
-end

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

@@ -108,6 +108,14 @@ RSpec.describe Gitlab::Database::LoadBalancing::Configuration do
   end
 
   describe '#load_balancing_enabled?' do
+    it 'returns false when running inside a Rake task' do
+      config = described_class.new(ActiveRecord::Base, %w[foo bar])
+
+      allow(Gitlab::Runtime).to receive(:rake?).and_return(true)
+
+      expect(config.load_balancing_enabled?).to eq(false)
+    end
+
     it 'returns true when hosts are configured' do
       config = described_class.new(ActiveRecord::Base, %w[foo bar])
 

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

@@ -47,16 +47,27 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
   end
 
   describe '#initialize' do
-    it 'ignores the hosts when the primary_only option is enabled' do
+    it 'ignores the hosts when load balancing is disabled' do
       config = Gitlab::Database::LoadBalancing::Configuration
         .new(ActiveRecord::Base, [db_host])
-      lb = described_class.new(config, primary_only: true)
+
+      allow(config).to receive(:load_balancing_enabled?).and_return(false)
+
+      lb = described_class.new(config)
       hosts = lb.host_list.hosts
 
       expect(hosts.length).to eq(1)
       expect(hosts.first)
         .to be_instance_of(Gitlab::Database::LoadBalancing::PrimaryHost)
     end
+
+    it 'sets the name of the connection that is used' do
+      config =
+        Gitlab::Database::LoadBalancing::Configuration.new(ActiveRecord::Base)
+      lb = described_class.new(config)
+
+      expect(lb.name).to eq(:main)
+    end
   end
 
   describe '#read' do
@@ -140,10 +151,13 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
         .to yield_with_args(ActiveRecord::Base.retrieve_connection)
     end
 
-    it 'uses the primary when the primary_only option is enabled' do
+    it 'uses the primary when load balancing is disabled' do
       config = Gitlab::Database::LoadBalancing::Configuration
         .new(ActiveRecord::Base)
-      lb = described_class.new(config, primary_only: true)
+
+      allow(config).to receive(:load_balancing_enabled?).and_return(false)
+
+      lb = described_class.new(config)
 
       # When no hosts are configured, we don't want to produce any warnings, as
       # they aren't useful/too noisy.

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

@@ -6,12 +6,12 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
   let(:app) { double(:app) }
   let(:middleware) { described_class.new(app) }
   let(:warden_user) { double(:warden, user: double(:user, id: 42)) }
-  let(:single_sticking_object) { Set.new([[:user, 42]]) }
+  let(:single_sticking_object) { Set.new([[ActiveRecord::Base, :user, 42]]) }
   let(:multiple_sticking_objects) do
     Set.new([
-      [:user, 42],
-      [:runner, '123456789'],
-      [:runner, '1234']
+      [ActiveRecord::Base, :user, 42],
+      [ActiveRecord::Base, :runner, '123456789'],
+      [ActiveRecord::Base, :runner, '1234']
     ])
   end
 
@@ -19,42 +19,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
     Gitlab::Database::LoadBalancing::Session.clear_session
   end
 
-  describe '.stick_or_unstick' do
-    it 'sticks or unsticks a single object and updates the Rack environment' do
-      expect(Gitlab::Database::LoadBalancing::Sticking)
-        .to receive(:unstick_or_continue_sticking)
-        .with(:user, 42)
-
-      env = {}
-
-      described_class.stick_or_unstick(env, :user, 42)
-
-      expect(env[described_class::STICK_OBJECT].to_a).to eq([[:user, 42]])
-    end
-
-    it 'sticks or unsticks multiple objects and updates the Rack environment' do
-      expect(Gitlab::Database::LoadBalancing::Sticking)
-        .to receive(:unstick_or_continue_sticking)
-        .with(:user, 42)
-        .ordered
-
-      expect(Gitlab::Database::LoadBalancing::Sticking)
-        .to receive(:unstick_or_continue_sticking)
-        .with(:runner, '123456789')
-        .ordered
-
-      env = {}
-
-      described_class.stick_or_unstick(env, :user, 42)
-      described_class.stick_or_unstick(env, :runner, '123456789')
-
-      expect(env[described_class::STICK_OBJECT].to_a).to eq([
-        [:user, 42],
-        [:runner, '123456789']
-      ])
-    end
-  end
-
   describe '#call' do
     it 'handles a request' do
       env = {}
@@ -77,7 +41,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
 
   describe '#unstick_or_continue_sticking' do
     it 'does not stick if no namespace and identifier could be found' do
-      expect(Gitlab::Database::LoadBalancing::Sticking)
+      expect(ApplicationRecord.sticking)
         .not_to receive(:unstick_or_continue_sticking)
 
       middleware.unstick_or_continue_sticking({})
@@ -86,9 +50,11 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
     it 'sticks to the primary if a warden user is found' do
       env = { 'warden' => warden_user }
 
-      expect(Gitlab::Database::LoadBalancing::Sticking)
-        .to receive(:unstick_or_continue_sticking)
-        .with(:user, 42)
+      Gitlab::Database::LoadBalancing.base_models.each do |model|
+        expect(model.sticking)
+          .to receive(:unstick_or_continue_sticking)
+          .with(:user, 42)
+      end
 
       middleware.unstick_or_continue_sticking(env)
     end
@@ -96,7 +62,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
     it 'sticks to the primary if a sticking namespace and identifier is found' do
       env = { described_class::STICK_OBJECT => single_sticking_object }
 
-      expect(Gitlab::Database::LoadBalancing::Sticking)
+      expect(ApplicationRecord.sticking)
         .to receive(:unstick_or_continue_sticking)
         .with(:user, 42)
 
@@ -106,17 +72,17 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
     it 'sticks to the primary if multiple sticking namespaces and identifiers were found' do
       env = { described_class::STICK_OBJECT => multiple_sticking_objects }
 
-      expect(Gitlab::Database::LoadBalancing::Sticking)
+      expect(ApplicationRecord.sticking)
         .to receive(:unstick_or_continue_sticking)
         .with(:user, 42)
         .ordered
 
-      expect(Gitlab::Database::LoadBalancing::Sticking)
+      expect(ApplicationRecord.sticking)
         .to receive(:unstick_or_continue_sticking)
         .with(:runner, '123456789')
         .ordered
 
-      expect(Gitlab::Database::LoadBalancing::Sticking)
+      expect(ApplicationRecord.sticking)
         .to receive(:unstick_or_continue_sticking)
         .with(:runner, '1234')
         .ordered
@@ -127,7 +93,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
 
   describe '#stick_if_necessary' do
     it 'does not stick to the primary if not necessary' do
-      expect(Gitlab::Database::LoadBalancing::Sticking)
+      expect(ApplicationRecord.sticking)
         .not_to receive(:stick_if_necessary)
 
       middleware.stick_if_necessary({})
@@ -136,9 +102,11 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
     it 'sticks to the primary if a warden user is found' do
       env = { 'warden' => warden_user }
 
-      expect(Gitlab::Database::LoadBalancing::Sticking)
-        .to receive(:stick_if_necessary)
-        .with(:user, 42)
+      Gitlab::Database::LoadBalancing.base_models.each do |model|
+        expect(model.sticking)
+          .to receive(:stick_if_necessary)
+          .with(:user, 42)
+      end
 
       middleware.stick_if_necessary(env)
     end
@@ -146,7 +114,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
     it 'sticks to the primary if a a single sticking object is found' do
       env = { described_class::STICK_OBJECT => single_sticking_object }
 
-      expect(Gitlab::Database::LoadBalancing::Sticking)
+      expect(ApplicationRecord.sticking)
         .to receive(:stick_if_necessary)
         .with(:user, 42)
 
@@ -156,17 +124,17 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
     it 'sticks to the primary if multiple sticking namespaces and identifiers were found' do
       env = { described_class::STICK_OBJECT => multiple_sticking_objects }
 
-      expect(Gitlab::Database::LoadBalancing::Sticking)
+      expect(ApplicationRecord.sticking)
         .to receive(:stick_if_necessary)
         .with(:user, 42)
         .ordered
 
-      expect(Gitlab::Database::LoadBalancing::Sticking)
+      expect(ApplicationRecord.sticking)
         .to receive(:stick_if_necessary)
         .with(:runner, '123456789')
         .ordered
 
-      expect(Gitlab::Database::LoadBalancing::Sticking)
+      expect(ApplicationRecord.sticking)
         .to receive(:stick_if_necessary)
         .with(:runner, '1234')
         .ordered
@@ -177,47 +145,34 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
 
   describe '#clear' do
     it 'clears the currently used host and session' do
-      lb = double(:lb)
       session = spy(:session)
 
-      allow(middleware).to receive(:load_balancer).and_return(lb)
-
-      expect(lb).to receive(:release_host)
-
       stub_const('Gitlab::Database::LoadBalancing::Session', session)
 
+      expect(Gitlab::Database::LoadBalancing).to receive(:release_hosts)
+
       middleware.clear
 
       expect(session).to have_received(:clear_session)
     end
   end
 
-  describe '.load_balancer' do
-    it 'returns a the load balancer' do
-      proxy = double(:proxy)
-
-      expect(Gitlab::Database::LoadBalancing).to receive(:proxy)
-        .and_return(proxy)
-
-      expect(proxy).to receive(:load_balancer)
-
-      middleware.load_balancer
-    end
-  end
-
-  describe '#sticking_namespaces_and_ids' do
+  describe '#sticking_namespaces' do
     context 'using a Warden request' do
       it 'returns the warden user if present' do
         env = { 'warden' => warden_user }
+        ids = Gitlab::Database::LoadBalancing.base_models.map do |model|
+          [model, :user, 42]
+        end
 
-        expect(middleware.sticking_namespaces_and_ids(env)).to eq([[:user, 42]])
+        expect(middleware.sticking_namespaces(env)).to eq(ids)
       end
 
       it 'returns an empty Array if no user was present' do
         warden = double(:warden, user: nil)
         env = { 'warden' => warden }
 
-        expect(middleware.sticking_namespaces_and_ids(env)).to eq([])
+        expect(middleware.sticking_namespaces(env)).to eq([])
       end
     end
 
@@ -225,17 +180,17 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
       it 'returns the sticking object' do
         env = { described_class::STICK_OBJECT => multiple_sticking_objects }
 
-        expect(middleware.sticking_namespaces_and_ids(env)).to eq([
-          [:user, 42],
-          [:runner, '123456789'],
-          [:runner, '1234']
+        expect(middleware.sticking_namespaces(env)).to eq([
+          [ActiveRecord::Base, :user, 42],
+          [ActiveRecord::Base, :runner, '123456789'],
+          [ActiveRecord::Base, :runner, '1234']
         ])
       end
     end
 
     context 'using a regular request' do
       it 'returns an empty Array' do
-        expect(middleware.sticking_namespaces_and_ids({})).to eq([])
+        expect(middleware.sticking_namespaces({})).to eq([])
       end
     end
   end

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

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::LoadBalancing::Setup do
+  describe '#setup' do
+    it 'sets up the load balancer' do
+      setup = described_class.new(ActiveRecord::Base)
+
+      expect(setup).to receive(:disable_prepared_statements)
+      expect(setup).to receive(:setup_load_balancer)
+      expect(setup).to receive(:setup_service_discovery)
+
+      setup.setup
+    end
+  end
+
+  describe '#disable_prepared_statements' do
+    it 'disables prepared statements and reconnects to the database' do
+      config = double(
+        :config,
+        configuration_hash: { host: 'localhost' },
+        env_name: 'test',
+        name: 'main'
+      )
+      model = double(:model, connection_db_config: config)
+
+      expect(ActiveRecord::DatabaseConfigurations::HashConfig)
+        .to receive(:new)
+        .with('test', 'main', { host: 'localhost', prepared_statements: false })
+        .and_call_original
+
+      # HashConfig doesn't implement its own #==, so we can't directly compare
+      # the expected value with a pre-defined one.
+      expect(model)
+        .to receive(:establish_connection)
+        .with(an_instance_of(ActiveRecord::DatabaseConfigurations::HashConfig))
+
+      described_class.new(model).disable_prepared_statements
+    end
+  end
+
+  describe '#setup_load_balancer' do
+    it 'sets up the load balancer' do
+      model = Class.new(ActiveRecord::Base)
+      setup = described_class.new(model)
+      config = Gitlab::Database::LoadBalancing::Configuration.new(model)
+      lb = instance_spy(Gitlab::Database::LoadBalancing::LoadBalancer)
+
+      allow(lb).to receive(:configuration).and_return(config)
+
+      expect(Gitlab::Database::LoadBalancing::LoadBalancer)
+        .to receive(:new)
+        .with(setup.configuration)
+        .and_return(lb)
+
+      setup.setup_load_balancer
+
+      expect(model.connection.load_balancer).to eq(lb)
+      expect(model.sticking)
+        .to be_an_instance_of(Gitlab::Database::LoadBalancing::Sticking)
+    end
+  end
+
+  describe '#setup_service_discovery' do
+    context 'when service discovery is disabled' do
+      it 'does nothing' do
+        expect(Gitlab::Database::LoadBalancing::ServiceDiscovery)
+          .not_to receive(:new)
+
+        described_class.new(ActiveRecord::Base).setup_service_discovery
+      end
+    end
+
+    context 'when service discovery is enabled' do
+      it 'immediately performs service discovery' do
+        model = ActiveRecord::Base
+        setup = described_class.new(model)
+        sv = instance_spy(Gitlab::Database::LoadBalancing::ServiceDiscovery)
+        lb = model.connection.load_balancer
+
+        allow(setup.configuration)
+          .to receive(:service_discovery_enabled?)
+          .and_return(true)
+
+        allow(Gitlab::Database::LoadBalancing::ServiceDiscovery)
+          .to receive(:new)
+          .with(lb, setup.configuration.service_discovery)
+          .and_return(sv)
+
+        expect(sv).to receive(:perform_service_discovery)
+        expect(sv).not_to receive(:start)
+
+        setup.setup_service_discovery
+      end
+
+      it 'starts service discovery if needed' do
+        model = ActiveRecord::Base
+        setup = described_class.new(model, start_service_discovery: true)
+        sv = instance_spy(Gitlab::Database::LoadBalancing::ServiceDiscovery)
+        lb = model.connection.load_balancer
+
+        allow(setup.configuration)
+          .to receive(:service_discovery_enabled?)
+          .and_return(true)
+
+        allow(Gitlab::Database::LoadBalancing::ServiceDiscovery)
+          .to receive(:new)
+          .with(lb, setup.configuration.service_discovery)
+          .and_return(sv)
+
+        expect(sv).to receive(:perform_service_discovery)
+        expect(sv).to receive(:start)
+
+        setup.setup_service_discovery
+      end
+    end
+  end
+end

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

@@ -5,7 +5,6 @@ require 'spec_helper'
 RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
   let(:middleware) { described_class.new }
 
-  let(:load_balancer) { Gitlab::Database::LoadBalancing.proxy.load_balancer }
   let(:worker_class) { 'TestDataConsistencyWorker' }
   let(:job) { { "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e" } }
 
@@ -84,9 +83,15 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
         end
 
         it 'passes database_replica_location' do
-          expected_location = { Gitlab::Database::MAIN_DATABASE_NAME.to_sym => location }
+          expected_location = {}
 
-          expect(load_balancer).to receive_message_chain(:host, "database_replica_location").and_return(location)
+          Gitlab::Database::LoadBalancing.each_load_balancer do |lb|
+            expect(lb.host)
+              .to receive(:database_replica_location)
+              .and_return(location)
+
+            expected_location[lb.name] = location
+          end
 
           run_middleware
 
@@ -102,9 +107,15 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
         end
 
         it 'passes primary write location', :aggregate_failures do
-          expected_location = { Gitlab::Database::MAIN_DATABASE_NAME.to_sym => location }
+          expected_location = {}
 
-          expect(load_balancer).to receive(:primary_write_location).and_return(location)
+          Gitlab::Database::LoadBalancing.each_load_balancer do |lb|
+            expect(lb)
+              .to receive(:primary_write_location)
+              .and_return(location)
+
+            expected_location[lb.name] = location
+          end
 
           run_middleware
 
@@ -136,8 +147,10 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
       let(:job) { { "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", 'wal_locations' => wal_locations } }
 
       before do
-        allow(load_balancer).to receive(:primary_write_location).and_return(new_location)
-        allow(load_balancer).to receive(:database_replica_location).and_return(new_location)
+        Gitlab::Database::LoadBalancing.each_load_balancer do |lb|
+          allow(lb).to receive(:primary_write_location).and_return(new_location)
+          allow(lb).to receive(:database_replica_location).and_return(new_location)
+        end
       end
 
       shared_examples_for 'does not set database location again' do |use_primary|

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

@@ -4,9 +4,6 @@ require 'spec_helper'
 
 RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware, :clean_gitlab_redis_queues do
   let(:middleware) { described_class.new }
-
-  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' } }
 
@@ -15,6 +12,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware, :clean_
     skip_default_enabled_yaml_check
 
     replication_lag!(false)
+    Gitlab::Database::LoadBalancing::Session.clear_session
   end
 
   after do
@@ -66,7 +64,10 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware, :clean_
       let(:wal_locations) { { Gitlab::Database::MAIN_DATABASE_NAME.to_sym => location } }
 
       it 'does not stick to the primary', :aggregate_failures do
-        expect(load_balancer).to receive(:select_up_to_date_host).with(location).and_return(true)
+        expect(ActiveRecord::Base.connection.load_balancer)
+          .to receive(:select_up_to_date_host)
+          .with(location)
+          .and_return(true)
 
         run_middleware do
           expect(Gitlab::Database::LoadBalancing::Session.current.use_primary?).not_to be_truthy
@@ -91,7 +92,12 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware, :clean_
         let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e', 'wal_locations' => wal_locations } }
 
         before do
-          allow(load_balancer).to receive(:select_up_to_date_host).with(location).and_return(true)
+          Gitlab::Database::LoadBalancing.each_load_balancer do |lb|
+            allow(lb)
+              .to receive(:select_up_to_date_host)
+              .with(location)
+              .and_return(true)
+          end
         end
 
         it_behaves_like 'replica is up to date', 'replica'
@@ -101,7 +107,10 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware, :clean_
         let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e', 'dedup_wal_locations' => wal_locations } }
 
         before do
-          allow(load_balancer).to receive(:select_up_to_date_host).with(wal_locations[:main]).and_return(true)
+          allow(ActiveRecord::Base.connection.load_balancer)
+            .to receive(:select_up_to_date_host)
+            .with(wal_locations[:main])
+            .and_return(true)
         end
 
         it_behaves_like 'replica is up to date', 'replica'
@@ -111,7 +120,10 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware, :clean_
         let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e', 'database_write_location' => '0/D525E3A8' } }
 
         before do
-          allow(load_balancer).to receive(:select_up_to_date_host).with('0/D525E3A8').and_return(true)
+          allow(ActiveRecord::Base.connection.load_balancer)
+            .to receive(:select_up_to_date_host)
+            .with('0/D525E3A8')
+            .and_return(true)
         end
 
         it_behaves_like 'replica is up to date', 'replica'
@@ -187,7 +199,9 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware, :clean_
 
       context 'when replica is not up to date' do
         before do
-          allow(load_balancer).to receive(:select_up_to_date_host).and_return(false)
+          Gitlab::Database::LoadBalancing.each_load_balancer do |lb|
+            allow(lb).to receive(:select_up_to_date_host).and_return(false)
+          end
         end
 
         include_examples 'stick to the primary', 'primary'
@@ -195,6 +209,45 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware, :clean_
     end
   end
 
+  describe '#databases_in_sync?' do
+    it 'treats load balancers without WAL entries as in sync' do
+      expect(middleware.send(:databases_in_sync?, {}))
+        .to eq(true)
+    end
+
+    it 'returns true when all load balancers are in sync' do
+      locations = {}
+
+      Gitlab::Database::LoadBalancing.each_load_balancer do |lb|
+        locations[lb.name] = 'foo'
+
+        expect(lb)
+          .to receive(:select_up_to_date_host)
+          .with('foo')
+          .and_return(true)
+      end
+
+      expect(middleware.send(:databases_in_sync?, locations))
+        .to eq(true)
+    end
+
+    it 'returns false when the load balancers are not in sync' do
+      locations = {}
+
+      Gitlab::Database::LoadBalancing.each_load_balancer do |lb|
+        locations[lb.name] = 'foo'
+
+        allow(lb)
+          .to receive(:select_up_to_date_host)
+          .with('foo')
+          .and_return(false)
+      end
+
+      expect(middleware.send(:databases_in_sync?, locations))
+        .to eq(false)
+    end
+  end
+
   def process_job(job)
     Sidekiq::JobRetry.new.local(worker_class, job.to_json, 'default') do
       worker_class.process_job(job)
@@ -208,6 +261,8 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware, :clean_
   end
 
   def replication_lag!(exists)
-    allow(load_balancer).to receive(:select_up_to_date_host).and_return(!exists)
+    Gitlab::Database::LoadBalancing.each_load_balancer do |lb|
+      allow(lb).to receive(:select_up_to_date_host).and_return(!exists)
+    end
   end
 end

spec/lib/gitlab/database/load_balancing_spec.rb

@@ -3,173 +3,48 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::LoadBalancing do
-  describe '.proxy' do
-    it 'returns the connection proxy' do
-      proxy = double(:connection_proxy)
+  describe '.base_models' do
+    it 'returns the models to apply load balancing to' do
+      models = described_class.base_models
 
-      allow(ActiveRecord::Base)
-        .to receive(:load_balancing_proxy)
-        .and_return(proxy)
+      expect(models).to include(ActiveRecord::Base)
 
-      expect(described_class.proxy).to eq(proxy)
-    end
-  end
-
-  describe '.configuration' do
-    it 'returns the configuration for the load balancer' do
-      raw = ActiveRecord::Base.connection_db_config.configuration_hash
-      cfg = described_class.configuration
-
-      # There isn't much to test here as the load balancing settings might not
-      # (and likely aren't) set when running tests.
-      expect(cfg.pool_size).to eq(raw[:pool])
-    end
-  end
-
-  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 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_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_replicas?).to eq(false)
+      if Gitlab::Database.has_config?(:ci)
+        expect(models).to include(Ci::CiDatabaseRecord)
       end
     end
 
-    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)
-
-        allow(described_class.configuration)
-          .to receive(:service_discovery_enabled?)
-          .and_return(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
+    it 'returns the models as a frozen array' do
+      expect(described_class.base_models).to be_frozen
     end
   end
 
-  describe '.configured?' do
-    it 'returns true when hosts are configured' do
-      allow(described_class.configuration)
-        .to receive(:hosts)
-        .and_return(%w[foo])
+  describe '.each_load_balancer' do
+    it 'yields every load balancer to the supplied block' do
+      lbs = []
 
-      expect(described_class.configured?).to eq(true)
-    end
-
-    it 'returns true when service discovery is enabled' do
-      allow(described_class.configuration).to receive(:hosts).and_return([])
-      allow(described_class.configuration)
-        .to receive(:service_discovery_enabled?)
-        .and_return(true)
-
-      expect(described_class.configured?).to eq(true)
-    end
-
-    it 'returns false when neither service discovery nor hosts are configured' 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.configured?).to eq(false)
-    end
-  end
-
-  describe '.start_service_discovery' do
-    it 'does not start if service discovery is disabled' do
-      expect(Gitlab::Database::LoadBalancing::ServiceDiscovery)
-        .not_to receive(:new)
+      described_class.each_load_balancer do |lb|
+        lbs << lb
+      end
 
-      described_class.start_service_discovery
+      expect(lbs.length).to eq(described_class.base_models.length)
     end
 
-    it 'starts service discovery if enabled' do
-      allow(described_class.configuration)
-        .to receive(:service_discovery_enabled?)
-        .and_return(true)
-
-      instance = double(:instance)
-      config = Gitlab::Database::LoadBalancing::Configuration
-        .new(ActiveRecord::Base)
-      lb = Gitlab::Database::LoadBalancing::LoadBalancer.new(config)
-      proxy = Gitlab::Database::LoadBalancing::ConnectionProxy.new(lb)
+    it 'returns an Enumerator when no block is given' do
+      res = described_class.each_load_balancer
 
-      allow(described_class)
-        .to receive(:proxy)
-        .and_return(proxy)
-
-      expect(Gitlab::Database::LoadBalancing::ServiceDiscovery)
-        .to receive(:new)
-        .with(lb, an_instance_of(Hash))
-        .and_return(instance)
-
-      expect(instance)
-        .to receive(:start)
-
-      described_class.start_service_discovery
+      expect(res.next)
+        .to be_an_instance_of(Gitlab::Database::LoadBalancing::LoadBalancer)
     end
   end
 
-  describe '.perform_service_discovery' do
-    it 'does nothing if service discovery is disabled' do
-      expect(Gitlab::Database::LoadBalancing::ServiceDiscovery)
-        .not_to receive(:new)
-
-      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)
+  describe '.release_hosts' do
+    it 'releases the host of every load balancer' do
+      described_class.each_load_balancer do |lb|
+        expect(lb).to receive(:release_host)
+      end
 
-      described_class.perform_service_discovery
+      described_class.release_hosts
     end
   end
 
@@ -227,7 +102,7 @@ RSpec.describe Gitlab::Database::LoadBalancing do
   # - In each test, we listen to the SQL queries (via sql.active_record
   # instrumentation) while triggering real queries from the defined model.
   # - We assert the desinations (replica/primary) of the queries in order.
-  describe 'LoadBalancing integration tests', :db_load_balancing, :delete do
+  describe 'LoadBalancing integration tests', :database_replica, :delete do
     before(:all) do
       ActiveRecord::Schema.define do
         create_table :load_balancing_test, force: true do |t|

spec/lib/gitlab/database_spec.rb

@@ -203,7 +203,7 @@ RSpec.describe Gitlab::Database do
         .to eq('main')
     end
 
-    context 'when replicas are configured', :db_load_balancing do
+    context 'when replicas are configured', :database_replica do
       it 'returns the name for a replica' do
         replica = ActiveRecord::Base.connection.load_balancer.host
 

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

@@ -158,7 +158,7 @@ RSpec.describe Gitlab::ImportExport::Json::StreamingSerializer do
     end
 
     describe 'load balancing' do
-      context 'when feature flag load_balancing_for_export_workers is enabled', :db_load_balancing do
+      context 'when feature flag load_balancing_for_export_workers is enabled' do
         before do
           stub_feature_flags(load_balancing_for_export_workers: true)
         end

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

@@ -195,7 +195,7 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
       with_them do
         let(:payload) { { name: name, sql: sql(sql_query, comments: comments), connection: connection } }
 
-        context 'query using a connection to a replica', :db_load_balancing do
+        context 'query using a connection to a replica' do
           before do
             allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).and_return(:replica)
           end

spec/lib/gitlab/sidekiq_logging/structured_logger_spec.rb

@@ -317,7 +317,7 @@ RSpec.describe Gitlab::SidekiqLogging::StructuredLogger do
         end
       end
 
-      context 'when load balancing is enabled', :db_load_balancing do
+      context 'when load balancing is enabled' do
         let(:db_config_name) do
           ::Gitlab::Database.db_config_name(ApplicationRecord.retrieve_connection)
         end

spec/models/ci/build_spec.rb

@@ -347,10 +347,10 @@ RSpec.describe Ci::Build do
   end
 
   describe '#stick_build_if_status_changed' do
-    it 'sticks the build if the status changed', :db_load_balancing do
+    it 'sticks the build if the status changed' do
       job = create(:ci_build, :pending)
 
-      expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:stick)
+      expect(ApplicationRecord.sticking).to receive(:stick)
         .with(:build, job.id)
 
       job.update!(status: :running)

spec/models/ci/pipeline_spec.rb

@@ -2790,7 +2790,16 @@ 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
+
+        # The number of extra load balancing queries depends on whether or not
+        # we use a load balancer for CI. That in turn depends on the contents of
+        # database.yml, so here we support both cases.
+        extra_load_balancer_queries =
+          if Gitlab::Database.has_config?(:ci)
+            6
+          else
+            3
+          end
 
         expect(control2.count).to eq(control1.count + extra_update_queries + extra_generic_commit_status_validation_queries + extra_load_balancer_queries)
       end

spec/models/ci/runner_spec.rb

@@ -397,7 +397,7 @@ RSpec.describe Ci::Runner do
     it 'sticks the runner to the primary and calls the original method' do
       runner = create(:ci_runner)
 
-      expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:stick)
+      expect(ApplicationRecord.sticking).to receive(:stick)
         .with(:runner, runner.id)
 
       expect(Gitlab::Workhorse).to receive(:set_key_and_notify)

spec/models/project_feature_usage_spec.rb

@@ -133,10 +133,8 @@ RSpec.describe ProjectFeatureUsage, type: :model do
 
       subject { project.feature_usage }
 
-      context 'database load balancing is configured', :db_load_balancing do
+      context 'database load balancing is configured' do
         before do
-          allow(ActiveRecord::Base).to receive(:connection).and_return(::Gitlab::Database::LoadBalancing.proxy)
-
           ::Gitlab::Database::LoadBalancing::Session.clear_session
         end
 

spec/models/project_spec.rb

@@ -3050,7 +3050,7 @@ RSpec.describe Project, factory_default: :keep do
     let(:project) { create(:project) }
 
     it 'marks the location with project ID' do
-      expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:mark_primary_write_location).with(:project, project.id)
+      expect(ApplicationRecord.sticking).to receive(:mark_primary_write_location).with(:project, project.id)
 
       project.mark_primary_write_location
     end

spec/requests/api/graphql/mutations/merge_requests/set_assignees_spec.rb

@@ -115,7 +115,7 @@ RSpec.describe 'Setting assignees of a merge request', :assume_throttled do
   context 'when passing append as true' do
     let(:mode) { Types::MutationOperationModeEnum.enum[:append] }
     let(:input) { { assignee_usernames: [assignee2.username], operation_mode: mode } }
-    let(:db_query_limit) { 21 }
+    let(:db_query_limit) { 22 }
 
     before do
       # In CE, APPEND is a NOOP as you can't have multiple assignees

spec/services/ci/drop_pipeline_service_spec.rb

@@ -50,13 +50,14 @@ RSpec.describe Ci::DropPipelineService do
       end.count
 
       writes_per_build = 2
+      load_balancer_queries = 3
       expected_reads_count = control_count - writes_per_build
 
       create_list(:ci_build, 5, :running, pipeline: cancelable_pipeline)
 
       expect do
         drop_pipeline!(cancelable_pipeline)
-      end.not_to exceed_query_limit(expected_reads_count + (5 * writes_per_build))
+      end.not_to exceed_query_limit(expected_reads_count + (5 * writes_per_build) + load_balancer_queries)
     end
   end
 end

spec/services/ci/register_job_service_spec.rb

@@ -14,7 +14,7 @@ module Ci
     let!(:pending_job) { create(:ci_build, :pending, :queued, pipeline: pipeline) }
 
     describe '#execute' do
-      context 'checks database loadbalancing stickiness', :db_load_balancing do
+      context 'checks database loadbalancing stickiness' do
         subject { described_class.new(shared_runner).execute }
 
         before do
@@ -22,14 +22,14 @@ module Ci
         end
 
         it 'result is valid if replica did caught-up' do
-          expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:all_caught_up?)
+          expect(ApplicationRecord.sticking).to receive(:all_caught_up?)
             .with(:runner, shared_runner.id) { true }
 
           expect(subject).to be_valid
         end
 
         it 'result is invalid if replica did not caught-up' do
-          expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:all_caught_up?)
+          expect(ApplicationRecord.sticking).to receive(:all_caught_up?)
             .with(:runner, shared_runner.id) { false }
 
           expect(subject).not_to be_valid

spec/services/user_project_access_changed_service_spec.rb

@@ -53,7 +53,7 @@ RSpec.describe UserProjectAccessChangedService do
     end
 
     it 'sticks all the updated users and returns the original result', :aggregate_failures do
-      expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:bulk_stick).with(:user, [1, 2])
+      expect(ApplicationRecord.sticking).to receive(:bulk_stick).with(:user, [1, 2])
 
       expect(service.execute).to eq(10)
     end

spec/services/users/activity_service_spec.rb

@@ -91,9 +91,9 @@ RSpec.describe Users::ActivityService do
     context 'when last activity is in the past' do
       let(:user) { create(:user, last_activity_on: Date.today - 1.week) }
 
-      context 'database load balancing is configured', :db_load_balancing do
+      context 'database load balancing is configured' do
         before do
-          allow(ActiveRecord::Base).to receive(:connection).and_return(::Gitlab::Database::LoadBalancing.proxy)
+          ::Gitlab::Database::LoadBalancing::Session.clear_session
         end
 
         let(:service) do

spec/support/database_load_balancing.rb

 # frozen_string_literal: true
 
 RSpec.configure do |config|
-  config.before(:each, :db_load_balancing) do
-    config = Gitlab::Database::LoadBalancing::Configuration
-      .new(ActiveRecord::Base, [Gitlab::Database.main.config['host']])
-    lb = ::Gitlab::Database::LoadBalancing::LoadBalancer.new(config)
-    proxy = ::Gitlab::Database::LoadBalancing::ConnectionProxy.new(lb)
+  config.around(:each, :database_replica) do |example|
+    old_proxies = []
 
-    allow(ActiveRecord::Base).to receive(:load_balancing_proxy).and_return(proxy)
+    Gitlab::Database::LoadBalancing.base_models.each do |model|
+      config = Gitlab::Database::LoadBalancing::Configuration
+        .new(model, [model.connection_db_config.configuration_hash[:host]])
+      lb = Gitlab::Database::LoadBalancing::LoadBalancer.new(config)
 
-    ::Gitlab::Database::LoadBalancing::Session.clear_session
+      old_proxies << [model, model.connection]
+
+      model.connection =
+        Gitlab::Database::LoadBalancing::ConnectionProxy.new(lb)
+    end
+
+    Gitlab::Database::LoadBalancing::Session.clear_session
     redis_shared_state_cleanup!
-  end
 
-  config.after(:each, :db_load_balancing) do
-    ::Gitlab::Database::LoadBalancing::Session.clear_session
+    example.run
+
+    Gitlab::Database::LoadBalancing::Session.clear_session
     redis_shared_state_cleanup!
+
+    old_proxies.each do |(model, proxy)|
+      model.connection = proxy
+    end
   end
 end

spec/support/shared_examples/lib/gitlab/ci/ci_trace_shared_examples.rb

@@ -35,8 +35,8 @@ RSpec.shared_examples 'common trace features' do
         stub_feature_flags(gitlab_ci_archived_trace_consistent_reads: trace.job.project)
       end
 
-      it 'calls ::Gitlab::Database::LoadBalancing::Sticking.unstick_or_continue_sticking' do
-        expect(::Gitlab::Database::LoadBalancing::Sticking).to receive(:unstick_or_continue_sticking)
+      it 'calls ::ApplicationRecord.sticking.unstick_or_continue_sticking' do
+        expect(::ApplicationRecord.sticking).to receive(:unstick_or_continue_sticking)
           .with(described_class::LOAD_BALANCING_STICKING_NAMESPACE, trace.job.id)
           .and_call_original
 
@@ -49,8 +49,8 @@ RSpec.shared_examples 'common trace features' do
         stub_feature_flags(gitlab_ci_archived_trace_consistent_reads: false)
       end
 
-      it 'does not call ::Gitlab::Database::LoadBalancing::Sticking.unstick_or_continue_sticking' do
-        expect(::Gitlab::Database::LoadBalancing::Sticking).not_to receive(:unstick_or_continue_sticking)
+      it 'does not call ::ApplicationRecord.sticking.unstick_or_continue_sticking' do
+        expect(::ApplicationRecord.sticking).not_to receive(:unstick_or_continue_sticking)
 
         trace.read { |stream| stream }
       end
@@ -305,8 +305,8 @@ RSpec.shared_examples 'common trace features' do
           stub_feature_flags(gitlab_ci_archived_trace_consistent_reads: trace.job.project)
         end
 
-        it 'calls ::Gitlab::Database::LoadBalancing::Sticking.stick' do
-          expect(::Gitlab::Database::LoadBalancing::Sticking).to receive(:stick)
+        it 'calls ::ApplicationRecord.sticking.stick' do
+          expect(::ApplicationRecord.sticking).to receive(:stick)
             .with(described_class::LOAD_BALANCING_STICKING_NAMESPACE, trace.job.id)
             .and_call_original
 
@@ -319,8 +319,8 @@ RSpec.shared_examples 'common trace features' do
           stub_feature_flags(gitlab_ci_archived_trace_consistent_reads: false)
         end
 
-        it 'does not call ::Gitlab::Database::LoadBalancing::Sticking.stick' do
-          expect(::Gitlab::Database::LoadBalancing::Sticking).not_to receive(:stick)
+        it 'does not call ::ApplicationRecord.sticking.stick' do
+          expect(::ApplicationRecord.sticking).not_to receive(:stick)
 
           subject
         end

spec/workers/authorized_project_update/user_refresh_from_replica_worker_spec.rb

@@ -44,7 +44,7 @@ RSpec.describe AuthorizedProjectUpdate::UserRefreshFromReplicaWorker do
       end
     end
 
-    context 'with load balancing enabled', :db_load_balancing do
+    context 'with load balancing enabled' do
       it 'reads from the replica database' do
         expect(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_replicas_for_read_queries).and_call_original
 

spec/workers/container_expiration_policy_worker_spec.rb

@@ -156,7 +156,7 @@ RSpec.describe ContainerExpirationPolicyWorker do
         subject
       end
 
-      context 'with load balancing enabled', :db_load_balancing do
+      context 'with load balancing enabled' do
         it 'reads the counts from the replica' do
           expect(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_replicas_for_read_queries).and_call_original