Remove `Gitlab::Database#create_connection_pool` and provide LB one

This changes how, and where the connection pool is created.
Since, we need Connection Pool be created in a context of LB
this moves this method to LB context, and properly mutates
config to also indicate a spec name, and replica usage.

lib/gitlab/database/connection.rb

@@ -152,20 +152,6 @@ module Gitlab
         end
       end
 
-      # pool_size - The size of the DB pool.
-      # host - An optional host name to use instead of the default one.
-      # port - An optional port to connect to.
-      def create_connection_pool(pool_size, host = nil, port = nil)
-        original_config = config
-        env_config = original_config.merge(pool: pool_size)
-
-        env_config[:host] = host if host
-        env_config[:port] = port if port
-
-        ActiveRecord::ConnectionAdapters::ConnectionHandler
-          .new.establish_connection(env_config)
-      end
-
       def cached_column_exists?(table_name, column_name)
         connection
           .schema_cache.columns_hash(table_name)

lib/gitlab/database/load_balancing/host.rb

@@ -29,7 +29,7 @@ module Gitlab
           @host = host
           @port = port
           @load_balancer = load_balancer
-          @pool = Database.main.create_connection_pool(LoadBalancing.pool_size, host, port)
+          @pool = load_balancer.create_replica_connection_pool(LoadBalancing.pool_size, host, port)
           @online = true
           @last_checked_at = Time.zone.now
 
@@ -41,7 +41,7 @@ module Gitlab
         #
         # timeout - The time after which the pool should be forcefully
         #           disconnected.
-        def disconnect!(timeout = 120)
+        def disconnect!(timeout: 120)
           start_time = ::Gitlab::Metrics::System.monotonic_time
 
           while (::Gitlab::Metrics::System.monotonic_time - start_time) <= timeout

lib/gitlab/database/load_balancing/load_balancer.rb

@@ -14,10 +14,14 @@ module Gitlab
 
         # hosts - The hostnames/addresses of the additional databases.
         def initialize(hosts = [], model = ActiveRecord::Base)
+          @model = model
           @host_list = HostList.new(hosts.map { |addr| Host.new(addr, self) })
           @connection_db_roles = {}.compare_by_identity
           @connection_db_roles_count = {}.compare_by_identity
-          @model = model
+        end
+
+        def disconnect!(timeout: 120)
+          host_list.hosts.each { |host| host.disconnect!(timeout: timeout) }
         end
 
         # Yields a connection that can be used for reads.
@@ -205,6 +209,30 @@ module Gitlab
           end
         end
 
+        # pool_size - The size of the DB pool.
+        # host - An optional host name to use instead of the default one.
+        # port - An optional port to connect to.
+        def create_replica_connection_pool(pool_size, host = nil, port = nil)
+          db_config = pool.db_config
+
+          env_config = db_config.configuration_hash.dup
+          env_config[:pool] = pool_size
+          env_config[:host] = host if host
+          env_config[:port] = port if port
+
+          replica_db_config = ActiveRecord::DatabaseConfigurations::HashConfig.new(
+            db_config.env_name,
+            db_config.name + "_replica",
+            env_config
+          )
+
+          # We cannot use ActiveRecord::Base.connection_handler.establish_connection
+          # as it will rewrite ActiveRecord::Base.connection
+          ActiveRecord::ConnectionAdapters::ConnectionHandler
+            .new
+            .establish_connection(replica_db_config)
+        end
+
         private
 
         # ActiveRecord::ConnectionAdapters::ConnectionHandler handles fetching,

lib/gitlab/database/load_balancing/service_discovery.rb

@@ -120,7 +120,7 @@ module Gitlab
           # host/connection. While this connection will be checked in and out,
           # it won't be explicitly disconnected.
           old_hosts.each do |host|
-            host.disconnect!(disconnect_timeout)
+            host.disconnect!(timeout: disconnect_timeout)
           end
         end
 

spec/lib/gitlab/database/connection_spec.rb

@@ -378,42 +378,6 @@ RSpec.describe Gitlab::Database::Connection do
     end
   end
 
-  describe '#create_connection_pool' do
-    it 'creates a new connection pool with specific pool size' do
-      pool = connection.create_connection_pool(5)
-
-      begin
-        expect(pool)
-          .to be_kind_of(ActiveRecord::ConnectionAdapters::ConnectionPool)
-
-        expect(pool.db_config.pool).to eq(5)
-      ensure
-        pool.disconnect!
-      end
-    end
-
-    it 'allows setting of a custom hostname' do
-      pool = connection.create_connection_pool(5, '127.0.0.1')
-
-      begin
-        expect(pool.db_config.host).to eq('127.0.0.1')
-      ensure
-        pool.disconnect!
-      end
-    end
-
-    it 'allows setting of a custom hostname and port' do
-      pool = connection.create_connection_pool(5, '127.0.0.1', 5432)
-
-      begin
-        expect(pool.db_config.host).to eq('127.0.0.1')
-        expect(pool.db_config.configuration_hash[:port]).to eq(5432)
-      ensure
-        pool.disconnect!
-      end
-    end
-  end
-
   describe '#cached_column_exists?' do
     it 'only retrieves data once' do
       expect(connection.scope.connection)

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

@@ -3,25 +3,17 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::LoadBalancing::HostList do
-  def expect_metrics(hosts)
-    expect(Gitlab::Metrics.registry.get(:db_load_balancing_hosts).get({})).to eq(hosts)
-  end
-
-  before do
-    allow(Gitlab::Database.main)
-      .to receive(:create_connection_pool)
-      .and_return(ActiveRecord::Base.connection_pool)
-  end
-
+  let(:db_host) { ActiveRecord::Base.connection_pool.db_config.host }
   let(:load_balancer) { double(:load_balancer) }
   let(:host_count) { 2 }
+  let(:hosts) { Array.new(host_count) { Gitlab::Database::LoadBalancing::Host.new(db_host, load_balancer, port: 5432) } }
+  let(:host_list) { described_class.new(hosts) }
 
-  let(:host_list) do
-    hosts = Array.new(host_count) do
-      Gitlab::Database::LoadBalancing::Host.new('localhost', load_balancer, port: 5432)
+  before do
+    # each call generate a new replica pool
+    allow(load_balancer).to receive(:create_replica_connection_pool) do
+      double(:replica_connection_pool)
     end
-
-    described_class.new(hosts)
   end
 
   describe '#initialize' do
@@ -42,8 +34,8 @@ RSpec.describe Gitlab::Database::LoadBalancing::HostList do
     context 'with ports' do
       it 'returns the host names of all hosts' do
         hosts = [
-          ['localhost', 5432],
-          ['localhost', 5432]
+          [db_host, 5432],
+          [db_host, 5432]
         ]
 
         expect(host_list.host_names_and_ports).to eq(hosts)
@@ -51,18 +43,12 @@ RSpec.describe Gitlab::Database::LoadBalancing::HostList do
     end
 
     context 'without ports' do
-      let(:host_list) do
-        hosts = Array.new(2) do
-          Gitlab::Database::LoadBalancing::Host.new('localhost', load_balancer)
-        end
-
-        described_class.new(hosts)
-      end
+      let(:hosts) { Array.new(2) { Gitlab::Database::LoadBalancing::Host.new(db_host, load_balancer) } }
 
       it 'returns the host names of all hosts' do
         hosts = [
-          ['localhost', nil],
-          ['localhost', nil]
+          [db_host, nil],
+          [db_host, nil]
         ]
 
         expect(host_list.host_names_and_ports).to eq(hosts)
@@ -71,10 +57,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::HostList do
   end
 
   describe '#manage_pool?' do
-    before do
-      allow(Gitlab::Database.main).to receive(:create_connection_pool) { double(:connection) }
-    end
-
     context 'when the testing pool belongs to one host of the host list' do
       it 'returns true' do
         pool = host_list.hosts.first.pool
@@ -185,4 +167,8 @@ RSpec.describe Gitlab::Database::LoadBalancing::HostList do
       end
     end
   end
+
+  def expect_metrics(hosts)
+    expect(Gitlab::Metrics.registry.get(:db_load_balancing_hosts).get({})).to eq(hosts)
+  end
 end

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

@@ -3,15 +3,16 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::LoadBalancing::Host do
-  let(:load_balancer) do
-    Gitlab::Database::LoadBalancing::LoadBalancer.new(%w[localhost])
-  end
+  let(:load_balancer) { Gitlab::Database::LoadBalancing::LoadBalancer.new }
 
-  let(:host) { load_balancer.host_list.hosts.first }
+  let(:host) do
+    Gitlab::Database::LoadBalancing::Host.new('localhost', load_balancer)
+  end
 
   before do
-    allow(Gitlab::Database.main).to receive(:create_connection_pool)
-      .and_return(ActiveRecord::Base.connection_pool)
+    allow(load_balancer).to receive(:create_replica_connection_pool) do
+      ActiveRecord::Base.connection_pool
+    end
   end
 
   def raise_and_wrap(wrapper, original)
@@ -63,7 +64,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Host do
       expect(host.pool)
         .to receive(:disconnect!)
 
-      host.disconnect!(1)
+      host.disconnect!(timeout: 1)
     end
   end
 

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

@@ -3,21 +3,22 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
-  let(:pool) { Gitlab::Database.main.create_connection_pool(2) }
   let(:conflict_error) { Class.new(RuntimeError) }
-
-  let(:lb) { described_class.new(%w(localhost localhost)) }
+  let(:db_host) { ActiveRecord::Base.connection_pool.db_config.host }
+  let(:lb) { described_class.new([db_host, db_host]) }
   let(:request_cache) { lb.send(:request_cache) }
 
   before do
-    allow(Gitlab::Database.main).to receive(:create_connection_pool)
-      .and_return(pool)
     stub_const(
       'Gitlab::Database::LoadBalancing::LoadBalancer::PG::TRSerializationFailure',
       conflict_error
     )
   end
 
+  after do |example|
+    lb.disconnect!(timeout: 0) unless example.metadata[:skip_disconnect]
+  end
+
   def raise_and_wrap(wrapper, original)
     raise original
   rescue original.class
@@ -239,7 +240,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
     context 'when the connection comes from the primary pool' do
       it 'returns :primary' do
         connection = double(:connection)
-        allow(connection).to receive(:pool).and_return(ActiveRecord::Base.connection_pool)
+        allow(connection).to receive(:pool).and_return(lb.send(:pool))
 
         expect(lb.db_role_for_connection(connection)).to be(:primary)
       end
@@ -271,8 +272,8 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
     end
 
     it 'does not create conflicts with other load balancers when caching hosts' do
-      lb1 = described_class.new(%w(localhost localhost), ActiveRecord::Base)
-      lb2 = described_class.new(%w(localhost localhost), Ci::CiDatabaseRecord)
+      lb1 = described_class.new([db_host, db_host], ActiveRecord::Base)
+      lb2 = described_class.new([db_host, db_host], Ci::CiDatabaseRecord)
 
       host1 = lb1.host
       host2 = lb2.host
@@ -456,4 +457,45 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
       end
     end
   end
+
+  describe '#create_replica_connection_pool' do
+    it 'creates a new connection pool with specific pool size and name' do
+      with_replica_pool(5, 'other_host') do |replica_pool|
+        expect(replica_pool)
+          .to be_kind_of(ActiveRecord::ConnectionAdapters::ConnectionPool)
+
+        expect(replica_pool.db_config.host).to eq('other_host')
+        expect(replica_pool.db_config.pool).to eq(5)
+        expect(replica_pool.db_config.name).to end_with("_replica")
+      end
+    end
+
+    it 'allows setting of a custom hostname and port' do
+      with_replica_pool(5, 'other_host', 5432) do |replica_pool|
+        expect(replica_pool.db_config.host).to eq('other_host')
+        expect(replica_pool.db_config.configuration_hash[:port]).to eq(5432)
+      end
+    end
+
+    it 'does not modify connection class pool' do
+      expect { with_replica_pool(5) { } }.not_to change { ActiveRecord::Base.connection_pool }
+    end
+
+    def with_replica_pool(*args)
+      pool = lb.create_replica_connection_pool(*args)
+      yield pool
+    ensure
+      pool&.disconnect!
+    end
+  end
+
+  describe '#disconnect!' do
+    it 'calls disconnect on all hosts with a timeout', :skip_disconnect do
+      expect_next_instances_of(Gitlab::Database::LoadBalancing::Host, 2) do |host|
+        expect(host).to receive(:disconnect!).with(timeout: 30)
+      end
+
+      lb.disconnect!(timeout: 30)
+    end
+  end
 end

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

@@ -169,7 +169,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
 
       expect(host)
         .to receive(:disconnect!)
-        .with(2)
+        .with(timeout: 2)
 
       service.replace_hosts([address_bar])
     end

spec/lib/gitlab/database/load_balancing_spec.rb

@@ -3,10 +3,6 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::LoadBalancing do
-  before do
-    stub_env('ENABLE_LOAD_BALANCING_FOR_FOSS', 'true')
-  end
-
   describe '.proxy' do
     before do
       @previous_proxy = ActiveRecord::Base.load_balancing_proxy