1 files changed, 24 insertions, 19 deletions
diff --git a/config/initializers/load_balancing.rb b/config/initializers/load_balancing.rb
index 2b58ae0f642..a31b11bb2be 100644
--- a/config/initializers/load_balancing.rb
+++ b/config/initializers/load_balancing.rb
@@ -1,28 +1,33 @@
 # frozen_string_literal: true
 
-ActiveRecord::Base.singleton_class.attr_accessor :load_balancing_proxy
+Gitlab::Application.configure do |config|
+  config.middleware.use(Gitlab::Database::LoadBalancing::RackMiddleware)
+end
 
-if Gitlab::Database::LoadBalancing.enable?
-  Gitlab::Database.main.disable_prepared_statements
+Gitlab::Database::LoadBalancing.base_models.each do |model|
+  # The load balancer needs to be configured immediately, and re-configured
+  # after forking. This ensures queries that run before forking use the load
+  # balancer, and queries running after a fork don't run into any errors when
+  # using dead database connections.
+  #
+  # See https://gitlab.com/gitlab-org/gitlab/-/merge_requests/63485 for more
+  # information.
+  Gitlab::Database::LoadBalancing::Setup.new(model).setup
 
-  Gitlab::Application.configure do |config|
-    config.middleware.use(Gitlab::Database::LoadBalancing::RackMiddleware)
+  # Database queries may be run before we fork, so we must set up the load
+  # balancer as early as possible. When we do fork, we need to make sure all the
+  # hosts are disconnected.
+  Gitlab::Cluster::LifecycleEvents.on_before_fork do
+    # When forking, we don't want to wait until the connections aren't in use
+    # any more, as this could delay the boot cycle.
+    model.connection.load_balancer.disconnect!(timeout: 0)
   end
 
-  # This hijacks the "connection" method to ensure both
-  # `ActiveRecord::Base.connection` and all models use the same load
-  # balancing proxy.
-  ActiveRecord::Base.singleton_class.prepend(Gitlab::Database::LoadBalancing::ActiveRecordProxy)
-
-  Gitlab::Database::LoadBalancing.configure_proxy
-
-  # This needs to be executed after fork of clustered processes
+  # Service discovery only needs to run in the worker processes, as the main one
+  # won't be running many (if any) database queries.
   Gitlab::Cluster::LifecycleEvents.on_worker_start do
-    # For Host-based LB, we need to re-connect as Rails discards connections on fork
-    Gitlab::Database::LoadBalancing.configure_proxy
-
-    # Service discovery must be started after configuring the proxy, as service
-    # discovery depends on this.
-    Gitlab::Database::LoadBalancing.start_service_discovery
+    Gitlab::Database::LoadBalancing::Setup
+      .new(model, start_service_discovery: true)
+      .setup
   end
 end