Fix wrong role classification in nested transaction

Issue https://gitlab.com/gitlab-org/gitlab/-/issues/322133

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

@@ -19,6 +19,7 @@ module Gitlab
         def initialize(hosts = [])
           @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
 
         # Yields a connection that can be used for reads.
@@ -34,11 +35,11 @@ module Gitlab
 
             begin
               connection = host.connection
-              @connection_db_roles[connection] = ROLE_REPLICA
+              track_connection_role(connection, ROLE_REPLICA)
 
               return yield connection
             rescue => error
-              @connection_db_roles.delete(connection) if connection.present?
+              untrack_connection_role(connection)
 
               if serialization_failure?(error)
                 # This error can occur when a query conflicts. See
@@ -83,7 +84,7 @@ module Gitlab
 
           read_write(&block)
         ensure
-          @connection_db_roles.delete(connection) if connection.present?
+          untrack_connection_role(connection)
         end
 
         # Yields a connection that can be used for both reads and writes.
@@ -94,12 +95,12 @@ module Gitlab
           # a few times.
           retry_with_backoff do
             connection = ActiveRecord::Base.retrieve_connection
-            @connection_db_roles[connection] = ROLE_PRIMARY
+            track_connection_role(connection, ROLE_PRIMARY)
 
             yield connection
           end
         ensure
-          @connection_db_roles.delete(connection) if connection.present?
+          untrack_connection_role(connection)
         end
 
         # Recognize the role (primary/replica) of the database this connection
@@ -203,6 +204,22 @@ module Gitlab
         def ensure_caching!
           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
       end
     end
   end

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

@@ -147,6 +147,20 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
       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
@@ -158,6 +172,20 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
       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)