Add latest changes from gitlab-org/gitlab@master

4 files changed, 46 insertions, 26 deletions

diff --git a/config/initializers/7_prometheus_metrics.rb b/config/initializers/7_prometheus_metrics.rb
index d2d546a5438..587d393fd77 100644
--- a/config/initializers/7_prometheus_metrics.rb
+++ b/config/initializers/7_prometheus_metrics.rb
@@ -48,7 +48,7 @@ if !Rails.env.test? && Gitlab::Metrics.prometheus_metrics_enabled?
 
     Gitlab::Metrics.gauge(:deployments, 'GitLab Version', {}, :max).set({ version: Gitlab::VERSION, revision: Gitlab.revision }, 1)
 
-    unless Gitlab::Runtime.sidekiq?
+    if Gitlab::Runtime.web_server?
       Gitlab::Metrics::RequestsRackMiddleware.initialize_metrics
     end
 
diff --git a/config/initializers/action_cable.rb b/config/initializers/action_cable.rb
index 16d29f5910f..a7ef5cc332c 100644
--- a/config/initializers/action_cable.rb
+++ b/config/initializers/action_cable.rb
@@ -19,4 +19,4 @@ ActionCable::SubscriptionAdapter::Redis.redis_connector = lambda do |config|
 end
 
 Gitlab::ActionCable::RequestStoreCallbacks.install
-Gitlab::Database::LoadBalancing::ActionCableCallbacks.install if Gitlab::Database::LoadBalancing.enable?
+Gitlab::Database::LoadBalancing::ActionCableCallbacks.install
diff --git a/config/initializers/load_balancing.rb b/config/initializers/load_balancing.rb
index 2b58ae0f642..29700b37c7d 100644
--- a/config/initializers/load_balancing.rb
+++ b/config/initializers/load_balancing.rb
@@ -2,27 +2,51 @@
 
 ActiveRecord::Base.singleton_class.attr_accessor :load_balancing_proxy
 
-if Gitlab::Database::LoadBalancing.enable?
-  Gitlab::Database.main.disable_prepared_statements
+Gitlab::Database.main.disable_prepared_statements
 
-  Gitlab::Application.configure do |config|
-    config.middleware.use(Gitlab::Database::LoadBalancing::RackMiddleware)
-  end
+Gitlab::Application.configure do |config|
+  config.middleware.use(Gitlab::Database::LoadBalancing::RackMiddleware)
+end
+
+# This hijacks the "connection" method to ensure both
+# `ActiveRecord::Base.connection` and all models use the same load
+# balancing proxy.
+ActiveRecord::Base.singleton_class.prepend(Gitlab::Database::LoadBalancing::ActiveRecordProxy)
+
+# The load balancer needs to be configured immediately, and re-configured after
+# forking. This ensures queries that run before forking use the load balancer,
+# and queries running after a fork don't run into any errors when using dead
+# database connections.
+#
+# See https://gitlab.com/gitlab-org/gitlab/-/merge_requests/63485 for more
+# information.
+setup = proc do
+  lb = Gitlab::Database::LoadBalancing::LoadBalancer.new(
+    Gitlab::Database::LoadBalancing.configuration,
+    primary_only: !Gitlab::Database::LoadBalancing.enable_replicas?
+  )
 
-  # This hijacks the "connection" method to ensure both
-  # `ActiveRecord::Base.connection` and all models use the same load
-  # balancing proxy.
-  ActiveRecord::Base.singleton_class.prepend(Gitlab::Database::LoadBalancing::ActiveRecordProxy)
+  ActiveRecord::Base.load_balancing_proxy =
+    Gitlab::Database::LoadBalancing::ConnectionProxy.new(lb)
 
-  Gitlab::Database::LoadBalancing.configure_proxy
+  # Populate service discovery immediately if it is configured
+  Gitlab::Database::LoadBalancing.perform_service_discovery
+end
+
+setup.call
 
-  # This needs to be executed after fork of clustered processes
-  Gitlab::Cluster::LifecycleEvents.on_worker_start do
-    # For Host-based LB, we need to re-connect as Rails discards connections on fork
-    Gitlab::Database::LoadBalancing.configure_proxy
+# Database queries may be run before we fork, so we must set up the load
+# balancer as early as possible. When we do fork, we need to make sure all the
+# hosts are disconnected.
+Gitlab::Cluster::LifecycleEvents.on_before_fork do
+  # When forking, we don't want to wait until the connections aren't in use any
+  # more, as this could delay the boot cycle.
+  Gitlab::Database::LoadBalancing.proxy.load_balancer.disconnect!(timeout: 0)
+end
 
-    # Service discovery must be started after configuring the proxy, as service
-    # discovery depends on this.
-    Gitlab::Database::LoadBalancing.start_service_discovery
-  end
+# Service discovery only needs to run in the worker processes, as the main one
+# won't be running many (if any) database queries.
+Gitlab::Cluster::LifecycleEvents.on_worker_start do
+  setup.call
+  Gitlab::Database::LoadBalancing.start_service_discovery
 end
diff --git a/config/initializers/zz_metrics.rb b/config/initializers/zz_metrics.rb
index 0fa67ea9f8c..8e69e1634f1 100644
--- a/config/initializers/zz_metrics.rb
+++ b/config/initializers/zz_metrics.rb
@@ -29,12 +29,8 @@ if Gitlab::Metrics.enabled? && !Rails.env.test? && !(Rails.env.development? && d
   Gitlab::Application.configure do |config|
     # We want to track certain metrics during the Load Balancing host resolving process.
     # Because of that, we need to have metrics code available earlier for Load Balancing.
-    if Gitlab::Database::LoadBalancing.enable?
-      config.middleware.insert_before Gitlab::Database::LoadBalancing::RackMiddleware,
-        Gitlab::Metrics::RackMiddleware
-    else
-      config.middleware.use(Gitlab::Metrics::RackMiddleware)
-    end
+    config.middleware.insert_before Gitlab::Database::LoadBalancing::RackMiddleware,
+      Gitlab::Metrics::RackMiddleware
 
     config.middleware.use(Gitlab::Middleware::RailsQueueDuration)
     config.middleware.use(Gitlab::Metrics::ElasticsearchRackMiddleware)