Commit bf2670a1 authored by Quang-Minh Nguyen's avatar Quang-Minh Nguyen

Fix wrong role classification in nested transaction

Issue https://gitlab.com/gitlab-org/gitlab/-/issues/322133
parent 1802e343
......@@ -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
......
......@@ -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)
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment