From 67d896a2d1c7ae633f6021af995efc8719d1bfb9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kamil=20Trzci=C5=84ski?= <ayufan@ayufan.eu>
Date: Wed, 22 Nov 2023 10:49:50 +0100
Subject: [PATCH 1/8] Move `Gitlab::Database::LoadBalancing` to Gem

---
 Gemfile                                       |  1 +
 Gemfile.lock                                  |  9 ++++
 config/initializers/load_balancing.rb         | 36 ++++++++++++++
 .../.rubocop.yml                              | 47 +++++++++++++++++++
 .../.rubocop_todo.yml                         | 41 ++++++++++++++++
 gems/gitlab-database-load_balancing/README.md |  2 +-
 .../lib}/gitlab/database/load_balancing.rb    | 35 ++++++++++++--
 .../load_balancing/action_cable_callbacks.rb  |  0
 .../database/load_balancing/callbacks.rb      | 43 +++++++++++++++++
 .../database/load_balancing/configuration.rb  | 12 +++--
 .../load_balancing/connection_proxy.rb        |  2 +
 .../gitlab/database/load_balancing/host.rb    | 13 ++---
 .../database/load_balancing/host_list.rb      |  3 +-
 .../database/load_balancing/load_balancer.rb  | 15 +++---
 .../database/load_balancing/primary_host.rb   |  2 +-
 .../database/load_balancing/resolver.rb       |  0
 .../load_balancing/service_discovery.rb       | 21 ++++-----
 .../service_discovery/sampler.rb              |  4 +-
 .../gitlab/database/load_balancing/session.rb |  8 ++--
 .../gitlab/database/load_balancing/setup.rb   |  4 +-
 .../database/load_balancing/srv_resolver.rb   |  7 ++-
 .../database/load_balancing/sticking.rb       | 16 ++-----
 .../spec}/fixtures/dns/a_rr.json              |  0
 .../a_with_aaaa_rr_in_additional_section.json |  0
 .../spec}/fixtures/dns/aaaa_rr.json           |  0
 .../srv_with_a_rr_in_additional_section.json  |  0
 .../action_cable_callbacks_spec.rb            |  4 +-
 .../load_balancing/configuration_spec.rb      |  8 ++--
 .../load_balancing/connection_proxy_spec.rb   |  8 ++--
 .../database/load_balancing/host_list_spec.rb |  6 ++-
 .../database/load_balancing/host_spec.rb      | 16 +++----
 .../load_balancing/load_balancer_spec.rb      | 10 ++--
 .../load_balancing/primary_host_spec.rb       |  2 +-
 .../database/load_balancing/resolver_spec.rb  |  0
 .../service_discovery/sampler_spec.rb         |  8 ++--
 .../load_balancing/service_discovery_spec.rb  | 16 +++----
 .../database/load_balancing/session_spec.rb   |  0
 .../database/load_balancing/setup_spec.rb     |  4 +-
 .../load_balancing/srv_resolver_spec.rb       |  5 +-
 .../database/load_balancing/sticking_spec.rb  | 20 +++-----
 .../transaction_leaking_spec.rb               | 14 +++---
 .../spec/spec_helper.rb                       | 14 ++++++
 .../spec/support/application_record.rb        | 11 +++++
 .../spec/support/database_replica.rb          | 28 +++++++++++
 lib/gitlab/database.rb                        | 17 ++-----
 45 files changed, 376 insertions(+), 136 deletions(-)
 create mode 100644 gems/gitlab-database-load_balancing/.rubocop_todo.yml
 rename {lib => gems/gitlab-database-load_balancing/lib}/gitlab/database/load_balancing.rb (59%)
 rename {lib => gems/gitlab-database-load_balancing/lib}/gitlab/database/load_balancing/action_cable_callbacks.rb (100%)
 create mode 100644 gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/callbacks.rb
 rename {lib => gems/gitlab-database-load_balancing/lib}/gitlab/database/load_balancing/configuration.rb (90%)
 rename {lib => gems/gitlab-database-load_balancing/lib}/gitlab/database/load_balancing/connection_proxy.rb (99%)
 rename {lib => gems/gitlab-database-load_balancing/lib}/gitlab/database/load_balancing/host.rb (95%)
 rename {lib => gems/gitlab-database-load_balancing/lib}/gitlab/database/load_balancing/host_list.rb (92%)
 rename {lib => gems/gitlab-database-load_balancing/lib}/gitlab/database/load_balancing/load_balancer.rb (96%)
 rename {lib => gems/gitlab-database-load_balancing/lib}/gitlab/database/load_balancing/primary_host.rb (96%)
 rename {lib => gems/gitlab-database-load_balancing/lib}/gitlab/database/load_balancing/resolver.rb (100%)
 rename {lib => gems/gitlab-database-load_balancing/lib}/gitlab/database/load_balancing/service_discovery.rb (94%)
 rename {lib => gems/gitlab-database-load_balancing/lib}/gitlab/database/load_balancing/service_discovery/sampler.rb (93%)
 rename {lib => gems/gitlab-database-load_balancing/lib}/gitlab/database/load_balancing/session.rb (95%)
 rename {lib => gems/gitlab-database-load_balancing/lib}/gitlab/database/load_balancing/setup.rb (94%)
 rename {lib => gems/gitlab-database-load_balancing/lib}/gitlab/database/load_balancing/srv_resolver.rb (84%)
 rename {lib => gems/gitlab-database-load_balancing/lib}/gitlab/database/load_balancing/sticking.rb (89%)
 rename {spec => gems/gitlab-database-load_balancing/spec}/fixtures/dns/a_rr.json (100%)
 rename {spec => gems/gitlab-database-load_balancing/spec}/fixtures/dns/a_with_aaaa_rr_in_additional_section.json (100%)
 rename {spec => gems/gitlab-database-load_balancing/spec}/fixtures/dns/aaaa_rr.json (100%)
 rename {spec => gems/gitlab-database-load_balancing/spec}/fixtures/dns/srv_with_a_rr_in_additional_section.json (100%)
 rename {spec/lib => gems/gitlab-database-load_balancing/spec}/gitlab/database/load_balancing/action_cable_callbacks_spec.rb (85%)
 rename {spec/lib => gems/gitlab-database-load_balancing/spec}/gitlab/database/load_balancing/configuration_spec.rb (94%)
 rename {spec/lib => gems/gitlab-database-load_balancing/spec}/gitlab/database/load_balancing/connection_proxy_spec.rb (97%)
 rename {spec/lib => gems/gitlab-database-load_balancing/spec}/gitlab/database/load_balancing/host_list_spec.rb (93%)
 rename {spec/lib => gems/gitlab-database-load_balancing/spec}/gitlab/database/load_balancing/host_spec.rb (96%)
 rename {spec/lib => gems/gitlab-database-load_balancing/spec}/gitlab/database/load_balancing/load_balancer_spec.rb (98%)
 rename {spec/lib => gems/gitlab-database-load_balancing/spec}/gitlab/database/load_balancing/primary_host_spec.rb (96%)
 rename {spec/lib => gems/gitlab-database-load_balancing/spec}/gitlab/database/load_balancing/resolver_spec.rb (100%)
 rename {spec/lib => gems/gitlab-database-load_balancing/spec}/gitlab/database/load_balancing/service_discovery/sampler_spec.rb (88%)
 rename {spec/lib => gems/gitlab-database-load_balancing/spec}/gitlab/database/load_balancing/service_discovery_spec.rb (96%)
 rename {spec/lib => gems/gitlab-database-load_balancing/spec}/gitlab/database/load_balancing/session_spec.rb (100%)
 rename {spec/lib => gems/gitlab-database-load_balancing/spec}/gitlab/database/load_balancing/setup_spec.rb (97%)
 rename {spec/lib => gems/gitlab-database-load_balancing/spec}/gitlab/database/load_balancing/srv_resolver_spec.rb (91%)
 rename {spec/lib => gems/gitlab-database-load_balancing/spec}/gitlab/database/load_balancing/sticking_spec.rb (87%)
 rename {spec/lib => gems/gitlab-database-load_balancing/spec}/gitlab/database/load_balancing/transaction_leaking_spec.rb (87%)
 create mode 100644 gems/gitlab-database-load_balancing/spec/support/application_record.rb
 create mode 100644 gems/gitlab-database-load_balancing/spec/support/database_replica.rb

diff --git a/Gemfile b/Gemfile
index 40194cd65b389f..b4ac85e747b2be 100644
--- a/Gemfile
+++ b/Gemfile
@@ -41,6 +41,7 @@ gem 'gitlab-safe_request_store', path: 'gems/gitlab-safe_request_store' # ruboco
 group :monorepo do
   gem 'gitlab-utils', path: 'gems/gitlab-utils' # rubocop:todo Gemfile/MissingFeatureCategory
   gem 'gitlab-backup-cli', path: 'gems/gitlab-backup-cli', feature_category: :backup_restore
+  gem 'gitlab-database-load_balancing', path: 'gems/gitlab-database-load_balancing' # rubocop:todo Gemfile/MissingFeatureCategory
 end
 
 gem 'gitlab-secret_detection', path: 'gems/gitlab-secret_detection', feature_category: :secret_detection
diff --git a/Gemfile.lock b/Gemfile.lock
index 9d40284e424ff9..70ae3c079711ea 100644
--- a/Gemfile.lock
+++ b/Gemfile.lock
@@ -29,6 +29,14 @@ PATH
     gitlab-backup-cli (0.0.1)
       thor (~> 1.3)
 
+PATH
+  remote: gems/gitlab-database-load_balancing
+  specs:
+    gitlab-database-load_balancing (0.1.0)
+      gitlab-net-dns (~> 0.9.2)
+      pg (~> 1.5.4)
+      rails (~> 7.0.8)
+
 PATH
   remote: gems/gitlab-http
   specs:
@@ -1866,6 +1874,7 @@ DEPENDENCIES
   gitlab-backup-cli!
   gitlab-chronic (~> 0.10.5)
   gitlab-dangerfiles (~> 4.6.0)
+  gitlab-database-load_balancing!
   gitlab-experiment (~> 0.8.0)
   gitlab-fog-azure-rm (~> 1.8.0)
   gitlab-http!
diff --git a/config/initializers/load_balancing.rb b/config/initializers/load_balancing.rb
index aa6216025732ce..1ebd33541320fe 100644
--- a/config/initializers/load_balancing.rb
+++ b/config/initializers/load_balancing.rb
@@ -13,6 +13,42 @@
   end
 end
 
+# Configure load balancing
+Gitlab::Database::LoadBalancing.enabled = !Gitlab::Runtime.rake?
+Gitlab::Database::LoadBalancing.default_pool_size = Gitlab::Database.default_pool_size
+
+Gitlab::Cluster::LifecycleEvents.on_worker_start do
+  Gitlab::Database::LoadBalancing.enabled = !Gitlab::Runtime.rake?
+  Gitlab::Database::LoadBalancing.default_pool_size = Gitlab::Database.default_pool_size
+end
+
+Gitlab::Database::LoadBalancing.base_models =
+  ::Gitlab::Database.database_base_models_using_load_balancing.values.freeze
+
+# Configure callbacks
+Gitlab::Database::LoadBalancing::Callbacks.configure! do |callbacks|
+  hosts_gauge = Gitlab::Metrics.gauge(:db_load_balancing_hosts, 'Current number of load balancing hosts')
+
+  callbacks.logger_proc = -> do
+    Gitlab::Database::LoadBalancing::Logger.build
+  end
+  callbacks.metrics_host_gauge_proc = ->(labels, value) do
+    hosts_gauge.set(labels, value)
+  end
+  callbacks.track_exception_proc = ->(exception) do
+    Gitlab::ErrorTracking.track_exception(exception)
+  end
+  callbacks.set_wal_for_proc = ->(key, value, ex:) do
+    Gitlab::Redis::DbLoadBalancing.with { |redis| redis.set(key, value, ex: ex) }
+  end
+  callbacks.get_wal_for_proc = ->(key) do
+    Gitlab::Redis::DbLoadBalancing.with { |redis| redis.get(key) }
+  end
+  callbacks.del_wal_for_proc = ->(key) do
+    Gitlab::Redis::DbLoadBalancing.with { |redis| redis.del(key) }
+  end
+end
+
 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
diff --git a/gems/gitlab-database-load_balancing/.rubocop.yml b/gems/gitlab-database-load_balancing/.rubocop.yml
index 583fa8227eec03..8340cf13f69c26 100644
--- a/gems/gitlab-database-load_balancing/.rubocop.yml
+++ b/gems/gitlab-database-load_balancing/.rubocop.yml
@@ -1,6 +1,53 @@
 inherit_from:
+  - .rubocop_todo.yml
   - ../config/rubocop.yml
 
 Gemfile/MissingFeatureCategory:
   Exclude:
     - 'Gemfile'
+
+Database/MultipleDatabases:
+  Exclude:
+    - 'lib/gitlab/database/load_balancing/load_balancer.rb'
+    - 'spec/gitlab/database/load_balancing/host_list_spec.rb'
+    - 'spec/gitlab/database/load_balancing/host_spec.rb'
+    - 'spec/gitlab/database/load_balancing/load_balancer_spec.rb'
+    - 'spec/gitlab/database/load_balancing/setup_spec.rb'
+    - 'spec/gitlab/database/load_balancing/sticking_spec.rb'
+
+Graphql/ResolverType:
+  Exclude:
+    - 'lib/gitlab/database/load_balancing/resolver.rb'
+    - 'lib/gitlab/database/load_balancing/srv_resolver.rb'
+
+Rails/SkipsModelValidations:
+  Exclude:
+    - 'spec/gitlab/database/load_balancing/connection_proxy_spec.rb'
+    - 'spec/gitlab/database/load_balancing_spec.rb'
+
+RSpec/ExpectInHook:
+  Exclude:
+    - 'spec/gitlab/database/load_balancing/host_spec.rb'
+    - 'spec/gitlab/database/load_balancing/load_balancer_spec.rb'
+
+RSpec/ImplicitSubject:
+  Exclude:
+    - 'spec/gitlab/database/load_balancing/configuration_spec.rb'
+
+RSpec/MultipleMemoizedHelpers:
+  Max: 9
+
+RSpec/VerifiedDoubles:
+  Exclude:
+    - 'spec/gitlab/database/load_balancing/configuration_spec.rb'
+    - 'spec/gitlab/database/load_balancing/connection_proxy_spec.rb'
+    - 'spec/gitlab/database/load_balancing/host_list_spec.rb'
+    - 'spec/gitlab/database/load_balancing/host_spec.rb'
+    - 'spec/gitlab/database/load_balancing/load_balancer_spec.rb'
+    - 'spec/gitlab/database/load_balancing/resolver_spec.rb'
+    - 'spec/gitlab/database/load_balancing/service_discovery_spec.rb'
+    - 'spec/gitlab/database/load_balancing/setup_spec.rb'
+    - 'spec/gitlab/database/load_balancing_spec.rb'
+
+Layout/LineLength:
+  Max: 129
diff --git a/gems/gitlab-database-load_balancing/.rubocop_todo.yml b/gems/gitlab-database-load_balancing/.rubocop_todo.yml
new file mode 100644
index 00000000000000..12a7b201f51cba
--- /dev/null
+++ b/gems/gitlab-database-load_balancing/.rubocop_todo.yml
@@ -0,0 +1,41 @@
+Lint/AssignmentInCondition:
+  Exclude:
+    - 'lib/gitlab/database/load_balancing/configuration.rb'
+    - 'lib/gitlab/database/load_balancing/host.rb'
+    - 'lib/gitlab/database/load_balancing/load_balancer.rb'
+
+Lint/DuplicateBranch:
+  Exclude:
+    - 'lib/gitlab/database/load_balancing/load_balancer.rb'
+
+Lint/EmptyBlock:
+  Exclude:
+    - 'spec/gitlab/database/load_balancing/load_balancer_spec.rb'
+
+Lint/RedundantCopDisableDirective:
+  Exclude:
+    - 'lib/gitlab/database/load_balancing.rb'
+
+Lint/UnusedMethodArgument:
+  Exclude:
+    - 'lib/gitlab/database/load_balancing/connection_proxy.rb'
+    - 'lib/gitlab/database/load_balancing/primary_host.rb'
+
+Naming/RescuedExceptionsVariableName:
+  Exclude:
+    - 'lib/gitlab/database/load_balancing/load_balancer.rb'
+    - 'lib/gitlab/database/load_balancing/service_discovery.rb'
+    - 'spec/gitlab/database/load_balancing/host_spec.rb'
+    - 'spec/gitlab/database/load_balancing/load_balancer_spec.rb'
+
+Style/GuardClause:
+  Exclude:
+    - 'lib/gitlab/database/load_balancing/load_balancer.rb'
+
+Style/Lambda:
+  Exclude:
+    - 'lib/gitlab/database/load_balancing/action_cable_callbacks.rb'
+
+Style/RedundantSelf:
+  Exclude:
+    - 'lib/gitlab/database/load_balancing/service_discovery.rb'
diff --git a/gems/gitlab-database-load_balancing/README.md b/gems/gitlab-database-load_balancing/README.md
index e1a963dba5087c..fb05fa373fe9dd 100644
--- a/gems/gitlab-database-load_balancing/README.md
+++ b/gems/gitlab-database-load_balancing/README.md
@@ -1,3 +1,3 @@
 # GitLab Database Load Balancing
 
-This gem is a stub for a move of all `Gitlab::Database::LoadBalancing` code.
+TODO
diff --git a/lib/gitlab/database/load_balancing.rb b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing.rb
similarity index 59%
rename from lib/gitlab/database/load_balancing.rb
rename to gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing.rb
index 5f9416fb4db160..c35712d8a7b3cd 100644
--- a/lib/gitlab/database/load_balancing.rb
+++ b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing.rb
@@ -1,5 +1,15 @@
 # frozen_string_literal: true
 
+require 'pg'
+require 'active_support/time_with_zone'
+require 'active_record'
+require 'active_record/database_configurations'
+require 'active_job'
+
+Dir['load_balancing/**/*.rb', base: File.dirname(__FILE__)].each do |file|
+  require_relative file
+end
+
 module Gitlab
   module Database
     module LoadBalancing
@@ -18,9 +28,9 @@ module LoadBalancing
         ActiveRecord::ConnectionNotEstablished
       ].freeze
 
-      def self.base_models
-        @base_models ||= ::Gitlab::Database.database_base_models_using_load_balancing.values.freeze
-      end
+      mattr_accessor :base_models
+      mattr_accessor :default_pool_size
+      mattr_accessor :enabled, default: true
 
       def self.each_load_balancer
         return to_enum(__method__) unless block_given?
@@ -49,7 +59,7 @@ def self.release_hosts
       def self.db_role_for_connection(connection)
         return ROLE_UNKNOWN if connection.is_a?(::Gitlab::Database::LoadBalancing::ConnectionProxy)
 
-        db_config = Database.db_config_for_connection(connection)
+        db_config = db_config_for_connection(connection)
         return ROLE_UNKNOWN unless db_config
 
         if db_config.name.ends_with?(LoadBalancer::REPLICA_SUFFIX)
@@ -58,6 +68,23 @@ def self.db_role_for_connection(connection)
           ROLE_PRIMARY
         end
       end
+
+      def self.db_config_for_connection(connection)
+        return unless connection
+
+        # For a ConnectionProxy we want to avoid ambiguous db_config as it may
+        # sometimes default to replica so we always return the primary config
+        # instead.
+        if connection.is_a?(::Gitlab::Database::LoadBalancing::ConnectionProxy) # rubocop:disable Style/IfUnlessModifier
+          return connection.load_balancer.configuration.db_config
+        end
+
+        # During application init we might receive `NullPool`
+        return unless connection.respond_to?(:pool) &&
+          connection.pool.respond_to?(:db_config)
+
+        connection.pool.db_config
+      end
     end
   end
 end
diff --git a/lib/gitlab/database/load_balancing/action_cable_callbacks.rb b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/action_cable_callbacks.rb
similarity index 100%
rename from lib/gitlab/database/load_balancing/action_cable_callbacks.rb
rename to gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/action_cable_callbacks.rb
diff --git a/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/callbacks.rb b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/callbacks.rb
new file mode 100644
index 00000000000000..086819e016ba93
--- /dev/null
+++ b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/callbacks.rb
@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module LoadBalancing
+      module Callbacks
+        mattr_accessor :logger_proc
+        mattr_accessor :metrics_host_gauge_proc, :track_exception_proc
+        mattr_accessor :set_wal_for_proc, :get_wal_for_proc, :del_wal_for_proc
+
+        NULL_LOGGER = Logger.new(File::NULL)
+
+        def self.configure!
+          yield(self)
+        end
+
+        def self.set_wal_for(key, wal, ex:)
+          set_wal_for_proc&.call(key, wal, ex: ex)
+        end
+
+        def self.get_wal_for(key)
+          get_wal_for_proc&.call(key)
+        end
+
+        def self.del_wal_for(key)
+          del_wal_for_proc&.call(key)
+        end
+
+        def self.logger
+          logger_proc&.call || NULL_LOGGER
+        end
+
+        def self.metrics_host_gauge(labels, value)
+          metrics_host_gauge_proc&.call(labels, value)
+        end
+
+        def self.track_exception(ex)
+          track_exception_proc&.call(ex)
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/database/load_balancing/configuration.rb b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/configuration.rb
similarity index 90%
rename from lib/gitlab/database/load_balancing/configuration.rb
rename to gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/configuration.rb
index 50472bd5780c60..4bcc3fde99c0ba 100644
--- a/lib/gitlab/database/load_balancing/configuration.rb
+++ b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/configuration.rb
@@ -1,13 +1,15 @@
 # frozen_string_literal: true
 
+require 'active_support/core_ext/numeric/bytes'
+
 module Gitlab
   module Database
     module LoadBalancing
       # Configuration settings for a single LoadBalancer instance.
       class Configuration
         attr_accessor :hosts, :max_replication_difference,
-                      :max_replication_lag_time, :replica_check_interval,
-                      :service_discovery
+          :max_replication_lag_time, :replica_check_interval,
+          :service_discovery
 
         # Creates a configuration object for the given ActiveRecord model.
         def self.for_model(model)
@@ -82,7 +84,7 @@ def pool_size
           # To support this scenario, we always attempt to read the pool size
           # from the model's configuration.
           @model.connection_db_config.configuration_hash[:pool] ||
-            Database.default_pool_size
+            Gitlab::Database::LoadBalancing.default_pool_size
         end
 
         # Returns `true` if the use of load balancing replicas should be
@@ -91,7 +93,7 @@ def pool_size
         # 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?
+          return false unless Gitlab::Database::LoadBalancing.enabled
 
           hosts.any? || service_discovery_enabled?
         end
@@ -99,7 +101,7 @@ def load_balancing_enabled?
         # This is disabled for Rake tasks to ensure e.g. database migrations
         # always produce consistent results.
         def service_discovery_enabled?
-          return false if Gitlab::Runtime.rake?
+          return false unless Gitlab::Database::LoadBalancing.enabled
 
           service_discovery[:record].present?
         end
diff --git a/lib/gitlab/database/load_balancing/connection_proxy.rb b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/connection_proxy.rb
similarity index 99%
rename from lib/gitlab/database/load_balancing/connection_proxy.rb
rename to gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/connection_proxy.rb
index 67cde3e687d3a4..7ae3e706b391d7 100644
--- a/lib/gitlab/database/load_balancing/connection_proxy.rb
+++ b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/connection_proxy.rb
@@ -149,3 +149,5 @@ def read_only_transaction?
     end
   end
 end
+
+# rubocop:enable GitlabSecurity/PublicSend
diff --git a/lib/gitlab/database/load_balancing/host.rb b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/host.rb
similarity index 95%
rename from lib/gitlab/database/load_balancing/host.rb
rename to gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/host.rb
index 880324f3ae4675..8acc7c01c19dfc 100644
--- a/lib/gitlab/database/load_balancing/host.rb
+++ b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/host.rb
@@ -7,7 +7,8 @@ module LoadBalancing
       class Host
         attr_reader :pool, :last_checked_at, :intervals, :load_balancer, :host, :port
 
-        delegate :connection, :release_connection, :enable_query_cache!, :disable_query_cache!, :query_cache_enabled, to: :pool
+        delegate :connection, :release_connection, :enable_query_cache!, :disable_query_cache!, :query_cache_enabled,
+          to: :pool
 
         CONNECTION_ERRORS = [
           ActionView::Template::Error,
@@ -77,9 +78,9 @@ def initialize(host, load_balancer, port: nil)
         # timeout - The time after which the pool should be forcefully
         #           disconnected.
         def disconnect!(timeout: 120)
-          start_time = ::Gitlab::Metrics::System.monotonic_time
+          start_time = ::Gitlab::Utils::System.monotonic_time
 
-          while (::Gitlab::Metrics::System.monotonic_time - start_time) <= timeout
+          while (::Gitlab::Utils::System.monotonic_time - start_time) <= timeout
             return if try_disconnect
 
             sleep(2)
@@ -104,7 +105,7 @@ def force_disconnect!
         end
 
         def offline!
-          ::Gitlab::Database::LoadBalancing::Logger.warn(
+          ::Gitlab::Database::LoadBalancing::Callbacks.logger.warn(
             event: :host_offline,
             message: 'Marking host as offline',
             db_host: @host,
@@ -124,7 +125,7 @@ def online?
 
           # Log that the host came back online if it was previously offline
           if @online && !was_online
-            ::Gitlab::Database::LoadBalancing::Logger.info(
+            ::Gitlab::Database::LoadBalancing::Callbacks.logger.info(
               event: :host_online,
               message: 'Host is online after replica status check',
               db_host: @host,
@@ -132,7 +133,7 @@ def online?
             )
           # Always log if the host goes offline
           elsif !@online
-            ::Gitlab::Database::LoadBalancing::Logger.warn(
+            ::Gitlab::Database::LoadBalancing::Callbacks.logger.warn(
               event: :host_offline,
               message: 'Host is offline after replica status check',
               db_host: @host,
diff --git a/lib/gitlab/database/load_balancing/host_list.rb b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/host_list.rb
similarity index 92%
rename from lib/gitlab/database/load_balancing/host_list.rb
rename to gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/host_list.rb
index fb3175c7d5d899..7a3750f637b430 100644
--- a/lib/gitlab/database/load_balancing/host_list.rb
+++ b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/host_list.rb
@@ -10,7 +10,6 @@ def initialize(hosts = [])
           @hosts = hosts.shuffle
           @index = 0
           @mutex = Mutex.new
-          @hosts_gauge = Gitlab::Metrics.gauge(:db_load_balancing_hosts, 'Current number of load balancing hosts')
 
           set_metrics!
         end
@@ -80,7 +79,7 @@ def next_host
         end
 
         def set_metrics!
-          @hosts_gauge.set({}, @hosts.length)
+          Callbacks.metrics_host_gauge({}, @hosts.length)
         end
       end
     end
diff --git a/lib/gitlab/database/load_balancing/load_balancer.rb b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/load_balancer.rb
similarity index 96%
rename from lib/gitlab/database/load_balancing/load_balancer.rb
rename to gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/load_balancer.rb
index 55a27f89b36c8a..c89491b0338ae9 100644
--- a/lib/gitlab/database/load_balancing/load_balancer.rb
+++ b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/load_balancer.rb
@@ -76,7 +76,7 @@ def read(&block)
                 # times before using the primary instead.
                 will_retry = conflict_retried < @host_list.length * 3
 
-                ::Gitlab::Database::LoadBalancing::Logger.warn(
+                ::Gitlab::Database::LoadBalancing::Callbacks.logger.warn(
                   event: :host_query_conflict,
                   message: 'Query conflict on host',
                   conflict_retried: conflict_retried,
@@ -101,7 +101,7 @@ def read(&block)
             end
           end
 
-          ::Gitlab::Database::LoadBalancing::Logger.warn(
+          ::Gitlab::Database::LoadBalancing::Callbacks.logger.warn(
             event: :no_secondaries_available,
             message: 'No secondaries were available, using primary instead',
             conflict_retried: conflict_retried,
@@ -132,7 +132,7 @@ def read_write
             transaction_open = connection.transaction_open?
 
             if attempt && attempt > 1
-              ::Gitlab::Database::LoadBalancing::Logger.warn(
+              ::Gitlab::Database::LoadBalancing::Callbacks.logger.warn(
                 event: :read_write_retry,
                 message: 'A read_write block was retried because of connection error'
               )
@@ -143,7 +143,7 @@ def read_write
             # No leaking will happen on the final attempt. Leaks are caused by subsequent retries
             not_final_attempt = attempt && attempt < attempts
             if transaction_open && connection_error?(e) && not_final_attempt
-              ::Gitlab::Database::LoadBalancing::Logger.warn(
+              ::Gitlab::Database::LoadBalancing::Callbacks.logger.warn(
                 event: :transaction_leak,
                 message: 'A write transaction has leaked during database fail-over'
               )
@@ -311,7 +311,6 @@ def create_replica_connection_pool(pool_size, host = nil, port = nil)
         # ActiveRecord::ConnectionAdapters::ConnectionHandler handles fetching,
         # and caching for connections pools for each "connection", so we
         # leverage that.
-        # rubocop:disable Database/MultipleDatabases
         def pool
           ActiveRecord::Base.connection_handler.retrieve_connection_pool(
             @configuration.connection_specification_name,
@@ -319,7 +318,6 @@ def pool
             shard: ActiveRecord::Base.default_shard
           ) || raise(::ActiveRecord::ConnectionNotEstablished)
         end
-        # rubocop:enable Database/MultipleDatabases
 
         def wal_diff(location1, location2)
           read_write do |connection|
@@ -368,7 +366,7 @@ def get_write_location(ar_connection)
                     END AS location;
               NEWSQL
             else
-              <<~SQL
+              <<~SQL.squish
                 SELECT pg_current_wal_insert_lsn()::text AS location
               SQL
             end
@@ -380,7 +378,8 @@ def get_write_location(ar_connection)
         def raise_if_concurrent_ruby!
           Gitlab::Utils.raise_if_concurrent_ruby!(:db)
         rescue StandardError => e
-          Gitlab::ErrorTracking.track_and_raise_for_dev_exception(e)
+          ::Gitlab::Database::LoadBalancing::Callbacks.track_exception(e)
+          raise if Rails.env.development? || Rails.env.test?
         end
       end
     end
diff --git a/lib/gitlab/database/load_balancing/primary_host.rb b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/primary_host.rb
similarity index 96%
rename from lib/gitlab/database/load_balancing/primary_host.rb
rename to gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/primary_host.rb
index fb52b384ddb4a4..48a6c034fe6b02 100644
--- a/lib/gitlab/database/load_balancing/primary_host.rb
+++ b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/primary_host.rb
@@ -49,7 +49,7 @@ def disconnect!(timeout: 120)
         end
 
         def offline!
-          ::Gitlab::Database::LoadBalancing::Logger.warn(
+          ::Gitlab::Database::LoadBalancing::Callbacks.logger.warn(
             event: :host_offline,
             message: 'Marking primary host as offline'
           )
diff --git a/lib/gitlab/database/load_balancing/resolver.rb b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/resolver.rb
similarity index 100%
rename from lib/gitlab/database/load_balancing/resolver.rb
rename to gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/resolver.rb
diff --git a/lib/gitlab/database/load_balancing/service_discovery.rb b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/service_discovery.rb
similarity index 94%
rename from lib/gitlab/database/load_balancing/service_discovery.rb
rename to gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/service_discovery.rb
index 3d11f0f88c1515..2d2a0cb3f3b35d 100644
--- a/lib/gitlab/database/load_balancing/service_discovery.rb
+++ b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/service_discovery.rb
@@ -17,8 +17,7 @@ class ServiceDiscovery
 
         attr_accessor :refresh_thread, :refresh_thread_last_run, :refresh_thread_interruption_logged
 
-        attr_reader :interval, :record, :record_type, :disconnect_timeout,
-                    :load_balancer
+        attr_reader :interval, :record, :record_type, :disconnect_timeout, :load_balancer
 
         MAX_SLEEP_ADJUSTMENT = 10
         MAX_DISCOVERY_RETRIES = 3
@@ -96,9 +95,9 @@ def perform_service_discovery
           rescue StandardError => error
             # Any exceptions that might occur should be reported to
             # Sentry, instead of silently terminating this thread.
-            Gitlab::ErrorTracking.track_exception(error)
+            ::Gitlab::Database::LoadBalancing::Callbacks.track_exception(error)
 
-            Gitlab::Database::LoadBalancing::Logger.error(
+            ::Gitlab::Database::LoadBalancing::Callbacks.logger.error(
               event: :service_discovery_failure,
               message: "Service discovery encountered an error: #{error.message}",
               host_list_length: load_balancer.host_list.length
@@ -123,7 +122,7 @@ def refresh_if_necessary
           current = addresses_from_load_balancer
 
           if from_dns != current
-            ::Gitlab::Database::LoadBalancing::Logger.info(
+            ::Gitlab::Database::LoadBalancing::Callbacks.logger.info(
               event: :host_list_update,
               message: "Updating the host list for service discovery",
               host_list_length: from_dns.length,
@@ -211,7 +210,7 @@ def log_refresh_thread_interruption
           return if refresh_thread_last_run.blank? || refresh_thread_interruption_logged ||
             (refresh_thread_last_run + DISCOVERY_THREAD_REFRESH_DELTA.minutes).future?
 
-          Gitlab::Database::LoadBalancing::Logger.error(
+          ::Gitlab::Database::LoadBalancing::Callbacks.logger.error(
             event: :service_discovery_refresh_thread_interrupt,
             refresh_thread_last_run: refresh_thread_last_run,
             thread_status: refresh_thread&.status&.to_s,
@@ -232,12 +231,12 @@ def record_type_for(type)
         def addresses_from_srv_record(response)
           srv_resolver = SrvResolver.new(resolver, response.additional)
 
-          response.answer.map do |r|
+          response.answer.filter_map do |r|
             address = srv_resolver.address_for(r.host.to_s)
             next unless address
 
             Address.new(address.to_s, r.port)
-          end.compact
+          end
         end
 
         def addresses_from_a_record(resources)
@@ -252,13 +251,13 @@ def sampler
         def disconnect_old_hosts(hosts)
           return unless hosts.present?
 
-          gentle_disconnect_start = ::Gitlab::Metrics::System.monotonic_time
+          gentle_disconnect_start = ::Gitlab::Utils::System.monotonic_time
           gentle_disconnect_deadline = gentle_disconnect_start + disconnect_timeout
 
           hosts_to_disconnect = hosts
 
           gentle_disconnect_duration = Benchmark.realtime do
-            while ::Gitlab::Metrics::System.monotonic_time < gentle_disconnect_deadline
+            while ::Gitlab::Utils::System.monotonic_time < gentle_disconnect_deadline
               hosts_to_disconnect = hosts_to_disconnect.reject(&:try_disconnect)
 
               break if hosts_to_disconnect.empty?
@@ -275,7 +274,7 @@ def disconnect_old_hosts(hosts)
           formatted_hosts = hosts_to_disconnect.map { |h| "#{h.host}:#{h.port}" }
           total_disconnect_duration = gentle_disconnect_duration + force_disconnect_duration
 
-          ::Gitlab::Database::LoadBalancing::Logger.info(
+          ::Gitlab::Database::LoadBalancing::Callbacks.logger.info(
             event: :host_list_disconnection,
             message: "Disconnected #{formatted_hosts} old load balancing hosts after #{total_disconnect_duration}s",
             gentle_disconnect_duration_s: gentle_disconnect_duration,
diff --git a/lib/gitlab/database/load_balancing/service_discovery/sampler.rb b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/service_discovery/sampler.rb
similarity index 93%
rename from lib/gitlab/database/load_balancing/service_discovery/sampler.rb
rename to gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/service_discovery/sampler.rb
index 4fc11cc1035f92..55e320fd42ecc7 100644
--- a/lib/gitlab/database/load_balancing/service_discovery/sampler.rb
+++ b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/service_discovery/sampler.rb
@@ -16,7 +16,7 @@ def initialize(max_replica_pools:, seed: Random.new_seed)
           def sample(addresses)
             return addresses if @max_replica_pools.nil? || addresses.count <= @max_replica_pools
 
-            ::Gitlab::Database::LoadBalancing::Logger.debug(
+            ::Gitlab::Database::LoadBalancing::Callbacks.logger.debug(
               event: :host_list_limit_exceeded,
               message: "Host list length exceeds max_replica_pools so random hosts will be chosen.",
               max_replica_pools: @max_replica_pools,
@@ -38,7 +38,7 @@ def sample(addresses)
             while selected_addresses.count < @max_replica_pools
               # Loop over all hostnames grabbing one address at a time to
               # evenly distribute across all hostnames
-              addresses_by_host.each do |host, addresses|
+              addresses_by_host.each do |_host, addresses|
                 next if addresses.empty?
 
                 selected_addresses << addresses.pop
diff --git a/lib/gitlab/database/load_balancing/session.rb b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/session.rb
similarity index 95%
rename from lib/gitlab/database/load_balancing/session.rb
rename to gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/session.rb
index 3682c9265c2212..8f5f1b06c19396 100644
--- a/lib/gitlab/database/load_balancing/session.rb
+++ b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/session.rb
@@ -41,7 +41,7 @@ def use_primary!
           @use_primary = true
         end
 
-        def use_primary(&blk)
+        def use_primary(&_blk)
           used_primary = @use_primary
           @use_primary = true
           yield
@@ -49,7 +49,7 @@ def use_primary(&blk)
           @use_primary = used_primary || @performed_write
         end
 
-        def ignore_writes(&block)
+        def ignore_writes(&_block)
           @ignore_writes = true
 
           yield
@@ -66,7 +66,7 @@ def ignore_writes(&block)
         #
         # Write and ambiguous queries inside this block are still handled by
         # the primary.
-        def use_replicas_for_read_queries(&blk)
+        def use_replicas_for_read_queries(&_blk)
           previous_flag = @use_replicas_for_read_queries
           @use_replicas_for_read_queries = true
           yield
@@ -89,7 +89,7 @@ def use_replicas_for_read_queries?
         # - If the queries are about to write
         # - The current session already performed writes
         # - It prefers to use primary, aka, use_primary or use_primary! were called
-        def fallback_to_replicas_for_ambiguous_queries(&blk)
+        def fallback_to_replicas_for_ambiguous_queries(&_blk)
           previous_flag = @fallback_to_replicas_for_ambiguous_queries
           @fallback_to_replicas_for_ambiguous_queries = true
           yield
diff --git a/lib/gitlab/database/load_balancing/setup.rb b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/setup.rb
similarity index 94%
rename from lib/gitlab/database/load_balancing/setup.rb
rename to gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/setup.rb
index 2e65e1c8e56fae..d5be4800c6d7f1 100644
--- a/lib/gitlab/database/load_balancing/setup.rb
+++ b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/setup.rb
@@ -18,7 +18,7 @@ def setup
           setup_connection_proxy
           setup_service_discovery
 
-          ::Gitlab::Database::LoadBalancing::Logger.debug(
+          ::Gitlab::Database::LoadBalancing::Callbacks.logger.debug(
             event: :setup,
             model: model.name,
             start_service_discovery: @start_service_discovery
@@ -30,7 +30,7 @@ def configure_connection
 
           hash = db_config_object.configuration_hash.merge(
             prepared_statements: false,
-            pool: Gitlab::Database.default_pool_size
+            pool: Gitlab::Database::LoadBalancing.default_pool_size
           )
 
           hash_config = ActiveRecord::DatabaseConfigurations::HashConfig.new(
diff --git a/lib/gitlab/database/load_balancing/srv_resolver.rb b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/srv_resolver.rb
similarity index 84%
rename from lib/gitlab/database/load_balancing/srv_resolver.rb
rename to gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/srv_resolver.rb
index 1f599ef4a273f9..312c95dd0283e5 100644
--- a/lib/gitlab/database/load_balancing/srv_resolver.rb
+++ b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/srv_resolver.rb
@@ -28,12 +28,11 @@ def address_for(host)
         private
 
         def addresses_from_additional
-          strong_memoize(:addresses_from_additional) do
-            additional.each_with_object({}) do |rr, h|
-              h[rr.name] = rr.address if rr.is_a?(Net::DNS::RR::A) || rr.is_a?(Net::DNS::RR::AAAA)
-            end
+          additional.each_with_object({}) do |rr, h|
+            h[rr.name] = rr.address if rr.is_a?(Net::DNS::RR::A) || rr.is_a?(Net::DNS::RR::AAAA)
           end
         end
+        strong_memoize_attr :addresses_from_additional
 
         def resolve_host(host)
           record = resolver.search(host, Net::DNS::ANY).answer.find do |rr|
diff --git a/lib/gitlab/database/load_balancing/sticking.rb b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/sticking.rb
similarity index 89%
rename from lib/gitlab/database/load_balancing/sticking.rb
rename to gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/sticking.rb
index df07f510c8d2f1..dfeb30c47ecc81 100644
--- a/lib/gitlab/database/load_balancing/sticking.rb
+++ b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing/sticking.rb
@@ -74,21 +74,15 @@ def with_primary_write_location
         end
 
         def unstick(namespace, id)
-          with_redis do |redis|
-            redis.del(redis_key_for(namespace, id))
-          end
+          Callbacks.del_wal_for(redis_key_for(namespace, id))
         end
 
         def set_write_location_for(namespace, id, location)
-          with_redis do |redis|
-            redis.set(redis_key_for(namespace, id), location, ex: EXPIRATION)
-          end
+          Callbacks.set_wal_for(redis_key_for(namespace, id), location, ex: EXPIRATION)
         end
 
         def last_write_location_for(namespace, id)
-          with_redis do |redis|
-            redis.get(redis_key_for(namespace, id))
-          end
+          Callbacks.get_wal_for(redis_key_for(namespace, id))
         end
 
         def redis_key_for(namespace, id)
@@ -96,10 +90,6 @@ def redis_key_for(namespace, id)
 
           "database-load-balancing/write-location/#{name}/#{namespace}/#{id}"
         end
-
-        def with_redis(&block)
-          Gitlab::Redis::DbLoadBalancing.with(&block)
-        end
       end
     end
   end
diff --git a/spec/fixtures/dns/a_rr.json b/gems/gitlab-database-load_balancing/spec/fixtures/dns/a_rr.json
similarity index 100%
rename from spec/fixtures/dns/a_rr.json
rename to gems/gitlab-database-load_balancing/spec/fixtures/dns/a_rr.json
diff --git a/spec/fixtures/dns/a_with_aaaa_rr_in_additional_section.json b/gems/gitlab-database-load_balancing/spec/fixtures/dns/a_with_aaaa_rr_in_additional_section.json
similarity index 100%
rename from spec/fixtures/dns/a_with_aaaa_rr_in_additional_section.json
rename to gems/gitlab-database-load_balancing/spec/fixtures/dns/a_with_aaaa_rr_in_additional_section.json
diff --git a/spec/fixtures/dns/aaaa_rr.json b/gems/gitlab-database-load_balancing/spec/fixtures/dns/aaaa_rr.json
similarity index 100%
rename from spec/fixtures/dns/aaaa_rr.json
rename to gems/gitlab-database-load_balancing/spec/fixtures/dns/aaaa_rr.json
diff --git a/spec/fixtures/dns/srv_with_a_rr_in_additional_section.json b/gems/gitlab-database-load_balancing/spec/fixtures/dns/srv_with_a_rr_in_additional_section.json
similarity index 100%
rename from spec/fixtures/dns/srv_with_a_rr_in_additional_section.json
rename to gems/gitlab-database-load_balancing/spec/fixtures/dns/srv_with_a_rr_in_additional_section.json
diff --git a/spec/lib/gitlab/database/load_balancing/action_cable_callbacks_spec.rb b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/action_cable_callbacks_spec.rb
similarity index 85%
rename from spec/lib/gitlab/database/load_balancing/action_cable_callbacks_spec.rb
rename to gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/action_cable_callbacks_spec.rb
index a57f02b22dfcbf..2c69f39f131c5f 100644
--- a/spec/lib/gitlab/database/load_balancing/action_cable_callbacks_spec.rb
+++ b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/action_cable_callbacks_spec.rb
@@ -8,7 +8,7 @@
       expect(Gitlab::Database::LoadBalancing).to receive(:release_hosts)
       expect(Gitlab::Database::LoadBalancing::Session).to receive(:clear_session)
 
-      described_class.wrapper.call(nil, lambda {})
+      described_class.wrapper.call(nil, -> {})
     end
 
     context 'with an exception' do
@@ -17,7 +17,7 @@
         expect(Gitlab::Database::LoadBalancing::Session).to receive(:clear_session)
 
         expect do
-          described_class.wrapper.call(nil, lambda { raise 'test_exception' })
+          described_class.wrapper.call(nil, -> { raise 'test_exception' })
         end.to raise_error('test_exception')
       end
     end
diff --git a/spec/lib/gitlab/database/load_balancing/configuration_spec.rb b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/configuration_spec.rb
similarity index 94%
rename from spec/lib/gitlab/database/load_balancing/configuration_spec.rb
rename to gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/configuration_spec.rb
index 7dc2e0be3e5e90..41534c17c06293 100644
--- a/spec/lib/gitlab/database/load_balancing/configuration_spec.rb
+++ b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/configuration_spec.rb
@@ -26,7 +26,7 @@
           use_tcp: false,
           max_replica_pools: nil
         )
-        expect(config.pool_size).to eq(Gitlab::Database.default_pool_size)
+        expect(config.pool_size).to eq(Gitlab::Database::LoadBalancing.default_pool_size)
       end
     end
 
@@ -110,7 +110,7 @@
     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)
+      allow(Gitlab::Database::LoadBalancing).to receive(:enabled).and_return(false)
 
       expect(config.load_balancing_enabled?).to eq(false)
     end
@@ -137,7 +137,7 @@
 
   describe '#service_discovery_enabled?' do
     it 'returns false when running inside a Rake task' do
-      allow(Gitlab::Runtime).to receive(:rake?).and_return(true)
+      allow(Gitlab::Database::LoadBalancing).to receive(:enabled).and_return(false)
 
       config = described_class.new(ActiveRecord::Base)
       config.service_discovery[:record] = 'foo'
@@ -184,7 +184,7 @@
       it 'returns the default pool size' do
         config = described_class.new(model)
 
-        expect(config.pool_size).to eq(Gitlab::Database.default_pool_size)
+        expect(config.pool_size).to eq(Gitlab::Database::LoadBalancing.default_pool_size)
       end
     end
   end
diff --git a/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/connection_proxy_spec.rb
similarity index 97%
rename from spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb
rename to gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/connection_proxy_spec.rb
index a2076f5b95037d..d879f46c6148f9 100644
--- a/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb
+++ b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/connection_proxy_spec.rb
@@ -16,7 +16,7 @@
   end
 
   describe '#select_all' do
-    let(:override_proxy) { ActiveRecord::Base.connection.class }
+    let(:override_proxy) { ApplicationRecord.connection.class }
 
     # We can't use :Gitlab::Utils::Override because this method is dynamically prepended
     it 'method signatures match' do
@@ -141,7 +141,7 @@
         .and_return(session)
     end
 
-    context 'session fallbacks ambiguous queries to replicas' do
+    context 'when session fallbacks ambiguous queries to replicas' do
       let(:replica) { double(:connection) }
 
       before do
@@ -174,7 +174,7 @@
       end
     end
 
-    context 'session does not fallback to replicas for ambiguous queries' do
+    context 'when session does not fallback to replicas for ambiguous queries' do
       let(:primary) { double(:connection) }
 
       before do
@@ -227,7 +227,7 @@
         .not_to raise_error
     end
 
-    context 'current session prefers to fallback ambiguous queries to replicas' do
+    context 'when current session prefers to fallback ambiguous queries to replicas' do
       let(:session) { double(:session) }
 
       before do
diff --git a/spec/lib/gitlab/database/load_balancing/host_list_spec.rb b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/host_list_spec.rb
similarity index 93%
rename from spec/lib/gitlab/database/load_balancing/host_list_spec.rb
rename to gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/host_list_spec.rb
index 9bb8116c434e25..0f4e21bb6d5e48 100644
--- a/spec/lib/gitlab/database/load_balancing/host_list_spec.rb
+++ b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/host_list_spec.rb
@@ -19,6 +19,8 @@
     allow(load_balancer).to receive(:create_replica_connection_pool) do
       double(:replica_connection_pool)
     end
+
+    allow(Gitlab::Database::LoadBalancing::Callbacks).to receive(:metrics_host_gauge)
   end
 
   describe '#initialize' do
@@ -136,6 +138,8 @@
   end
 
   def expect_metrics(hosts)
-    expect(Gitlab::Metrics.registry.get(:db_load_balancing_hosts).get({})).to eq(hosts)
+    expect(Gitlab::Database::LoadBalancing::Callbacks).to have_received(:metrics_host_gauge)
+      .with({}, hosts)
+      .at_least(1).times
   end
 end
diff --git a/spec/lib/gitlab/database/load_balancing/host_spec.rb b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/host_spec.rb
similarity index 96%
rename from spec/lib/gitlab/database/load_balancing/host_spec.rb
rename to gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/host_spec.rb
index 89cecaff075ebf..f3008bd528097d 100644
--- a/spec/lib/gitlab/database/load_balancing/host_spec.rb
+++ b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/host_spec.rb
@@ -83,7 +83,7 @@ def wrapped_exception(wrapper, original)
     it 'marks the host as offline' do
       expect(host.pool).to receive(:disconnect!)
 
-      expect(Gitlab::Database::LoadBalancing::Logger).to receive(:warn)
+      expect(Gitlab::Database::LoadBalancing::Callbacks.logger).to receive(:warn)
         .with(hash_including(event: :host_offline))
         .and_call_original
 
@@ -99,8 +99,8 @@ def wrapped_exception(wrapper, original)
 
       it 'returns the latest status' do
         expect(host).not_to receive(:refresh_status)
-        expect(Gitlab::Database::LoadBalancing::Logger).not_to receive(:info)
-        expect(Gitlab::Database::LoadBalancing::Logger).not_to receive(:warn)
+        expect(Gitlab::Database::LoadBalancing::Callbacks.logger).not_to receive(:info)
+        expect(Gitlab::Database::LoadBalancing::Callbacks.logger).not_to receive(:warn)
 
         expect(host).to be_online
       end
@@ -109,8 +109,8 @@ def wrapped_exception(wrapper, original)
         host.offline!
 
         expect(host).not_to receive(:refresh_status)
-        expect(Gitlab::Database::LoadBalancing::Logger).not_to receive(:info)
-        expect(Gitlab::Database::LoadBalancing::Logger).not_to receive(:warn)
+        expect(Gitlab::Database::LoadBalancing::Callbacks.logger).not_to receive(:info)
+        expect(Gitlab::Database::LoadBalancing::Callbacks.logger).not_to receive(:warn)
 
         expect(host).not_to be_online
       end
@@ -131,7 +131,7 @@ def wrapped_exception(wrapper, original)
         # Hosts are online by default
 
         it 'does not log the online event' do
-          expect(Gitlab::Database::LoadBalancing::Logger)
+          expect(Gitlab::Database::LoadBalancing::Callbacks.logger)
             .not_to receive(:info)
             .with(hash_including(event: :host_online))
 
@@ -145,7 +145,7 @@ def wrapped_exception(wrapper, original)
         end
 
         it 'logs the online event' do
-          expect(Gitlab::Database::LoadBalancing::Logger)
+          expect(Gitlab::Database::LoadBalancing::Callbacks.logger)
             .to receive(:info)
             .with(hash_including(event: :host_online))
             .and_call_original
@@ -160,7 +160,7 @@ def wrapped_exception(wrapper, original)
         end
 
         it 'marks the host offline' do
-          expect(Gitlab::Database::LoadBalancing::Logger).to receive(:warn)
+          expect(Gitlab::Database::LoadBalancing::Callbacks.logger).to receive(:warn)
             .with(hash_including(event: :host_offline))
             .and_call_original
 
diff --git a/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/load_balancer_spec.rb
similarity index 98%
rename from spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb
rename to gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/load_balancer_spec.rb
index 3c14dc23a8083a..24bdb7ee6f1a93 100644
--- a/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb
+++ b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/load_balancer_spec.rb
@@ -191,7 +191,7 @@ def twice_wrapped_exception(top, middle, original)
       lb.read { raise conflict_error }
     end
 
-    context 'only primary is configured' do
+    context 'when only primary is configured' do
       let(:lb) do
         config = Gitlab::Database::LoadBalancing::Configuration.new(ActiveRecord::Base)
         allow(config).to receive(:load_balancing_enabled?).and_return(false)
@@ -230,7 +230,7 @@ def twice_wrapped_exception(top, middle, original)
 
       # When no hosts are configured, we don't want to produce any warnings, as
       # they aren't useful/too noisy.
-      expect(Gitlab::Database::LoadBalancing::Logger).not_to receive(:warn)
+      expect(Gitlab::Database::LoadBalancing::Callbacks.logger).not_to receive(:warn)
 
       expect { |b| lb.read(&b) }
         .to yield_with_args(ActiveRecord::Base.retrieve_connection)
@@ -594,7 +594,11 @@ def with_replica_pool(*args)
     it 'returns the diff between two write locations' do
       loc1 = lb.send(:get_write_location, lb.pool.connection)
 
-      create(:user) # This ensures we get a new WAL location
+      ActiveRecord::Schema.define do
+        create_table :_test_load_balancing_test, force: true do |t|
+          t.string :name, null: true
+        end
+      end
 
       loc2 = lb.send(:get_write_location, lb.pool.connection)
       diff = lb.wal_diff(loc2, loc1)
diff --git a/spec/lib/gitlab/database/load_balancing/primary_host_spec.rb b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/primary_host_spec.rb
similarity index 96%
rename from spec/lib/gitlab/database/load_balancing/primary_host_spec.rb
rename to gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/primary_host_spec.rb
index 53605d14c1726d..d1e1b3eda46106 100644
--- a/spec/lib/gitlab/database/load_balancing/primary_host_spec.rb
+++ b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/primary_host_spec.rb
@@ -52,7 +52,7 @@
 
   describe '#offline!' do
     it 'logs the event but does nothing else' do
-      expect(Gitlab::Database::LoadBalancing::Logger).to receive(:warn)
+      expect(Gitlab::Database::LoadBalancing::Callbacks.logger).to receive(:warn)
         .with(hash_including(event: :host_offline))
         .and_call_original
 
diff --git a/spec/lib/gitlab/database/load_balancing/resolver_spec.rb b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/resolver_spec.rb
similarity index 100%
rename from spec/lib/gitlab/database/load_balancing/resolver_spec.rb
rename to gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/resolver_spec.rb
diff --git a/spec/lib/gitlab/database/load_balancing/service_discovery/sampler_spec.rb b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/service_discovery/sampler_spec.rb
similarity index 88%
rename from spec/lib/gitlab/database/load_balancing/service_discovery/sampler_spec.rb
rename to gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/service_discovery/sampler_spec.rb
index 1a49aa2871fab4..b29ea1bc46c779 100644
--- a/spec/lib/gitlab/database/load_balancing/service_discovery/sampler_spec.rb
+++ b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/service_discovery/sampler_spec.rb
@@ -26,10 +26,10 @@
 
     it 'samples random ports across all hosts' do
       expect(sampler.sample(addresses)).to eq([
-                                                address_class.new("127.0.0.1", 6432),
-                                                address_class.new("127.0.0.2", 6435),
-                                                address_class.new("127.0.0.1", 6435)
-                                              ])
+        address_class.new("127.0.0.1", 6432),
+        address_class.new("127.0.0.2", 6435),
+        address_class.new("127.0.0.1", 6435)
+      ])
     end
 
     it 'returns the same answer for the same input when called multiple times' do
diff --git a/spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/service_discovery_spec.rb
similarity index 96%
rename from spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb
rename to gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/service_discovery_spec.rb
index 442fa678d4efc0..1ff6cb381dd1ad 100644
--- a/spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb
+++ b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/service_discovery_spec.rb
@@ -88,7 +88,7 @@
 
         expect(service).not_to receive(:sleep)
 
-        expect(Gitlab::ErrorTracking).not_to receive(:track_exception)
+        expect(Gitlab::Database::LoadBalancing::Callbacks).not_to receive(:track_exception)
 
         service.perform_service_discovery
       end
@@ -96,7 +96,7 @@
 
     context 'with StandardError' do
       before do
-        allow(Gitlab::ErrorTracking).to receive(:track_exception)
+        allow(Gitlab::Database::LoadBalancing::Callbacks).to receive(:track_exception)
         allow(service).to receive(:sleep)
       end
 
@@ -139,7 +139,7 @@
           .to receive(:refresh_if_necessary)
                 .and_raise(error).exactly(described_class::MAX_DISCOVERY_RETRIES).times
 
-        expect(Gitlab::ErrorTracking)
+        expect(Gitlab::Database::LoadBalancing::Callbacks)
           .to receive(:track_exception)
                 .with(error).exactly(described_class::MAX_DISCOVERY_RETRIES).times
 
@@ -251,7 +251,7 @@
       end
 
       it 'does not log any load balancing event' do
-        expect(::Gitlab::Database::LoadBalancing::Logger).not_to receive(:info)
+        expect(::Gitlab::Database::LoadBalancing::Callbacks.logger).not_to receive(:info)
 
         service.replace_hosts([address_foo, address_bar])
       end
@@ -460,7 +460,7 @@
       it 'does not log any interruption' do
         expect(service.refresh_thread_last_run).to be_nil
 
-        expect(Gitlab::Database::LoadBalancing::Logger).not_to receive(:error)
+        expect(Gitlab::Database::LoadBalancing::Callbacks.logger).not_to receive(:error)
 
         subject
       end
@@ -470,7 +470,7 @@
       let(:last_run_timestamp) { Time.current }
 
       it 'does not log if last run time plus delta is in future' do
-        expect(Gitlab::Database::LoadBalancing::Logger).not_to receive(:error)
+        expect(Gitlab::Database::LoadBalancing::Callbacks.logger).not_to receive(:error)
 
         subject
       end
@@ -482,7 +482,7 @@
         it 'does not log if the interruption is already logged' do
           service.refresh_thread_interruption_logged = true
 
-          expect(Gitlab::Database::LoadBalancing::Logger).not_to receive(:error)
+          expect(Gitlab::Database::LoadBalancing::Callbacks.logger).not_to receive(:error)
 
           subject
         end
@@ -490,7 +490,7 @@
         it 'logs the error if the interruption was not logged before' do
           expect(service.refresh_thread_interruption_logged).not_to be_present
 
-          expect(Gitlab::Database::LoadBalancing::Logger).to receive(:error).with(
+          expect(Gitlab::Database::LoadBalancing::Callbacks.logger).to receive(:error).with(
             event: :service_discovery_refresh_thread_interrupt,
             refresh_thread_last_run: last_run_timestamp,
             thread_status: refresh_thread.status.to_s,
diff --git a/spec/lib/gitlab/database/load_balancing/session_spec.rb b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/session_spec.rb
similarity index 100%
rename from spec/lib/gitlab/database/load_balancing/session_spec.rb
rename to gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/session_spec.rb
diff --git a/spec/lib/gitlab/database/load_balancing/setup_spec.rb b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/setup_spec.rb
similarity index 97%
rename from spec/lib/gitlab/database/load_balancing/setup_spec.rb
rename to gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/setup_spec.rb
index fa6d71bca7fab9..2e62dd0f95d443 100644
--- a/spec/lib/gitlab/database/load_balancing/setup_spec.rb
+++ b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/setup_spec.rb
@@ -30,7 +30,7 @@
         .with('test', 'main', {
           host: 'localhost',
           prepared_statements: false,
-          pool: Gitlab::Database.default_pool_size
+          pool: Gitlab::Database::LoadBalancing.default_pool_size
         })
         .and_call_original
 
@@ -119,7 +119,7 @@
     end
   end
 
-  context 'uses correct base models', :reestablished_active_record_base do
+  context 'when using correct base models', :reestablished_active_record_base do
     using RSpec::Parameterized::TableSyntax
 
     let(:ci_class) do
diff --git a/spec/lib/gitlab/database/load_balancing/srv_resolver_spec.rb b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/srv_resolver_spec.rb
similarity index 91%
rename from spec/lib/gitlab/database/load_balancing/srv_resolver_spec.rb
rename to gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/srv_resolver_spec.rb
index 6ac0608d4851de..27c8a0b7c26420 100644
--- a/spec/lib/gitlab/database/load_balancing/srv_resolver_spec.rb
+++ b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/srv_resolver_spec.rb
@@ -52,8 +52,9 @@
   end
 
   def dns_response_packet_from_fixture(fixture_name)
-    fixture         = File.read(Rails.root + "spec/fixtures/dns/#{fixture_name}.json")
-    encoded_payload = Gitlab::Json.parse(fixture)['payload']
+    path            = File.join(File.dirname(__FILE__), "../../../fixtures/dns/#{fixture_name}.json")
+    fixture         = File.read(path)
+    encoded_payload = JSON.parse(fixture)['payload']
     payload         = Base64.decode64(encoded_payload)
 
     Net::DNS::Packet.parse(payload)
diff --git a/spec/lib/gitlab/database/load_balancing/sticking_spec.rb b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/sticking_spec.rb
similarity index 87%
rename from spec/lib/gitlab/database/load_balancing/sticking_spec.rb
rename to gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/sticking_spec.rb
index 8c2901c3b8994e..0e6dd4b354a6e8 100644
--- a/spec/lib/gitlab/database/load_balancing/sticking_spec.rb
+++ b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/sticking_spec.rb
@@ -11,16 +11,12 @@
     described_class.new(load_balancer)
   end
 
-  let(:redis) { instance_double(::Gitlab::Redis::MultiStore) }
-
   before do
-    allow(::Gitlab::Redis::DbLoadBalancing).to receive(:with).and_yield(redis)
-
     allow(ActiveRecord::Base.load_balancer)
       .to receive(:primary_write_location)
       .and_return(primary_write_location)
 
-    allow(redis).to receive(:get)
+    allow(Gitlab::Database::LoadBalancing::Callbacks).to receive(:get_wal_for)
       .with("database-load-balancing/write-location/#{load_balancer.name}/user/42")
       .and_return(last_write_location)
   end
@@ -47,7 +43,7 @@
         it 'returns false, does not unstick and calls use_primary!' do
           expect(load_balancer).not_to receive(:select_up_to_date_host)
 
-          expect(redis).not_to receive(:del)
+          expect(Gitlab::Database::LoadBalancing::Callbacks).not_to receive(:del_wal_for)
           expect(::Gitlab::Database::LoadBalancing::Session.current).to receive(:use_primary!)
 
           expect(sticking.find_caught_up_replica(:user, 42, use_primary_on_empty_location: true)).to eq(false)
@@ -60,8 +56,7 @@
         expect(load_balancer).to receive(:select_up_to_date_host).with(last_write_location)
           .and_return(::Gitlab::Database::LoadBalancing::LoadBalancer::ALL_CAUGHT_UP)
 
-        expect(redis)
-          .to receive(:del)
+        expect(Gitlab::Database::LoadBalancing::Callbacks).to receive(:del_wal_for)
           .with("database-load-balancing/write-location/#{load_balancer.name}/user/42")
 
         expect(sticking.find_caught_up_replica(:user, 42)).to eq(true)
@@ -73,7 +68,7 @@
         expect(load_balancer).to receive(:select_up_to_date_host).with(last_write_location)
           .and_return(::Gitlab::Database::LoadBalancing::LoadBalancer::ANY_CAUGHT_UP)
 
-        expect(redis).not_to receive(:del)
+        expect(Gitlab::Database::LoadBalancing::Callbacks).not_to receive(:del_wal_for)
 
         expect(sticking.find_caught_up_replica(:user, 42)).to eq(true)
       end
@@ -86,7 +81,7 @@
       end
 
       it 'returns false, does not unstick and calls use_primary!' do
-        expect(redis).not_to receive(:del)
+        expect(Gitlab::Database::LoadBalancing::Callbacks).not_to receive(:del_wal_for)
         expect(::Gitlab::Database::LoadBalancing::Session.current).to receive(:use_primary!)
 
         expect(sticking.find_caught_up_replica(:user, 42)).to eq(false)
@@ -94,7 +89,7 @@
 
       context 'when use_primary_on_failure is false' do
         it 'does not call use_primary!' do
-          expect(redis).not_to receive(:del)
+          expect(Gitlab::Database::LoadBalancing::Callbacks).not_to receive(:del_wal_for)
           expect(::Gitlab::Database::LoadBalancing::Session.current).not_to receive(:use_primary!)
 
           expect(sticking.find_caught_up_replica(:user, 42, use_primary_on_failure: false)).to eq(false)
@@ -110,8 +105,7 @@
         .and_return(false)
 
       ids.each do |id|
-        expect(redis)
-          .to receive(:set)
+        expect(Gitlab::Database::LoadBalancing::Callbacks).to receive(:set_wal_for)
           .with("database-load-balancing/write-location/#{load_balancer.name}/user/#{id}", 'the-primary-lsn', ex: 30)
       end
 
diff --git a/spec/lib/gitlab/database/load_balancing/transaction_leaking_spec.rb b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/transaction_leaking_spec.rb
similarity index 87%
rename from spec/lib/gitlab/database/load_balancing/transaction_leaking_spec.rb
rename to gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/transaction_leaking_spec.rb
index 56fbaef031dd26..16a09912df823b 100644
--- a/spec/lib/gitlab/database/load_balancing/transaction_leaking_spec.rb
+++ b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing/transaction_leaking_spec.rb
@@ -2,7 +2,7 @@
 
 require 'spec_helper'
 
-RSpec.describe 'Load balancer behavior with errors inside a transaction', :redis, :delete, feature_category: :database do # rubocop:disable Layout/LineLength
+RSpec.describe 'Load balancer behavior with errors inside a transaction', :redis, :delete, feature_category: :database do
   include StubENV
   let(:model) { ActiveRecord::Base }
   let(:db_host) { model.connection_pool.db_config.host }
@@ -17,7 +17,7 @@
 
     Gitlab::Database::LoadBalancing::Setup.new(model).setup
 
-    model.connection.execute(<<~SQL)
+    model.connection.execute(<<~SQL.squish)
       CREATE TABLE IF NOT EXISTS #{test_table_name} (id SERIAL PRIMARY KEY, value INTEGER)
     SQL
 
@@ -27,7 +27,7 @@
   end
 
   after do
-    model.connection.execute(<<~SQL)
+    model.connection.execute(<<~SQL.squish)
       DROP TABLE IF EXISTS #{test_table_name}
     SQL
 
@@ -43,7 +43,7 @@ def execute(conn)
     # This will result in a PG error, which is not raised.
     # Instead, we retry the statement on a fresh connection (where the pid is different and it does nothing)
     # and the load balancer continues with a fresh connection and no transaction if a transaction was open previously
-    conn.execute(<<~SQL)
+    conn.execute(<<~SQL.squish)
       SELECT CASE
       WHEN pg_backend_pid() = #{backend_pid} THEN
         pg_terminate_backend(#{backend_pid})
@@ -57,7 +57,7 @@ def execute(conn)
 
   context 'in a transaction' do
     it 'raises an exception when a retry would occur' do
-      expect(::Gitlab::Database::LoadBalancing::Logger)
+      expect(::Gitlab::Database::LoadBalancing::Callbacks.logger)
         .not_to receive(:warn).with(hash_including(event: :transaction_leak))
 
       expect do
@@ -70,9 +70,9 @@ def execute(conn)
 
   context 'without a transaction' do
     it 'retries' do
-      expect(::Gitlab::Database::LoadBalancing::Logger)
+      expect(::Gitlab::Database::LoadBalancing::Callbacks.logger)
         .not_to receive(:warn).with(hash_including(event: :transaction_leak))
-      expect(::Gitlab::Database::LoadBalancing::Logger)
+      expect(::Gitlab::Database::LoadBalancing::Callbacks.logger)
         .to receive(:warn).with(hash_including(event: :read_write_retry))
 
       expect { execute(model.connection) }.not_to raise_error
diff --git a/gems/gitlab-database-load_balancing/spec/spec_helper.rb b/gems/gitlab-database-load_balancing/spec/spec_helper.rb
index 71faf49cce0c78..c60e502291be6f 100644
--- a/gems/gitlab-database-load_balancing/spec/spec_helper.rb
+++ b/gems/gitlab-database-load_balancing/spec/spec_helper.rb
@@ -8,6 +8,11 @@
 require 'gitlab/utils/all'
 require 'gitlab/safe_request_store'
 
+require "gitlab/database/load_balancing"
+
+require_relative 'support/database_replica'
+require_relative 'support/application_record'
+
 RSpec.configure do |config|
   include StubRails
   include NextInstanceOf
@@ -25,4 +30,13 @@
   config.around(:example, :request_store) do |example|
     ::Gitlab::SafeRequestStore.ensure_request_store { example.run }
   end
+
+  config.before(:all) do
+    Gitlab::Database::LoadBalancing.default_pool_size = 20
+    Gitlab::Database::LoadBalancing.base_models = [ActiveRecord::Base, Ci::ApplicationRecord]
+
+    Gitlab::Database::LoadBalancing.base_models.each do |model|
+      Gitlab::Database::LoadBalancing::Setup.new(model).setup
+    end
+  end
 end
diff --git a/gems/gitlab-database-load_balancing/spec/support/application_record.rb b/gems/gitlab-database-load_balancing/spec/support/application_record.rb
new file mode 100644
index 00000000000000..acea00b5072cc1
--- /dev/null
+++ b/gems/gitlab-database-load_balancing/spec/support/application_record.rb
@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+class ApplicationRecord < ActiveRecord::Base
+  self.abstract_class = true
+end
+
+module Ci
+  class ApplicationRecord < ApplicationRecord
+    self.abstract_class = true
+  end
+end
diff --git a/gems/gitlab-database-load_balancing/spec/support/database_replica.rb b/gems/gitlab-database-load_balancing/spec/support/database_replica.rb
new file mode 100644
index 00000000000000..b37540105a8e10
--- /dev/null
+++ b/gems/gitlab-database-load_balancing/spec/support/database_replica.rb
@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+RSpec.configure do |config|
+  config.around(:each, :database_replica) do |example|
+    old_proxies = {}
+
+    Gitlab::Database::LoadBalancing.base_models.each do |model|
+      old_proxies[model] = [model.load_balancer, model.connection, model.sticking]
+
+      config = Gitlab::Database::LoadBalancing::Configuration
+        .new(model, [model.connection_db_config.configuration_hash[:host]])
+
+      model.load_balancer = Gitlab::Database::LoadBalancing::LoadBalancer.new(config)
+      model.sticking = Gitlab::Database::LoadBalancing::Sticking.new(model.load_balancer)
+      model.connection = Gitlab::Database::LoadBalancing::ConnectionProxy.new(model.load_balancer)
+    end
+
+    Gitlab::Database::LoadBalancing::Session.clear_session
+
+    example.run
+
+    Gitlab::Database::LoadBalancing::Session.clear_session
+
+    old_proxies.each do |model, proxy|
+      model.load_balancer, model.connection, model.sticking = proxy
+    end
+  end
+end
diff --git a/lib/gitlab/database.rb b/lib/gitlab/database.rb
index 6222155d812f78..ee47762fc2e690 100644
--- a/lib/gitlab/database.rb
+++ b/lib/gitlab/database.rb
@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require 'gitlab/database/load_balancing'
+
 module Gitlab
   module Database
     MAIN_DATABASE_NAME = 'main'
@@ -269,20 +271,7 @@ def self.gitlab_schemas_for_connection(connection)
     end
 
     def self.db_config_for_connection(connection)
-      return unless connection
-
-      # For a ConnectionProxy we want to avoid ambiguous db_config as it may
-      # sometimes default to replica so we always return the primary config
-      # instead.
-      if connection.is_a?(::Gitlab::Database::LoadBalancing::ConnectionProxy)
-        return connection.load_balancer.configuration.db_config
-      end
-
-      # During application init we might receive `NullPool`
-      return unless connection.respond_to?(:pool) &&
-        connection.pool.respond_to?(:db_config)
-
-      connection.pool.db_config
+      Gitlab::Database::LoadBalancing.db_config_for_connection(connection)
     end
 
     # At the moment, the connection can only be retrieved by
-- 
GitLab


From 75dec5a521a9bd39a70b8d9ce0a375220b8339c8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kamil=20Trzci=C5=84ski?= <ayufan@ayufan.eu>
Date: Wed, 22 Nov 2023 10:52:42 +0100
Subject: [PATCH 2/8] Move `spec/gitlab/database/load_balancing_spec.rb`

---
 .../spec}/gitlab/database/load_balancing_spec.rb                  | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 rename {spec/lib => gems/gitlab-database-load_balancing/spec}/gitlab/database/load_balancing_spec.rb (100%)

diff --git a/spec/lib/gitlab/database/load_balancing_spec.rb b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing_spec.rb
similarity index 100%
rename from spec/lib/gitlab/database/load_balancing_spec.rb
rename to gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing_spec.rb
-- 
GitLab


From 779bfc6f8d9fdef784245b32dc6cee3a5e0f7a4a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kamil=20Trzci=C5=84ski?= <ayufan@ayufan.eu>
Date: Wed, 22 Nov 2023 10:58:40 +0100
Subject: [PATCH 3/8] Move some LB specs to initializer specs

---
 .../gitlab/database/load_balancing_spec.rb     | 16 ----------------
 spec/initializers/load_balancing_spec.rb       | 18 ++++++++++++++++++
 2 files changed, 18 insertions(+), 16 deletions(-)

diff --git a/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing_spec.rb b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing_spec.rb
index f65c27688b8e2a..a48d4d875b2437 100644
--- a/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing_spec.rb
+++ b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing_spec.rb
@@ -3,22 +3,6 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::LoadBalancing, :suppress_gitlab_schemas_validate_connection, feature_category: :cell do
-  describe '.base_models' do
-    it 'returns the models to apply load balancing to' do
-      models = described_class.base_models
-
-      expect(models).to include(ActiveRecord::Base)
-
-      if Gitlab::Database.has_config?(:ci)
-        expect(models).to include(Ci::ApplicationRecord)
-      end
-    end
-
-    it 'returns the models as a frozen array' do
-      expect(described_class.base_models).to be_frozen
-    end
-  end
-
   describe '.each_load_balancer' do
     it 'yields every load balancer to the supplied block' do
       lbs = []
diff --git a/spec/initializers/load_balancing_spec.rb b/spec/initializers/load_balancing_spec.rb
index eddedcb2f386fe..db572708d53a77 100644
--- a/spec/initializers/load_balancing_spec.rb
+++ b/spec/initializers/load_balancing_spec.rb
@@ -97,4 +97,22 @@ def simulate_puma_worker
       end
     end
   end
+
+  describe Gitlab::Database::LoadBalancing do
+    describe '.base_models' do
+      it 'returns the models to apply load balancing to' do
+        models = described_class.base_models
+
+        expect(models).to include(ActiveRecord::Base)
+
+        if Gitlab::Database.has_config?(:ci)
+          expect(models).to include(Ci::ApplicationRecord)
+        end
+      end
+
+      it 'returns the models as a frozen array' do
+        expect(described_class.base_models).to be_frozen
+      end
+    end
+  end
 end
-- 
GitLab


From 64da6710f0ff722a3eff0ed6dc3e1e057e8b3c36 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kamil=20Trzci=C5=84ski?= <ayufan@ayufan.eu>
Date: Wed, 22 Nov 2023 11:01:30 +0100
Subject: [PATCH 4/8] Fix Rubocop

---
 gems/gitlab-database-load_balancing/.rubocop.yml              | 1 +
 gems/gitlab-database-load_balancing/.rubocop_todo.yml         | 1 +
 .../spec/gitlab/database/load_balancing_spec.rb               | 4 ++--
 3 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/gems/gitlab-database-load_balancing/.rubocop.yml b/gems/gitlab-database-load_balancing/.rubocop.yml
index 8340cf13f69c26..26b40847d087e9 100644
--- a/gems/gitlab-database-load_balancing/.rubocop.yml
+++ b/gems/gitlab-database-load_balancing/.rubocop.yml
@@ -9,6 +9,7 @@ Gemfile/MissingFeatureCategory:
 Database/MultipleDatabases:
   Exclude:
     - 'lib/gitlab/database/load_balancing/load_balancer.rb'
+    - 'spec/gitlab/database/load_balancing_spec.rb'
     - 'spec/gitlab/database/load_balancing/host_list_spec.rb'
     - 'spec/gitlab/database/load_balancing/host_spec.rb'
     - 'spec/gitlab/database/load_balancing/load_balancer_spec.rb'
diff --git a/gems/gitlab-database-load_balancing/.rubocop_todo.yml b/gems/gitlab-database-load_balancing/.rubocop_todo.yml
index 12a7b201f51cba..94be7317466be8 100644
--- a/gems/gitlab-database-load_balancing/.rubocop_todo.yml
+++ b/gems/gitlab-database-load_balancing/.rubocop_todo.yml
@@ -38,4 +38,5 @@ Style/Lambda:
 
 Style/RedundantSelf:
   Exclude:
+    - 'spec/gitlab/database/load_balancing_spec.rb'
     - 'lib/gitlab/database/load_balancing/service_discovery.rb'
diff --git a/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing_spec.rb b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing_spec.rb
index a48d4d875b2437..e0e11ed7e079f3 100644
--- a/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing_spec.rb
+++ b/gems/gitlab-database-load_balancing/spec/gitlab/database/load_balancing_spec.rb
@@ -477,7 +477,7 @@
       end
     end
 
-    context 'custom connection handling' do
+    context 'with custom connection handling' do
       where(:queries, :expected_role) do
         [
           # Reload cache. The schema loading queries should be handled by
@@ -542,7 +542,7 @@
       end
     end
 
-    context 'a write inside a transaction inside fallback_to_replicas_for_ambiguous_queries block' do
+    context 'with a write inside a transaction inside fallback_to_replicas_for_ambiguous_queries block' do
       it 'raises an exception' do
         expect do
           ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
-- 
GitLab


From 659a7db988274c2389176e21910c7466aee81eb8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kamil=20Trzci=C5=84ski?= <ayufan@ayufan.eu>
Date: Wed, 22 Nov 2023 12:31:48 +0100
Subject: [PATCH 5/8] Tidy initializer code

---
 config/initializers/load_balancing.rb         | 62 ++++++++-----------
 .../lib/gitlab/database/load_balancing.rb     |  4 ++
 spec/initializers/load_balancing_spec.rb      |  5 +-
 3 files changed, 31 insertions(+), 40 deletions(-)

diff --git a/config/initializers/load_balancing.rb b/config/initializers/load_balancing.rb
index 1ebd33541320fe..65295b400d9cd7 100644
--- a/config/initializers/load_balancing.rb
+++ b/config/initializers/load_balancing.rb
@@ -1,51 +1,41 @@
 # frozen_string_literal: true
 
-Gitlab::Application.configure do |config|
-  config.middleware.use(Gitlab::Database::LoadBalancing::RackMiddleware)
+# Configure Load Balancing
+Gitlab::Database::LoadBalancing.configure! do |load_balancer|
+  load_balancer.enabled = !Gitlab::Runtime.rake?
+  load_balancer.default_pool_size = Gitlab::Database.default_pool_size
+  load_balancer.base_models = ::Gitlab::Database.database_base_models_using_load_balancing.values.freeze
 
-  # We need re-rerun the setup when code reloads in development
-  config.reloader.to_prepare do
-    if Gitlab.dev_or_test_env?
-      Gitlab::Database::LoadBalancing.base_models.each do |model|
-        Gitlab::Database::LoadBalancing::Setup.new(model).setup
-      end
-    end
+  Gitlab::Cluster::LifecycleEvents.on_worker_start do
+    load_balancer.enabled = !Gitlab::Runtime.rake?
+    load_balancer.default_pool_size = Gitlab::Database.default_pool_size
   end
 end
 
-# Configure load balancing
-Gitlab::Database::LoadBalancing.enabled = !Gitlab::Runtime.rake?
-Gitlab::Database::LoadBalancing.default_pool_size = Gitlab::Database.default_pool_size
-
-Gitlab::Cluster::LifecycleEvents.on_worker_start do
-  Gitlab::Database::LoadBalancing.enabled = !Gitlab::Runtime.rake?
-  Gitlab::Database::LoadBalancing.default_pool_size = Gitlab::Database.default_pool_size
-end
-
-Gitlab::Database::LoadBalancing.base_models =
-  ::Gitlab::Database.database_base_models_using_load_balancing.values.freeze
-
-# Configure callbacks
+# Configure Load Balancing interface callbacks
 Gitlab::Database::LoadBalancing::Callbacks.configure! do |callbacks|
   hosts_gauge = Gitlab::Metrics.gauge(:db_load_balancing_hosts, 'Current number of load balancing hosts')
 
-  callbacks.logger_proc = -> do
-    Gitlab::Database::LoadBalancing::Logger.build
-  end
-  callbacks.metrics_host_gauge_proc = ->(labels, value) do
-    hosts_gauge.set(labels, value)
-  end
-  callbacks.track_exception_proc = ->(exception) do
-    Gitlab::ErrorTracking.track_exception(exception)
-  end
+  callbacks.logger_proc = -> { Gitlab::Database::LoadBalancing::Logger.build }
+  callbacks.metrics_host_gauge_proc = ->(labels, value) { hosts_gauge.set(labels, value) }
+  callbacks.track_exception_proc = ->(exception) { Gitlab::ErrorTracking.track_exception(exception) }
+  callbacks.get_wal_for_proc = ->(key) { Gitlab::Redis::DbLoadBalancing.with { |redis| redis.get(key) } }
+  callbacks.del_wal_for_proc = ->(key) { Gitlab::Redis::DbLoadBalancing.with { |redis| redis.del(key) } }
   callbacks.set_wal_for_proc = ->(key, value, ex:) do
     Gitlab::Redis::DbLoadBalancing.with { |redis| redis.set(key, value, ex: ex) }
   end
-  callbacks.get_wal_for_proc = ->(key) do
-    Gitlab::Redis::DbLoadBalancing.with { |redis| redis.get(key) }
-  end
-  callbacks.del_wal_for_proc = ->(key) do
-    Gitlab::Redis::DbLoadBalancing.with { |redis| redis.del(key) }
+end
+
+Gitlab::Application.configure do |config|
+  config.middleware.use(Gitlab::Database::LoadBalancing::RackMiddleware)
+
+  # We need re-rerun the setup when code reloads in development
+  config.reloader.to_prepare do
+    if Gitlab.dev_or_test_env?
+      Gitlab::Database::LoadBalancing.base_models.each do |model|
+        Gitlab::Database::LoadBalancing::Setup.new(model).setup
+      end
+    end
   end
 end
 
diff --git a/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing.rb b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing.rb
index c35712d8a7b3cd..083af88d7698e1 100644
--- a/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing.rb
+++ b/gems/gitlab-database-load_balancing/lib/gitlab/database/load_balancing.rb
@@ -32,6 +32,10 @@ module LoadBalancing
       mattr_accessor :default_pool_size
       mattr_accessor :enabled, default: true
 
+      def self.configure!
+        yield(self)
+      end
+
       def self.each_load_balancer
         return to_enum(__method__) unless block_given?
 
diff --git a/spec/initializers/load_balancing_spec.rb b/spec/initializers/load_balancing_spec.rb
index db572708d53a77..3512326d0de497 100644
--- a/spec/initializers/load_balancing_spec.rb
+++ b/spec/initializers/load_balancing_spec.rb
@@ -104,10 +104,7 @@ def simulate_puma_worker
         models = described_class.base_models
 
         expect(models).to include(ActiveRecord::Base)
-
-        if Gitlab::Database.has_config?(:ci)
-          expect(models).to include(Ci::ApplicationRecord)
-        end
+        expect(models).to include(Ci::ApplicationRecord) if Gitlab::Database.has_config?(:ci)
       end
 
       it 'returns the models as a frozen array' do
-- 
GitLab


From 54acbace494c7123d41ee7e204b137f7f1c9acfb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kamil=20Trzci=C5=84ski?= <ayufan@ayufan.eu>
Date: Thu, 23 Nov 2023 13:38:57 +0100
Subject: [PATCH 6/8] Improve spec_helper

---
 .../gitlab-database-load_balancing/spec/spec_helper.rb | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/gems/gitlab-database-load_balancing/spec/spec_helper.rb b/gems/gitlab-database-load_balancing/spec/spec_helper.rb
index c60e502291be6f..835ce299e823c2 100644
--- a/gems/gitlab-database-load_balancing/spec/spec_helper.rb
+++ b/gems/gitlab-database-load_balancing/spec/spec_helper.rb
@@ -32,11 +32,13 @@
   end
 
   config.before(:all) do
-    Gitlab::Database::LoadBalancing.default_pool_size = 20
-    Gitlab::Database::LoadBalancing.base_models = [ActiveRecord::Base, Ci::ApplicationRecord]
+    Gitlab::Database::LoadBalancing.configure! do |load_balancing|
+      load_balancing.default_pool_size = 20
+      load_balancing.base_models = [ActiveRecord::Base, Ci::ApplicationRecord]
 
-    Gitlab::Database::LoadBalancing.base_models.each do |model|
-      Gitlab::Database::LoadBalancing::Setup.new(model).setup
+      load_balancing.base_models.each do |model|
+        Gitlab::Database::LoadBalancing::Setup.new(model).setup
+      end
     end
   end
 end
-- 
GitLab


From 941388451f3e5ec97521c48d2a39552cedd8a137 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kamil=20Trzci=C5=84ski?= <ayufan@ayufan.eu>
Date: Thu, 23 Nov 2023 14:32:15 +0000
Subject: [PATCH 7/8] Improve initializer specs

---
 Gemfile                                  |  2 +-
 spec/initializers/load_balancing_spec.rb | 69 ++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/Gemfile b/Gemfile
index b4ac85e747b2be..f92bc38598820e 100644
--- a/Gemfile
+++ b/Gemfile
@@ -41,7 +41,7 @@ gem 'gitlab-safe_request_store', path: 'gems/gitlab-safe_request_store' # ruboco
 group :monorepo do
   gem 'gitlab-utils', path: 'gems/gitlab-utils' # rubocop:todo Gemfile/MissingFeatureCategory
   gem 'gitlab-backup-cli', path: 'gems/gitlab-backup-cli', feature_category: :backup_restore
-  gem 'gitlab-database-load_balancing', path: 'gems/gitlab-database-load_balancing' # rubocop:todo Gemfile/MissingFeatureCategory
+  gem 'gitlab-database-load_balancing', path: 'gems/gitlab-database-load_balancing', feature_category: :database
 end
 
 gem 'gitlab-secret_detection', path: 'gems/gitlab-secret_detection', feature_category: :secret_detection
diff --git a/spec/initializers/load_balancing_spec.rb b/spec/initializers/load_balancing_spec.rb
index 3512326d0de497..11ab5ba91fb795 100644
--- a/spec/initializers/load_balancing_spec.rb
+++ b/spec/initializers/load_balancing_spec.rb
@@ -112,4 +112,73 @@ def simulate_puma_worker
       end
     end
   end
+
+  describe Gitlab::Database::LoadBalancing::Callbacks do
+    describe '.metrics_host_gauge' do
+      it 'sets prometheus metric' do
+        expect(described_class).to receive(:metrics_host_gauge)
+          .with({}, 1).and_call_original
+
+        # Metric will be set once configured host list:
+        # - 1 host since we only have primary
+        # - this test will fail if running CI against replicas
+        model = Gitlab::Database::LoadBalancing.base_models.first
+        Gitlab::Database::LoadBalancing::Setup.new(model).setup
+
+        metric = ::Prometheus::Client.registry.get(:db_load_balancing_hosts)
+        expect(metric).not_to be_nil
+        expect(metric.values[{}].get).to eq(1)
+      end
+    end
+
+    context 'when sticking connection', :database_replica, :redis do
+      let(:key) { 'database-load-balancing/write-location/main/user/100' }
+      let(:wal) { '0/D525E3A8' }
+
+      it 'does set latest WAL and unstick once caught-up' do
+        allow(ApplicationRecord.load_balancer).to receive(:primary_write_location).and_return(wal)
+
+        expect(described_class).to receive(:set_wal_for)
+          .with(key, wal, ex: 30)
+          .and_call_original
+
+        expect { ApplicationRecord.sticking.stick(:user, 100) }
+          .to change { described_class.get_wal_for(key) }.from(nil).to(wal)
+
+        expect(described_class).to receive(:del_wal_for)
+          .with(key).and_call_original
+
+        expect { ApplicationRecord.sticking.find_caught_up_replica(:user, 100) }
+          .to change { described_class.get_wal_for(key) }.from(wal).to(nil)
+      end
+
+      describe '.set_wal_for' do
+        it 'does get latest set value' do
+          expect { described_class.set_wal_for(key, wal, ex: 30) }
+            .to change { described_class.get_wal_for(key) }.from(nil).to(wal)
+        end
+      end
+
+      describe '.del_wal_for' do
+        it 'does delete set value' do
+          described_class.set_wal_for(key, wal, ex: 30)
+
+          expect { described_class.del_wal_for(key) }
+            .to change { described_class.get_wal_for(key) }.from(wal).to(nil)
+        end
+      end
+    end
+
+    describe '.track_exception' do
+      it 'forwards exception to ErrorTracking' do
+        expect(Gitlab::ErrorTracking).to receive(:track_exception)
+          .with(Gitlab::Utils::ConcurrentRubyThreadIsUsedError)
+
+        Gitlab::Utils.restrict_within_concurrent_ruby do
+          expect { ApplicationRecord.connection.execute("SELECT 1") }
+            .to raise_error(Gitlab::Utils::ConcurrentRubyThreadIsUsedError)
+        end
+      end
+    end
+  end
 end
-- 
GitLab


From 34a76c2e4e2ef895201c843a58e310d82e885849 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kamil=20Trzci=C5=84ski?= <ayufan@ayufan.eu>
Date: Tue, 28 Nov 2023 11:51:55 +0100
Subject: [PATCH 8/8] Fix rspec

---
 spec/initializers/load_balancing_spec.rb | 25 ++++++++++++------------
 1 file changed, 13 insertions(+), 12 deletions(-)

diff --git a/spec/initializers/load_balancing_spec.rb b/spec/initializers/load_balancing_spec.rb
index 11ab5ba91fb795..2f7fd712f2b6a5 100644
--- a/spec/initializers/load_balancing_spec.rb
+++ b/spec/initializers/load_balancing_spec.rb
@@ -2,18 +2,16 @@
 
 require 'spec_helper'
 
-RSpec.describe 'load_balancing', :delete, :reestablished_active_record_base, feature_category: :cell do
-  subject(:initialize_load_balancer) do
-    load Rails.root.join('config/initializers/load_balancing.rb')
-  end
-
-  before do
-    # Stub out middleware call, as not idempotent
-    allow(Gitlab::Application.instance.middleware).to receive(:use)
-  end
+RSpec.describe 'load_balancing', feature_category: :cell do
+  context 'with replica hosts configured', :delete, :reestablished_active_record_base do
+    subject(:initialize_load_balancer) do
+      load Rails.root.join('config/initializers/load_balancing.rb')
+    end
 
-  context 'with replica hosts configured' do
     before do
+      # Stub out middleware call, as not idempotent
+      allow(Gitlab::Application.instance.middleware).to receive(:use)
+
       # Setup host-based load balancing
       # Patch in our load balancer config, simply pointing at the test database twice
       allow(Gitlab::Database::LoadBalancing::Configuration).to receive(:for_model) do |base_model|
@@ -115,7 +113,7 @@ def simulate_puma_worker
 
   describe Gitlab::Database::LoadBalancing::Callbacks do
     describe '.metrics_host_gauge' do
-      it 'sets prometheus metric' do
+      it 'sets prometheus metric', :reestablished_active_record_base do
         expect(described_class).to receive(:metrics_host_gauge)
           .with({}, 1).and_call_original
 
@@ -131,7 +129,7 @@ def simulate_puma_worker
       end
     end
 
-    context 'when sticking connection', :database_replica, :redis do
+    context 'when sticking connection', :database_replica, :redis, :request_store do
       let(:key) { 'database-load-balancing/write-location/main/user/100' }
       let(:wal) { '0/D525E3A8' }
 
@@ -148,6 +146,9 @@ def simulate_puma_worker
         expect(described_class).to receive(:del_wal_for)
           .with(key).and_call_original
 
+        expect(ApplicationRecord.load_balancer).to receive(:select_up_to_date_host)
+          .with(wal).and_return(Gitlab::Database::LoadBalancing::LoadBalancer::ALL_CAUGHT_UP)
+
         expect { ApplicationRecord.sticking.find_caught_up_replica(:user, 100) }
           .to change { described_class.get_wal_for(key) }.from(wal).to(nil)
       end
-- 
GitLab