Simplify DB connection pool size configuration (optional)

config/initializers/database_config.rb

@@ -20,11 +20,44 @@ Gitlab.ee do
   end
 end
 
-# When running on multi-threaded runtimes like Puma or Sidekiq,
-# set the number of threads per process as the minimum DB connection pool size.
-# This is to avoid connectivity issues as was documented here:
-# https://github.com/rails/rails/pull/23057
-if Gitlab::Runtime.multi_threaded?
+# TODO get rid of feature flag https://gitlab.com/gitlab-com/gl-infra/scalability/-/issues/495
+if ENV['DB_USE_NEW_POOL_SIZE_LOGIC'] == '1'
+  # Because of the way Ruby on Rails manages database connections, it is
+  # important that we have at least as many connections as we have
+  # threads. While there is a 'pool' setting in database.yml, it is not
+  # very practical because you need to maintain it in tandem with the
+  # number of application threads. Because of this we override the number
+  # of allowed connections in the database connection pool based on the
+  # configured number of application threads.
+  #
+  # Gitlab::Runtime.max_threads is the number of "user facing" application
+  # threads the process has been configured with. We also have auxiliary
+  # threads that use database connections. Because it is not practical to
+  # keep an accurate count of the number auxiliary threads as the
+  # application evolves over time, we just add a fixed headroom to the
+  # number of user-facing threads. It is OK if this number is too large
+  # because connections are instantiated lazily.
+
+  headroom = (ENV["DB_POOL_HEADROOM"].presence || 10).to_i
+  calculated_pool_size = Gitlab::Runtime.max_threads + headroom
+
+  db_config = Gitlab::Database.config ||
+      Rails.application.config.database_configuration[Rails.env]
+
+  db_config['pool'] = calculated_pool_size
+  ActiveRecord::Base.establish_connection(db_config)
+
+  Gitlab.ee do
+    if Gitlab::Runtime.sidekiq? && Gitlab::Geo.geo_database_configured?
+      Rails.configuration.geo_database['pool'] = calculated_pool_size
+      Geo::TrackingBase.establish_connection(Rails.configuration.geo_database)
+    end
+  end
+elsif Gitlab::Runtime.multi_threaded?
+  # When running on multi-threaded runtimes like Puma or Sidekiq,
+  # set the number of threads per process as the minimum DB connection pool size.
+  # This is to avoid connectivity issues as was documented here:
+  # https://github.com/rails/rails/pull/23057
   max_threads = Gitlab::Runtime.max_threads
   db_config = Gitlab::Database.config ||
       Rails.application.config.database_configuration[Rails.env]

ee/spec/initializers/database_config_spec.rb

@@ -39,6 +39,28 @@ RSpec.describe 'Database config initializer for GitLab EE' do
     end
   end
 
+  context "with new pool size logic" do
+    let(:max_threads) { 8 }
+
+    before do
+      stub_env('DB_USE_NEW_POOL_SIZE_LOGIC', '1')
+      allow(Gitlab::Runtime).to receive(:max_threads).and_return(max_threads)
+      allow(ActiveRecord::Base).to receive(:establish_connection)
+
+      expect(Geo::TrackingBase).to receive(:establish_connection)
+    end
+
+    context "and the runtime is Sidekiq" do
+      before do
+        allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(true)
+      end
+
+      it "sets Geo DB connection pool size to the max number of worker threads" do
+        expect { subject }.to change { Rails.configuration.geo_database['pool'] }.from(1).to(18)
+      end
+    end
+  end
+
   def stub_geo_database_config(pool_size:)
     config = {
       'adapter' => 'postgresql',

spec/initializers/database_config_spec.rb

@@ -71,6 +71,60 @@ RSpec.describe 'Database config initializer' do
     end
   end
 
+  context "with new pool size logic" do
+    let(:max_threads) { 8 }
+
+    before do
+      stub_env('DB_USE_NEW_POOL_SIZE_LOGIC', '1')
+      allow(Gitlab::Runtime).to receive(:max_threads).and_return(max_threads)
+    end
+
+    context "and no existing pool size is set" do
+      before do
+        stub_database_config(pool_size: nil)
+      end
+
+      it "sets it based on the max number of worker threads" do
+        expect { subject }.to change { Gitlab::Database.config['pool'] }.from(nil).to(18)
+      end
+    end
+
+    context "and the existing pool size is smaller than the max number of worker threads" do
+      before do
+        stub_database_config(pool_size: 1)
+      end
+
+      it "sets it based on the max number of worker threads" do
+        expect { subject }.to change { Gitlab::Database.config['pool'] }.from(1).to(18)
+      end
+    end
+
+    context "and the existing pool size is larger than the max number of worker threads" do
+      before do
+        stub_database_config(pool_size: 100)
+      end
+
+      it "sets it based on the max number of worker threads" do
+        expect { subject }.to change { Gitlab::Database.config['pool'] }.from(100).to(18)
+      end
+    end
+
+    context "when specifying headroom through an ENV variable" do
+      let(:headroom) { 15 }
+
+      before do
+        stub_database_config(pool_size: 1)
+        stub_env("DB_POOL_HEADROOM", headroom)
+      end
+
+      it "adds headroom on top of the calculated size" do
+        expect { subject }.to change { Gitlab::Database.config['pool'] }
+                                .from(1)
+                                .to(max_threads + headroom)
+      end
+    end
+  end
+
   def stub_database_config(pool_size:)
     config = {
       'adapter' => 'postgresql',