Consolidate DB pool size settings into initializer

This merges settings from the existing initializer
and the sidekiq initializer into a single file.

Logs any forced increases in db pool size.

config/initializers/database_config.rb

 # frozen_string_literal: true
 
-# when running on puma, scale connection pool size with the number
-# of threads per worker process
-if Gitlab::Runtime.puma?
+def log_pool_size(db, previous_pool_size, current_pool_size)
+  log_message = ["#{db} connection pool size: #{current_pool_size}"]
+
+  if previous_pool_size && current_pool_size > previous_pool_size
+    log_message << "(increased from #{previous_pool_size} to match thread count)"
+  end
+
+  Gitlab::AppLogger.debug(log_message.join(' '))
+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?
+  max_threads = Gitlab::Runtime.max_threads
   db_config = Gitlab::Database.config ||
       Rails.application.config.database_configuration[Rails.env]
-  puma_options = Puma.cli_config.options
+  previous_db_pool_size = db_config['pool']
 
-  # We use either the maximum number of threads per worker process, or
-  # the user specified value, whichever is larger.
-  desired_pool_size = [db_config['pool'].to_i, puma_options[:max_threads]].max
+  db_config['pool'] = [db_config['pool'].to_i, max_threads].max
 
-  db_config['pool'] = desired_pool_size
-
-  # recreate the connection pool from the new config
   ActiveRecord::Base.establish_connection(db_config)
+
+  current_db_pool_size = ActiveRecord::Base.connection.pool.size
+
+  log_pool_size('DB', previous_db_pool_size, current_db_pool_size)
+
+  Gitlab.ee do
+    if Gitlab::Runtime.sidekiq? && Gitlab::Geo.geo_database_configured?
+      previous_geo_db_pool_size = Rails.configuration.geo_database['pool']
+      Rails.configuration.geo_database['pool'] = max_threads
+      Geo::TrackingBase.establish_connection(Rails.configuration.geo_database)
+      current_geo_db_pool_size = Geo::TrackingBase.connection_pool.size
+      log_pool_size('Geo DB', previous_geo_db_pool_size, current_geo_db_pool_size)
+    end
+  end
 end

config/initializers/sidekiq.rb

@@ -88,23 +88,10 @@ Sidekiq.configure_server do |config|
 
   Gitlab::SidekiqVersioning.install!
 
-  db_config = Gitlab::Database.config ||
-    Rails.application.config.database_configuration[Rails.env]
-  db_config['pool'] = Sidekiq.options[:concurrency]
-  ActiveRecord::Base.establish_connection(db_config)
-  Rails.logger.debug("Connection Pool size for Sidekiq Server is now: #{ActiveRecord::Base.connection.pool.instance_variable_get('@size')}") # rubocop:disable Gitlab/RailsLogger
-
   Gitlab.ee do
     Gitlab::Mirror.configure_cron_job!
 
     Gitlab::Geo.configure_cron_jobs!
-
-    if Gitlab::Geo.geo_database_configured?
-      Rails.configuration.geo_database['pool'] = Sidekiq.options[:concurrency]
-      Geo::TrackingBase.establish_connection(Rails.configuration.geo_database)
-
-      Rails.logger.debug("Connection Pool size for Sidekiq Server is now: #{Geo::TrackingBase.connection_pool.size} (Geo tracking database)") # rubocop:disable Gitlab/RailsLogger
-    end
   end
 
   # Avoid autoload issue such as 'Mail::Parsers::AddressStruct'

ee/spec/initializers/database_config_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe 'Database config initializer for GitLab EE' do
+  subject do
+    load Rails.root.join('config/initializers/database_config.rb')
+  end
+
+  before do
+    stub_geo_database_config(pool_size: 1)
+  end
+
+  context "when using multi-threaded runtime" do
+    let(:max_threads) { 8 }
+
+    before do
+      allow(Gitlab::Runtime).to receive(:multi_threaded?).and_return(true)
+      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(max_threads)
+      end
+    end
+  end
+
+  context "when using single-threaded runtime" do
+    it "does nothing" do
+      expect { subject }.not_to change { Rails.configuration.geo_database['pool'] }
+    end
+  end
+
+  def stub_geo_database_config(pool_size:)
+    config = {
+      'adapter' => 'postgresql',
+      'host' => 'db.host.com',
+      'pool' => pool_size
+    }.compact
+
+    allow(Rails.configuration).to receive(:geo_database).and_return(config)
+  end
+end

lib/gitlab/runtime.rb

@@ -67,6 +67,16 @@ module Gitlab
       def process_name
         File.basename($0)
       end
+
+      def max_threads
+        if puma?
+          Puma.cli_config.options[:max_threads]
+        elsif sidekiq?
+          Sidekiq.options[:concurrency]
+        else
+          1
+        end
+      end
     end
   end
 end

spec/initializers/database_config_spec.rb

@@ -11,14 +11,12 @@ describe 'Database config initializer' do
     allow(ActiveRecord::Base).to receive(:establish_connection)
   end
 
-  context "when using Puma" do
-    let(:puma) { double('puma') }
-    let(:puma_options) { { max_threads: 8 } }
+  context "when using multi-threaded runtime" do
+    let(:max_threads) { 8 }
 
     before do
-      allow(Gitlab::Runtime).to receive(:puma?).and_return(true)
-      stub_const("Puma", puma)
-      allow(puma).to receive_message_chain(:cli_config, :options).and_return(puma_options)
+      allow(Gitlab::Runtime).to receive(:multi_threaded?).and_return(true)
+      allow(Gitlab::Runtime).to receive(:max_threads).and_return(max_threads)
     end
 
     context "and no existing pool size is set" do
@@ -27,23 +25,23 @@ describe 'Database config initializer' do
       end
 
       it "sets it to the max number of worker threads" do
-        expect { subject }.to change { Gitlab::Database.config['pool'] }.from(nil).to(8)
+        expect { subject }.to change { Gitlab::Database.config['pool'] }.from(nil).to(max_threads)
       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: 7)
+        stub_database_config(pool_size: max_threads - 1)
       end
 
       it "sets it to the max number of worker threads" do
-        expect { subject }.to change { Gitlab::Database.config['pool'] }.from(7).to(8)
+        expect { subject }.to change { Gitlab::Database.config['pool'] }.by(1)
       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: 9)
+        stub_database_config(pool_size: max_threads + 1)
       end
 
       it "keeps the configured pool size" do
@@ -52,11 +50,7 @@ describe 'Database config initializer' do
     end
   end
 
-  context "when not using Puma" do
-    before do
-      stub_database_config(pool_size: 7)
-    end
-
+  context "when using single-threaded runtime" do
     it "does nothing" do
       expect { subject }.not_to change { Gitlab::Database.config['pool'] }
     end

spec/lib/gitlab/runtime_spec.rb

@@ -26,9 +26,15 @@ describe Gitlab::Runtime do
 
   context "puma" do
     let(:puma_type) { double('::Puma') }
+    let(:options) do
+      {
+        max_threads: 2
+      }
+    end
 
     before do
       stub_const('::Puma', puma_type)
+      allow(puma_type).to receive_message_chain(:cli_config, :options).and_return(options)
     end
 
     it "identifies itself" do
@@ -43,6 +49,10 @@ describe Gitlab::Runtime do
       expect(subject.rake?).to be(false)
       expect(subject.rspec?).to be(false)
     end
+
+    it "reports its maximum concurrency" do
+      expect(subject.max_threads).to eq(2)
+    end
   end
 
   context "unicorn" do
@@ -66,14 +76,24 @@ describe Gitlab::Runtime do
       expect(subject.rake?).to be(false)
       expect(subject.rspec?).to be(false)
     end
+
+    it "reports its maximum concurrency" do
+      expect(subject.max_threads).to eq(1)
+    end
   end
 
   context "sidekiq" do
     let(:sidekiq_type) { double('::Sidekiq') }
+    let(:options) do
+      {
+        concurrency: 2
+      }
+    end
 
     before do
       stub_const('::Sidekiq', sidekiq_type)
       allow(sidekiq_type).to receive(:server?).and_return(true)
+      allow(sidekiq_type).to receive(:options).and_return(options)
     end
 
     it "identifies itself" do
@@ -88,6 +108,10 @@ describe Gitlab::Runtime do
       expect(subject.rake?).to be(false)
       expect(subject.rspec?).to be(false)
     end
+
+    it "reports its maximum concurrency" do
+      expect(subject.max_threads).to eq(2)
+    end
   end
 
   context "console" do
@@ -109,6 +133,10 @@ describe Gitlab::Runtime do
       expect(subject.rake?).to be(false)
       expect(subject.rspec?).to be(false)
     end
+
+    it "reports its maximum concurrency" do
+      expect(subject.max_threads).to eq(1)
+    end
   end
 
   context "rspec" do
@@ -127,5 +155,9 @@ describe Gitlab::Runtime do
       expect(subject.rake?).to be(false)
       expect(subject.puma?).to be(false)
     end
+
+    it "reports its maximum concurrency" do
+      expect(subject.max_threads).to eq(1)
+    end
   end
 end