Commit 15fedb87 authored by Patrick Bair's avatar Patrick Bair

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
parents 0cf05541 e5b07e73
......@@ -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
......
......@@ -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
......
......@@ -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] ||= {}
......
......@@ -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
......
......@@ -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)
......
......@@ -296,55 +296,37 @@ 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)
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)
proxy.load_balancer.read_write do |connection|
expect(described_class.db_role_for_connection(connection)).to eq(:primary)
end
end
end
end
......
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