Add support for load balancing multiple databases

This adds support for using the database load balancer with multiple
databases. Load balancing is applied to two classes:

- ActiveRecord::Base
- Ci::CiDatabaseRecord

Each class has its own load balancer, configuration, service discovery,
etc. Load balancing for the CI class is only enabled when a CI
configuration exists, as it can reuse the main load balancer when
there's no dedicated CI database.

Sticking technically supports multiple databases, but in practise we
apply the same sticking rules to all databases. This is due to how
LoadBalancing::Session is used: there is only one instance per
request/Sidekiq job, and it's not aware of what database connections did
what. This means that a write to database A will result in GitLab
sticking to the primaries of _all_ databases. The choice for this is
simple: it requires fewer code changes, and allows us to introduce
multiple database support in smaller increments.

One change we made to sticking is to turn the Sticking module into a
class, and attach an instance to every base module that has its own load
balancer. This makes it easier to apply sticking on a per-database level
in the future, without having to type
`Gitlab::Database::LoadBalancing::Sticking...` every time.

Sticking also supports reading and writing of data using the old Redis
key names. This ensures sticking continues to work during a deployment,
as during this window we'll run two different versions in production.
Once the code has been deployed to GitLab.com and has been confirmed to
work, we'll remove support for reading/writing the old keys.

Sidekiq also supports load balancing multiple databases. If a load
balancer/database doesn't have any WAL data in the Sidekiq job, we treat
the database as being in sync. This way we can support Sidekiq jobs
using both the old and new load balancing data.

See https://gitlab.com/gitlab-org/gitlab/-/issues/331776 for more
details.

Changelog: added

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/sticking_spec.rb

@@ -3,19 +3,60 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
+  let(:sticking) do
+    described_class.new(ActiveRecord::Base.connection.load_balancer)
+  end
+
   after do
     Gitlab::Database::LoadBalancing::Session.clear_session
   end
 
-  describe '.stick_if_necessary' do
+  describe '#stick_or_unstick_request' do
+    it 'sticks or unsticks a single object and updates the Rack environment' do
+      expect(sticking)
+        .to receive(:unstick_or_continue_sticking)
+        .with(:user, 42)
+
+      env = {}
+
+      sticking.stick_or_unstick_request(env, :user, 42)
+
+      expect(env[Gitlab::Database::LoadBalancing::RackMiddleware::STICK_OBJECT].to_a)
+        .to eq([[ActiveRecord::Base, :user, 42]])
+    end
+
+    it 'sticks or unsticks multiple objects and updates the Rack environment' do
+      expect(sticking)
+        .to receive(:unstick_or_continue_sticking)
+        .with(:user, 42)
+        .ordered
+
+      expect(sticking)
+        .to receive(:unstick_or_continue_sticking)
+        .with(:runner, '123456789')
+        .ordered
+
+      env = {}
+
+      sticking.stick_or_unstick_request(env, :user, 42)
+      sticking.stick_or_unstick_request(env, :runner, '123456789')
+
+      expect(env[Gitlab::Database::LoadBalancing::RackMiddleware::STICK_OBJECT].to_a).to eq([
+        [ActiveRecord::Base, :user, 42],
+        [ActiveRecord::Base, :runner, '123456789']
+      ])
+    end
+  end
+
+  describe '#stick_if_necessary' do
     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(sticking).not_to receive(:stick)
 
-      described_class.stick_if_necessary(:user, 42)
+      sticking.stick_if_necessary(:user, 42)
     end
 
     it 'sticks to the primary if a write was performed' do
@@ -23,20 +64,21 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
         .to receive(:performed_write?)
         .and_return(true)
 
-      expect(described_class).to receive(:stick).with(:user, 42)
+      expect(sticking)
+        .to receive(:stick)
+        .with(:user, 42)
 
-      described_class.stick_if_necessary(:user, 42)
+      sticking.stick_if_necessary(:user, 42)
     end
   end
 
-  describe '.all_caught_up?' do
-    let(:lb) { double(:lb) }
+  describe '#all_caught_up?' do
+    let(:lb) { ActiveRecord::Base.connection.load_balancer }
     let(:last_write_location) { 'foo' }
 
     before do
-      allow(described_class).to receive(:load_balancer).and_return(lb)
-
-      allow(described_class).to receive(:last_write_location_for)
+      allow(sticking)
+        .to receive(:last_write_location_for)
         .with(:user, 42)
         .and_return(last_write_location)
     end
@@ -45,13 +87,9 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
       let(:last_write_location) { nil }
 
       it 'returns true' do
-        allow(described_class).to receive(:last_write_location_for)
-          .with(:user, 42)
-          .and_return(nil)
-
         expect(lb).not_to receive(:select_up_to_date_host)
 
-        expect(described_class.all_caught_up?(:user, 42)).to eq(true)
+        expect(sticking.all_caught_up?(:user, 42)).to eq(true)
       end
     end
 
@@ -61,9 +99,11 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
       end
 
       it 'returns true, and unsticks' do
-        expect(described_class).to receive(:unstick).with(:user, 42)
+        expect(sticking)
+          .to receive(:unstick)
+          .with(:user, 42)
 
-        expect(described_class.all_caught_up?(:user, 42)).to eq(true)
+        expect(sticking.all_caught_up?(:user, 42)).to eq(true)
       end
 
       it 'notifies with the proper event payload' do
@@ -72,7 +112,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
           .with('caught_up_replica_pick.load_balancing', { result: true })
           .and_call_original
 
-        described_class.all_caught_up?(:user, 42)
+        sticking.all_caught_up?(:user, 42)
       end
     end
 
@@ -82,7 +122,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
       end
 
       it 'returns false' do
-        expect(described_class.all_caught_up?(:user, 42)).to eq(false)
+        expect(sticking.all_caught_up?(:user, 42)).to eq(false)
       end
 
       it 'notifies with the proper event payload' do
@@ -91,42 +131,43 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
           .with('caught_up_replica_pick.load_balancing', { result: false })
           .and_call_original
 
-        described_class.all_caught_up?(:user, 42)
+        sticking.all_caught_up?(:user, 42)
       end
     end
   end
 
-  describe '.unstick_or_continue_sticking' do
-    let(:lb) { double(:lb) }
-
-    before do
-      allow(described_class).to receive(:load_balancer).and_return(lb)
-    end
+  describe '#unstick_or_continue_sticking' do
+    let(:lb) { ActiveRecord::Base.connection.load_balancer }
 
     it 'simply returns if no write location could be found' do
-      allow(described_class).to receive(:last_write_location_for)
+      allow(sticking)
+        .to receive(:last_write_location_for)
         .with(:user, 42)
         .and_return(nil)
 
       expect(lb).not_to receive(:select_up_to_date_host)
 
-      described_class.unstick_or_continue_sticking(:user, 42)
+      sticking.unstick_or_continue_sticking(:user, 42)
     end
 
     it 'unsticks if all secondaries have caught up' do
-      allow(described_class).to receive(:last_write_location_for)
+      allow(sticking)
+        .to receive(:last_write_location_for)
         .with(:user, 42)
         .and_return('foo')
 
       allow(lb).to receive(:select_up_to_date_host).with('foo').and_return(true)
 
-      expect(described_class).to receive(:unstick).with(:user, 42)
+      expect(sticking)
+        .to receive(:unstick)
+        .with(:user, 42)
 
-      described_class.unstick_or_continue_sticking(:user, 42)
+      sticking.unstick_or_continue_sticking(:user, 42)
     end
 
     it 'continues using the primary if the secondaries have not yet caught up' do
-      allow(described_class).to receive(:last_write_location_for)
+      allow(sticking)
+        .to receive(:last_write_location_for)
         .with(:user, 42)
         .and_return('foo')
 
@@ -135,21 +176,22 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
       expect(Gitlab::Database::LoadBalancing::Session.current)
         .to receive(:use_primary!)
 
-      described_class.unstick_or_continue_sticking(:user, 42)
+      sticking.unstick_or_continue_sticking(:user, 42)
     end
   end
 
   RSpec.shared_examples 'sticking' do
     before do
-      lb = double(:lb, primary_write_location: 'foo')
-
-      allow(described_class).to receive(:load_balancer).and_return(lb)
+      allow(ActiveRecord::Base.connection.load_balancer)
+        .to receive(:primary_write_location)
+        .and_return('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')
+        expect(sticking)
+          .to receive(:set_write_location_for)
+          .with(:user, id, 'foo')
       end
 
       expect(Gitlab::Database::LoadBalancing::Session.current)
@@ -159,99 +201,106 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
     end
   end
 
-  describe '.stick' do
+  describe '#stick' do
     it_behaves_like 'sticking' do
       let(:ids) { [42] }
-      subject { described_class.stick(:user, ids.first) }
+      subject { sticking.stick(:user, ids.first) }
     end
   end
 
-  describe '.bulk_stick' do
+  describe '#bulk_stick' do
     it_behaves_like 'sticking' do
       let(:ids) { [42, 43] }
-      subject { described_class.bulk_stick(:user, ids) }
+      subject { sticking.bulk_stick(:user, ids) }
     end
   end
 
-  describe '.mark_primary_write_location' do
+  describe '#mark_primary_write_location' do
     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)
+      allow(ActiveRecord::Base.connection.load_balancer)
+        .to receive(:primary_write_location)
+        .and_return('foo')
 
-      expect(described_class).to receive(:set_write_location_for)
+      expect(sticking)
+        .to receive(:set_write_location_for)
         .with(:user, 42, 'foo')
 
-      described_class.mark_primary_write_location(:user, 42)
+      sticking.mark_primary_write_location(:user, 42)
     end
   end
 
-  describe '.unstick' do
+  describe '#unstick' do
     it 'removes the sticking data from Redis' do
-      described_class.set_write_location_for(:user, 4, 'foo')
-      described_class.unstick(:user, 4)
+      sticking.set_write_location_for(:user, 4, 'foo')
+      sticking.unstick(:user, 4)
 
-      expect(described_class.last_write_location_for(:user, 4)).to be_nil
+      expect(sticking.last_write_location_for(:user, 4)).to be_nil
     end
-  end
-
-  describe '.last_write_location_for' do
-    it 'returns the last WAL write location for a user' do
-      described_class.set_write_location_for(:user, 4, 'foo')
 
-      expect(described_class.last_write_location_for(:user, 4)).to eq('foo')
-    end
-  end
+    it 'removes the old key' do
+      Gitlab::Redis::SharedState.with do |redis|
+        redis.set(sticking.send(:old_redis_key_for, :user, 4), 'foo', ex: 30)
+      end
 
-  describe '.redis_key_for' do
-    it 'returns a String' do
-      expect(described_class.redis_key_for(:user, 42))
-        .to eq('database-load-balancing/write-location/user/42')
+      sticking.unstick(:user, 4)
+      expect(sticking.last_write_location_for(:user, 4)).to be_nil
     end
   end
 
-  describe '.load_balancer' do
-    it 'returns a the load balancer' do
-      proxy = double(:proxy)
+  describe '#last_write_location_for' do
+    it 'returns the last WAL write location for a user' do
+      sticking.set_write_location_for(:user, 4, 'foo')
 
-      expect(Gitlab::Database::LoadBalancing).to receive(:proxy)
-        .and_return(proxy)
+      expect(sticking.last_write_location_for(:user, 4)).to eq('foo')
+    end
 
-      expect(proxy).to receive(:load_balancer)
+    it 'falls back to reading the old key' do
+      Gitlab::Redis::SharedState.with do |redis|
+        redis.set(sticking.send(:old_redis_key_for, :user, 4), 'foo', ex: 30)
+      end
 
-      described_class.load_balancer
+      expect(sticking.last_write_location_for(:user, 4)).to eq('foo')
     end
   end
 
-  describe '.select_caught_up_replicas' do
-    let(:lb) { double(:lb) }
-
-    before do
-      allow(described_class).to receive(:load_balancer).and_return(lb)
+  describe '#redis_key_for' do
+    it 'returns a String' do
+      expect(sticking.redis_key_for(:user, 42))
+        .to eq('database-load-balancing/write-location/main/user/42')
     end
+  end
+
+  describe '#select_caught_up_replicas' do
+    let(:lb) { ActiveRecord::Base.connection.load_balancer }
 
     context 'with no write location' do
       before do
-        allow(described_class).to receive(:last_write_location_for)
-          .with(:project, 42).and_return(nil)
+        allow(sticking)
+          .to receive(:last_write_location_for)
+          .with(:project, 42)
+          .and_return(nil)
       end
 
       it 'returns false and does not try to find caught up hosts' do
         expect(lb).not_to receive(:select_up_to_date_host)
-        expect(described_class.select_caught_up_replicas(:project, 42)).to be false
+        expect(sticking.select_caught_up_replicas(:project, 42)).to be false
       end
     end
 
     context 'with write location' do
       before do
-        allow(described_class).to receive(:last_write_location_for)
-          .with(:project, 42).and_return('foo')
+        allow(sticking)
+          .to receive(:last_write_location_for)
+          .with(:project, 42)
+          .and_return('foo')
       end
 
       it 'returns true, selects hosts, and unsticks if any secondary has caught up' do
         expect(lb).to receive(:select_up_to_date_host).and_return(true)
-        expect(described_class).to receive(:unstick).with(:project, 42)
-        expect(described_class.select_caught_up_replicas(:project, 42)).to be true
+        expect(sticking)
+          .to receive(:unstick)
+          .with(:project, 42)
+        expect(sticking.select_caught_up_replicas(:project, 42)).to be true
       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