Add an additional layer to classify a conneciton pool

Issue https://gitlab.com/gitlab-com/gl-infra/scalability/-/issues/965

ee/lib/gitlab/database/load_balancing/host_list.rb

@@ -8,11 +8,13 @@ module Gitlab
         # hosts - The list of secondary hosts to add.
         def initialize(hosts = [])
           @hosts = hosts.shuffle
+          @pools = {}.compare_by_identity
           @index = 0
           @mutex = Mutex.new
           @hosts_gauge = Gitlab::Metrics.gauge(:db_load_balancing_hosts, 'Current number of load balancing hosts')
 
           set_metrics!
+          update_pools
         end
 
         def hosts
@@ -27,9 +29,14 @@ module Gitlab
           @mutex.synchronize { @hosts.map { |host| [host.host, host.port] } }
         end
 
+        def manage_pool?(pool)
+          @pools[pool] == true
+        end
+
         def hosts=(hosts)
           @mutex.synchronize do
             @hosts = hosts.shuffle
+            update_pools
             @index = 0
           end
 
@@ -71,6 +78,12 @@ module Gitlab
         def set_metrics!
           @hosts_gauge.set({}, @hosts.length)
         end
+
+        def update_pools
+          @pools = @hosts.each_with_object({}) do |host, flags|
+            flags[host.pool] = true
+          end
+        end
       end
     end
   end

ee/lib/gitlab/database/load_balancing/load_balancer.rb

@@ -107,7 +107,9 @@ module Gitlab
         # is connecting to. If the connection is not issued by this load
         # balancer, return nil
         def db_role_for_connection(connection)
-          @connection_db_roles[connection]
+          return @connection_db_roles[connection] if @connection_db_roles[connection]
+          return ROLE_REPLICA if @host_list.manage_pool?(connection.pool)
+          return ROLE_PRIMARY if connection.pool == ActiveRecord::Base.connection_pool
         end
 
         # Returns a host to use for queries.

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

@@ -69,6 +69,55 @@ RSpec.describe Gitlab::Database::LoadBalancing::HostList do
     end
   end
 
+  describe '#manage_pool?' do
+    before do
+      allow(Gitlab::Database)
+        .to receive(:create_connection_pool)
+        .and_return(
+          double(:connection),
+          double(:connection),
+          double(:connection),
+          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
+
+        expect(host_list.manage_pool?(pool)).to be(true)
+      end
+    end
+
+    context 'when the testing pool belongs to a former host of the host list' do
+      it 'returns false' do
+        pool = host_list.hosts.first.pool
+        host_list.hosts = [
+          Gitlab::Database::LoadBalancing::Host.new('foo', load_balancer)
+        ]
+
+        expect(host_list.manage_pool?(pool)).to be(false)
+      end
+    end
+
+    context 'when the testing bool belongs to a new host of the host list' do
+      it 'returns true' do
+        host = Gitlab::Database::LoadBalancing::Host.new('foo', load_balancer)
+        host_list.hosts = [host]
+
+        expect(host_list.manage_pool?(host.pool)).to be(true)
+      end
+    end
+
+    context 'when the testing pool does not have any relation with the host list' do
+      it 'returns false' do
+        host = Gitlab::Database::LoadBalancing::Host.new('foo', load_balancer)
+
+        expect(host_list.manage_pool?(host.pool)).to be(false)
+      end
+    end
+  end
+
   describe '#hosts=' do
     it 'updates the list of hosts to use' do
       host_list.hosts = [

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

@@ -226,9 +226,29 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
       end
     end
 
-    context 'when the connection does not come from the load balancer' do
+    context 'when the connection comes from a pool managed by the host list' do
+      it 'returns :replica' do
+        connection = double(:connection)
+        allow(connection).to receive(:pool).and_return(lb.host_list.hosts.first.pool)
+
+        expect(lb.db_role_for_connection(connection)).to be(:replica)
+      end
+    end
+
+    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)
+
+        expect(lb.db_role_for_connection(connection)).to be(:primary)
+      end
+    end
+
+    context 'when the connection does not come from any known pool' do
       it 'returns nil' do
         connection = double(:connection)
+        pool = double(:connection_pool)
+        allow(connection).to receive(:pool).and_return(pool)
 
         expect(lb.db_role_for_connection(connection)).to be(nil)
       end