Merge branch 'many-dbs-connection-roles' into 'master'

Change `db_role_for_connection` to support many load balancers

See merge request gitlab-org/gitlab!67339

lib/gitlab/database/load_balancing.rb

@@ -127,9 +127,17 @@ module Gitlab
       # recognize the connection, this method returns the primary role
       # directly. In future, we may need to check for other sources.
       def self.db_role_for_connection(connection)
-        return ROLE_PRIMARY if !enable? || proxy.blank?
+        return ROLE_UNKNOWN unless connection
 
-        proxy.load_balancer.db_role_for_connection(connection)
+        # The connection proxy does not have a role assigned
+        # as this is dependent on a execution context
+        return ROLE_UNKNOWN if connection.is_a?(ConnectionProxy)
+
+        if connection.pool.db_config.name.ends_with?(LoadBalancer::REPLICA_SUFFIX)
+          ROLE_REPLICA
+        else
+          ROLE_PRIMARY
+        end
       end
     end
   end

lib/gitlab/database/load_balancing/host_list.rb

@@ -8,13 +8,11 @@ module Gitlab
         # hosts - The list of secondary hosts to add.
         def initialize(hosts = [])
           @hosts = hosts.shuffle
-          @pools = Set.new
           @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
@@ -35,15 +33,10 @@ module Gitlab
           @mutex.synchronize { @hosts.map { |host| [host.host, host.port] } }
         end
 
-        def manage_pool?(pool)
-          @pools.include?(pool)
-        end
-
         def hosts=(hosts)
           @mutex.synchronize do
             @hosts = hosts
             unsafe_shuffle
-            update_pools
           end
 
           set_metrics!
@@ -89,10 +82,6 @@ module Gitlab
         def set_metrics!
           @hosts_gauge.set({}, @hosts.length)
         end
-
-        def update_pools
-          @pools = Set.new(@hosts.map(&:pool))
-        end
       end
     end
   end

lib/gitlab/database/load_balancing/load_balancer.rb

@@ -10,14 +10,14 @@ module Gitlab
       class LoadBalancer
         CACHE_KEY = :gitlab_load_balancer_host
 
+        REPLICA_SUFFIX = '_replica'
+
         attr_reader :host_list
 
         # 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
         end
 
         def disconnect!(timeout: 120)
@@ -29,7 +29,6 @@ module Gitlab
         # If no secondaries were available this method will use the primary
         # instead.
         def read(&block)
-          connection = nil
           conflict_retried = 0
 
           while host
@@ -37,12 +36,8 @@ module Gitlab
 
             begin
               connection = host.connection
-              track_connection_role(connection, ROLE_REPLICA)
-
               return yield connection
             rescue StandardError => error
-              untrack_connection_role(connection)
-
               if serialization_failure?(error)
                 # This error can occur when a query conflicts. See
                 # https://www.postgresql.org/docs/current/static/hot-standby.html#HOT-STANDBY-CONFLICT
@@ -85,8 +80,6 @@ module Gitlab
           )
 
           read_write(&block)
-        ensure
-          untrack_connection_role(connection)
         end
 
         # Yields a connection that can be used for both reads and writes.
@@ -97,21 +90,8 @@ module Gitlab
           # a few times.
           retry_with_backoff do
             connection = pool.connection
-            track_connection_role(connection, ROLE_PRIMARY)
-
             yield connection
           end
-        ensure
-          untrack_connection_role(connection)
-        end
-
-        # Recognize the role (primary/replica) of the database this connection
-        # is connecting to. If the connection is not issued by this load
-        # balancer, return nil
-        def db_role_for_connection(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 == pool
         end
 
         # Returns a host to use for queries.
@@ -222,7 +202,7 @@ module Gitlab
 
           replica_db_config = ActiveRecord::DatabaseConfigurations::HashConfig.new(
             db_config.env_name,
-            db_config.name + "_replica",
+            db_config.name + REPLICA_SUFFIX,
             env_config
           )
 
@@ -250,22 +230,6 @@ module Gitlab
           host.enable_query_cache! unless host.query_cache_enabled
         end
 
-        def track_connection_role(connection, role)
-          @connection_db_roles[connection] = role
-          @connection_db_roles_count[connection] ||= 0
-          @connection_db_roles_count[connection] += 1
-        end
-
-        def untrack_connection_role(connection)
-          return if connection.blank? || @connection_db_roles_count[connection].blank?
-
-          @connection_db_roles_count[connection] -= 1
-          if @connection_db_roles_count[connection] <= 0
-            @connection_db_roles.delete(connection)
-            @connection_db_roles_count.delete(connection)
-          end
-        end
-
         def request_cache
           base = RequestStore[:gitlab_load_balancer] ||= {}
           base[pool] ||= {}

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

@@ -56,44 +56,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::HostList do
     end
   end
 
-  describe '#manage_pool?' do
-    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 pool 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 'returns a copy of the host' do
       first = host_list.hosts

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

@@ -137,126 +137,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
     end
   end
 
-  describe '#db_role_for_connection' do
-    context 'when the load balancer creates the connection with #read' do
-      it 'returns :replica' do
-        role = nil
-        lb.read do |connection|
-          role = lb.db_role_for_connection(connection)
-        end
-
-        expect(role).to be(:replica)
-      end
-    end
-
-    context 'when the load balancer uses nested #read' do
-      it 'returns :replica' do
-        roles = []
-        lb.read do |connection_1|
-          lb.read do |connection_2|
-            roles << lb.db_role_for_connection(connection_2)
-          end
-          roles << lb.db_role_for_connection(connection_1)
-        end
-
-        expect(roles).to eq([:replica, :replica])
-      end
-    end
-
-    context 'when the load balancer creates the connection with #read_write' do
-      it 'returns :primary' do
-        role = nil
-        lb.read_write do |connection|
-          role = lb.db_role_for_connection(connection)
-        end
-
-        expect(role).to be(:primary)
-      end
-    end
-
-    context 'when the load balancer uses nested #read_write' do
-      it 'returns :primary' do
-        roles = []
-        lb.read_write do |connection_1|
-          lb.read_write do |connection_2|
-            roles << lb.db_role_for_connection(connection_2)
-          end
-          roles << lb.db_role_for_connection(connection_1)
-        end
-
-        expect(roles).to eq([:primary, :primary])
-      end
-    end
-
-    context 'when the load balancer falls back the connection creation to primary' do
-      it 'returns :primary' do
-        allow(lb).to receive(:serialization_failure?).and_return(true)
-
-        role = nil
-        raised = 7 # 2 hosts = 6 retries
-
-        lb.read do |connection|
-          if raised > 0
-            raised -= 1
-            raise
-          end
-
-          role = lb.db_role_for_connection(connection)
-        end
-
-        expect(role).to be(:primary)
-      end
-    end
-
-    context 'when the load balancer uses replica after recovery from a failure' do
-      it 'returns :replica' do
-        allow(lb).to receive(:connection_error?).and_return(true)
-
-        role = nil
-        raised = false
-
-        lb.read do |connection|
-          unless raised
-            raised = true
-            raise
-          end
-
-          role = lb.db_role_for_connection(connection)
-        end
-
-        expect(role).to be(:replica)
-      end
-    end
-
-    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(lb.send(: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
-    end
-  end
-
   describe '#host' do
     it 'returns the secondary host to use' do
       expect(lb.host).to be_an_instance_of(Gitlab::Database::LoadBalancing::Host)

spec/lib/gitlab/database/load_balancing_spec.rb

@@ -296,56 +296,38 @@ RSpec.describe Gitlab::Database::LoadBalancing do
   end
 
   describe '.db_role_for_connection' do
-    let(:connection) { double(:conneciton) }
-
     context 'when the load balancing is not configured' do
-      before do
-        allow(described_class).to receive(:enable?).and_return(false)
-      end
+      let(:connection) { ActiveRecord::Base.connection }
 
       it 'returns primary' do
-        expect(described_class.db_role_for_connection(connection)).to be(:primary)
+        expect(described_class.db_role_for_connection(connection)).to eq(:primary)
       end
     end
 
     context 'when the load balancing is configured' do
-      let(:proxy) { described_class::ConnectionProxy.new(%w(foo)) }
-      let(:load_balancer) { described_class::LoadBalancer.new(%w(foo)) }
+      let(:db_host) { ActiveRecord::Base.connection_pool.db_config.host }
+      let(:proxy) { described_class::ConnectionProxy.new([db_host]) }
 
-      before do
-        allow(described_class).to receive(:enable?).and_return(true)
-        allow(described_class).to receive(:proxy).and_return(proxy)
-        allow(proxy).to receive(:load_balancer).and_return(load_balancer)
+      context 'when a proxy connection is used' do
+        it 'returns :unknown' do
+          expect(described_class.db_role_for_connection(proxy)).to eq(:unknown)
+        end
       end
 
-      context 'when the load balancer returns :replica' do
+      context 'when a read connection is used' do
         it 'returns :replica' do
-          allow(load_balancer).to receive(:db_role_for_connection).and_return(:replica)
-
-          expect(described_class.db_role_for_connection(connection)).to be(:replica)
-
-          expect(load_balancer).to have_received(:db_role_for_connection).with(connection)
+          proxy.load_balancer.read do |connection|
+            expect(described_class.db_role_for_connection(connection)).to eq(:replica)
+          end
         end
       end
 
-      context 'when the load balancer returns :primary' do
+      context 'when a read_write connection is used' do
         it 'returns :primary' do
-          allow(load_balancer).to receive(:db_role_for_connection).and_return(:primary)
-
-          expect(described_class.db_role_for_connection(connection)).to be(:primary)
-
-          expect(load_balancer).to have_received(:db_role_for_connection).with(connection)
+          proxy.load_balancer.read_write do |connection|
+            expect(described_class.db_role_for_connection(connection)).to eq(:primary)
           end
         end
-
-      context 'when the load balancer returns nil' do
-        it 'returns nil' do
-          allow(load_balancer).to receive(:db_role_for_connection).and_return(nil)
-
-          expect(described_class.db_role_for_connection(connection)).to be(nil)
-
-          expect(load_balancer).to have_received(:db_role_for_connection).with(connection)
-        end
       end
     end
   end