Make `LoadBalancer` to configure pool sizes of all classes

This changes `database_config.rb` and moves connection
pool size configuration to `Load Balancer`.

Changelog: fixed

config/initializers/database_config.rb

@@ -10,7 +10,9 @@ Gitlab.ee do
   end
 end
 
-ActiveRecord::Base.establish_connection(Gitlab::Database.main.db_config_with_default_pool_size)
+unless Gitlab::Utils.to_boolean(ENV["GITLAB_LB_CONFIGURE_CONNECTION"], default: true)
+  ActiveRecord::Base.establish_connection(Gitlab::Database.main.db_config_with_default_pool_size)
+end
 
 Gitlab.ee do
   if Gitlab::Runtime.sidekiq? && Gitlab::Geo.geo_database_configured?

lib/gitlab/database/connection.rb

@@ -59,6 +59,7 @@ module Gitlab
         adapter_name.casecmp('postgresql') == 0
       end
 
+      # TODO: To be removed with GITLAB_LB_CONFIGURE_CONNECTION
       def db_config_with_default_pool_size
         db_config_object = scope.connection_db_config
         config = db_config_object
@@ -72,20 +73,6 @@ module Gitlab
         )
       end
 
-      # Disables prepared statements for the current database connection.
-      def disable_prepared_statements
-        db_config_object = scope.connection_db_config
-        config = db_config_object.configuration_hash.merge(prepared_statements: false)
-
-        hash_config = ActiveRecord::DatabaseConfigurations::HashConfig.new(
-          db_config_object.env_name,
-          db_config_object.name,
-          config
-        )
-
-        scope.establish_connection(hash_config)
-      end
-
       # Check whether the underlying database is in read-only mode
       def db_read_only?
         pg_is_in_recovery =

lib/gitlab/database/load_balancing/setup.rb

@@ -14,20 +14,30 @@ module Gitlab
         end
 
         def setup
-          disable_prepared_statements
+          configure_connection
           setup_load_balancer
           setup_service_discovery
         end
 
-        def disable_prepared_statements
+        def configure_connection
           db_config_object = @model.connection_db_config
-          config =
-            db_config_object.configuration_hash.merge(prepared_statements: false)
+
+          hash =
+            if Gitlab::Utils.to_boolean(ENV["GITLAB_LB_CONFIGURE_CONNECTION"], default: true)
+              db_config_object.configuration_hash.merge(
+                prepared_statements: false,
+                pool: Gitlab::Database.default_pool_size
+              )
+            else
+              db_config_object.configuration_hash.merge(
+                prepared_statements: false
+              )
+            end
 
           hash_config = ActiveRecord::DatabaseConfigurations::HashConfig.new(
             db_config_object.env_name,
             db_config_object.name,
-            config
+            hash
           )
 
           @model.establish_connection(hash_config)

spec/initializers/database_config_spec.rb

@@ -9,6 +9,7 @@ RSpec.describe 'Database config initializer', :reestablished_active_record_base
 
   before do
     allow(Gitlab::Runtime).to receive(:max_threads).and_return(max_threads)
+    stub_env('GITLAB_LB_CONFIGURE_CONNECTION', 'false')
   end
 
   let(:max_threads) { 8 }
@@ -41,6 +42,16 @@ RSpec.describe 'Database config initializer', :reestablished_active_record_base
 
       expect(ActiveRecord::Base.connection_db_config.pool).to eq(18)
     end
+
+    context 'when GITLAB_LB_CONFIGURE_CONNECTION=true' do
+      before do
+        stub_env('GITLAB_LB_CONFIGURE_CONNECTION', 'true')
+      end
+
+      it 'does not overwrite custom pool settings' do
+        expect { subject }.not_to change { ActiveRecord::Base.connection_db_config.pool }
+      end
+    end
   end
 
   context "when specifying headroom through an ENV variable" do

spec/lib/gitlab/database/connection_spec.rb

@@ -126,39 +126,6 @@ RSpec.describe Gitlab::Database::Connection do
     end
   end
 
-  describe '#disable_prepared_statements', :reestablished_active_record_base do
-    it 'disables prepared statements' do
-      connection.scope.establish_connection(
-        ::Gitlab::Database.main.config.merge(prepared_statements: true)
-      )
-
-      expect(connection.scope.connection.prepared_statements).to eq(true)
-
-      connection.disable_prepared_statements
-
-      expect(connection.scope.connection.prepared_statements).to eq(false)
-    end
-
-    it 'retains the connection name' do
-      connection.disable_prepared_statements
-
-      expect(connection.scope.connection_db_config.name).to eq('main')
-    end
-
-    context 'with dynamic connection pool size' do
-      before do
-        connection.scope.establish_connection(connection.config.merge(pool: 7))
-      end
-
-      it 'retains the set pool size' do
-        connection.disable_prepared_statements
-
-        expect(connection.scope.connection.prepared_statements).to eq(false)
-        expect(connection.scope.connection.pool.size).to eq(7)
-      end
-    end
-  end
-
   describe '#db_read_only?' do
     it 'detects a read-only database' do
       allow(connection.scope.connection)

spec/lib/gitlab/database/load_balancing/setup_spec.rb

@@ -7,7 +7,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
     it 'sets up the load balancer' do
       setup = described_class.new(ActiveRecord::Base)
 
-      expect(setup).to receive(:disable_prepared_statements)
+      expect(setup).to receive(:configure_connection)
       expect(setup).to receive(:setup_load_balancer)
       expect(setup).to receive(:setup_service_discovery)
 
@@ -15,11 +15,11 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
     end
   end
 
-  describe '#disable_prepared_statements' do
-    it 'disables prepared statements and reconnects to the database' do
+  describe '#configure_connection' do
+    it 'configures pool, prepared statements and reconnects to the database' do
       config = double(
         :config,
-        configuration_hash: { host: 'localhost' },
+        configuration_hash: { host: 'localhost', pool: 2, prepared_statements: true },
         env_name: 'test',
         name: 'main'
       )
@@ -27,7 +27,11 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
 
       expect(ActiveRecord::DatabaseConfigurations::HashConfig)
         .to receive(:new)
-        .with('test', 'main', { host: 'localhost', prepared_statements: false })
+        .with('test', 'main', {
+          host: 'localhost',
+          prepared_statements: false,
+          pool: Gitlab::Database.default_pool_size
+        })
         .and_call_original
 
       # HashConfig doesn't implement its own #==, so we can't directly compare
@@ -36,7 +40,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
         .to receive(:establish_connection)
         .with(an_instance_of(ActiveRecord::DatabaseConfigurations::HashConfig))
 
-      described_class.new(model).disable_prepared_statements
+      described_class.new(model).configure_connection
     end
   end