Commit 00a0eb68 authored by Tiger Watson's avatar Tiger Watson

Merge branch 'load-balancer-connection-arguments' into 'master'

Refactor LoadBalancing::LoadBalancer to support scoping it to different models

See merge request gitlab-org/gitlab!67057
parents 184cd5cd 0a4e9841
...@@ -7,19 +7,17 @@ module Gitlab ...@@ -7,19 +7,17 @@ module Gitlab
# #
# Each host in the load balancer uses the same credentials as the primary # Each host in the load balancer uses the same credentials as the primary
# database. # database.
#
# This class *requires* that `ActiveRecord::Base.retrieve_connection`
# always returns a connection to the primary.
class LoadBalancer class LoadBalancer
CACHE_KEY = :gitlab_load_balancer_host CACHE_KEY = :gitlab_load_balancer_host
attr_reader :host_list attr_reader :host_list
# hosts - The hostnames/addresses of the additional databases. # hosts - The hostnames/addresses of the additional databases.
def initialize(hosts = []) def initialize(hosts = [], model = ActiveRecord::Base)
@host_list = HostList.new(hosts.map { |addr| Host.new(addr, self) }) @host_list = HostList.new(hosts.map { |addr| Host.new(addr, self) })
@connection_db_roles = {}.compare_by_identity @connection_db_roles = {}.compare_by_identity
@connection_db_roles_count = {}.compare_by_identity @connection_db_roles_count = {}.compare_by_identity
@model = model
end end
# Yields a connection that can be used for reads. # Yields a connection that can be used for reads.
...@@ -94,7 +92,7 @@ module Gitlab ...@@ -94,7 +92,7 @@ module Gitlab
# Instead of immediately grinding to a halt we'll retry the operation # Instead of immediately grinding to a halt we'll retry the operation
# a few times. # a few times.
retry_with_backoff do retry_with_backoff do
connection = ActiveRecord::Base.retrieve_connection connection = pool.connection
track_connection_role(connection, ROLE_PRIMARY) track_connection_role(connection, ROLE_PRIMARY)
yield connection yield connection
...@@ -109,7 +107,7 @@ module Gitlab ...@@ -109,7 +107,7 @@ module Gitlab
def db_role_for_connection(connection) def db_role_for_connection(connection)
return @connection_db_roles[connection] if @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_REPLICA if @host_list.manage_pool?(connection.pool)
return ROLE_PRIMARY if connection.pool == ActiveRecord::Base.connection_pool return ROLE_PRIMARY if connection.pool == pool
end end
# Returns a host to use for queries. # Returns a host to use for queries.
...@@ -117,21 +115,21 @@ module Gitlab ...@@ -117,21 +115,21 @@ module Gitlab
# Hosts are scoped per thread so that multiple threads don't # Hosts are scoped per thread so that multiple threads don't
# accidentally re-use the same host + connection. # accidentally re-use the same host + connection.
def host def host
RequestStore[CACHE_KEY] ||= @host_list.next request_cache[CACHE_KEY] ||= @host_list.next
end end
# Releases the host and connection for the current thread. # Releases the host and connection for the current thread.
def release_host def release_host
if host = RequestStore[CACHE_KEY] if host = request_cache[CACHE_KEY]
host.disable_query_cache! host.disable_query_cache!
host.release_connection host.release_connection
end end
RequestStore.delete(CACHE_KEY) request_cache.delete(CACHE_KEY)
end end
def release_primary_connection def release_primary_connection
ActiveRecord::Base.connection_pool.release_connection pool.release_connection
end end
# Returns the transaction write location of the primary. # Returns the transaction write location of the primary.
...@@ -152,7 +150,7 @@ module Gitlab ...@@ -152,7 +150,7 @@ module Gitlab
return false unless host return false unless host
RequestStore[CACHE_KEY] = host request_cache[CACHE_KEY] = host
true true
end end
...@@ -209,6 +207,17 @@ module Gitlab ...@@ -209,6 +207,17 @@ module Gitlab
private private
# ActiveRecord::ConnectionAdapters::ConnectionHandler handles fetching,
# and caching for connections pools for each "connection", so we
# leverage that.
def pool
ActiveRecord::Base.connection_handler.retrieve_connection_pool(
@model.connection_specification_name,
role: ActiveRecord::Base.writing_role,
shard: ActiveRecord::Base.default_shard
)
end
def ensure_caching! def ensure_caching!
host.enable_query_cache! unless host.query_cache_enabled host.enable_query_cache! unless host.query_cache_enabled
end end
...@@ -228,6 +237,11 @@ module Gitlab ...@@ -228,6 +237,11 @@ module Gitlab
@connection_db_roles_count.delete(connection) @connection_db_roles_count.delete(connection)
end end
end end
def request_cache
base = RequestStore[:gitlab_load_balancer] ||= {}
base[pool] ||= {}
end
end end
end end
end end
......
...@@ -7,6 +7,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do ...@@ -7,6 +7,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
let(:conflict_error) { Class.new(RuntimeError) } let(:conflict_error) { Class.new(RuntimeError) }
let(:lb) { described_class.new(%w(localhost localhost)) } let(:lb) { described_class.new(%w(localhost localhost)) }
let(:request_cache) { lb.send(:request_cache) }
before do before do
allow(Gitlab::Database.main).to receive(:create_connection_pool) allow(Gitlab::Database.main).to receive(:create_connection_pool)
...@@ -123,8 +124,9 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do ...@@ -123,8 +124,9 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
describe '#read_write' do describe '#read_write' do
it 'yields a connection for a write' do it 'yields a connection for a write' do
expect { |b| lb.read_write(&b) } connection = ActiveRecord::Base.connection_pool.connection
.to yield_with_args(ActiveRecord::Base.retrieve_connection)
expect { |b| lb.read_write(&b) }.to yield_with_args(connection)
end end
it 'uses a retry with exponential backoffs' do it 'uses a retry with exponential backoffs' do
...@@ -260,13 +262,24 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do ...@@ -260,13 +262,24 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
end end
it 'stores the host in a thread-local variable' do it 'stores the host in a thread-local variable' do
RequestStore.delete(described_class::CACHE_KEY) request_cache.delete(described_class::CACHE_KEY)
expect(lb.host_list).to receive(:next).once.and_call_original expect(lb.host_list).to receive(:next).once.and_call_original
lb.host lb.host
lb.host lb.host
end end
it 'does not create conflicts with other load balancers when caching hosts' do
lb1 = described_class.new(%w(localhost localhost), ActiveRecord::Base)
lb2 = described_class.new(%w(localhost localhost), Ci::CiDatabaseRecord)
host1 = lb1.host
host2 = lb2.host
expect(lb1.send(:request_cache)[described_class::CACHE_KEY]).to eq(host1)
expect(lb2.send(:request_cache)[described_class::CACHE_KEY]).to eq(host2)
end
end end
describe '#release_host' do describe '#release_host' do
...@@ -277,7 +290,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do ...@@ -277,7 +290,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
lb.release_host lb.release_host
expect(RequestStore[described_class::CACHE_KEY]).to be_nil expect(request_cache[described_class::CACHE_KEY]).to be_nil
end end
end end
...@@ -415,7 +428,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do ...@@ -415,7 +428,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
describe '#select_up_to_date_host' do describe '#select_up_to_date_host' do
let(:location) { 'AB/12345'} let(:location) { 'AB/12345'}
let(:hosts) { lb.host_list.hosts } let(:hosts) { lb.host_list.hosts }
let(:set_host) { RequestStore[described_class::CACHE_KEY] } let(:set_host) { request_cache[described_class::CACHE_KEY] }
subject { lb.select_up_to_date_host(location) } subject { lb.select_up_to_date_host(location) }
......
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