Commit 89543370 authored by Adam Hegyi's avatar Adam Hegyi

Merge branch 'lb-reuses-similar-connections' into 'master'

Make `LoadBalancer` to re-use very similar connection

See merge request gitlab-org/gitlab!72698
parents 09583ed3 28ea4ecc
......@@ -7,7 +7,7 @@ module Gitlab
class Configuration
attr_accessor :hosts, :max_replication_difference,
:max_replication_lag_time, :replica_check_interval,
:service_discovery, :model
:service_discovery
# Creates a configuration object for the given ActiveRecord model.
def self.for_model(model)
......@@ -41,6 +41,8 @@ module Gitlab
end
end
config.reuse_primary_connection!
config
end
......@@ -59,6 +61,24 @@ module Gitlab
disconnect_timeout: 120,
use_tcp: false
}
# Temporary model for GITLAB_LOAD_BALANCING_REUSE_PRIMARY_
# To be removed with FF
@primary_model = nil
end
def db_config_name
@model.connection_db_config.name.to_sym
end
# With connection re-use the primary connection can be overwritten
# to be used from different model
def primary_connection_specification_name
(@primary_model || @model).connection_specification_name
end
def replica_db_config
@model.connection_db_config
end
def pool_size
......@@ -86,6 +106,30 @@ module Gitlab
def service_discovery_enabled?
service_discovery[:record].present?
end
# TODO: This is temporary code to allow re-use of primary connection
# if the two connections are pointing to the same host. This is needed
# to properly support transaction visibility.
#
# This behavior is required to support [Phase 3](https://gitlab.com/groups/gitlab-org/-/epics/6160#progress).
# This method is meant to be removed as soon as it is finished.
#
# The remapping is done as-is:
# export GITLAB_LOAD_BALANCING_REUSE_PRIMARY_<name-of-connection>=<new-name-of-connection>
#
# Ex.:
# export GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci=main
#
def reuse_primary_connection!
new_connection = ENV["GITLAB_LOAD_BALANCING_REUSE_PRIMARY_#{db_config_name}"]
return unless new_connection.present?
@primary_model = Gitlab::Database.database_base_models[new_connection.to_sym]
unless @primary_model
raise "Invalid value for 'GITLAB_LOAD_BALANCING_REUSE_PRIMARY_#{db_config_name}=#{new_connection}'"
end
end
end
end
end
......
......@@ -12,7 +12,7 @@ module Gitlab
REPLICA_SUFFIX = '_replica'
attr_reader :name, :host_list, :configuration
attr_reader :host_list, :configuration
# configuration - An instance of `LoadBalancing::Configuration` that
# contains the configuration details (such as the hosts)
......@@ -26,8 +26,10 @@ module Gitlab
else
HostList.new(configuration.hosts.map { |addr| Host.new(addr, self) })
end
end
@name = @configuration.model.connection_db_config.name.to_sym
def name
@configuration.db_config_name
end
def primary_only?
......@@ -227,7 +229,7 @@ module Gitlab
# host - An optional host name to use instead of the default one.
# port - An optional port to connect to.
def create_replica_connection_pool(pool_size, host = nil, port = nil)
db_config = pool.db_config
db_config = @configuration.replica_db_config
env_config = db_config.configuration_hash.dup
env_config[:pool] = pool_size
......@@ -252,7 +254,7 @@ module Gitlab
# leverage that.
def pool
ActiveRecord::Base.connection_handler.retrieve_connection_pool(
@configuration.model.connection_specification_name,
@configuration.primary_connection_specification_name,
role: ActiveRecord::Base.writing_role,
shard: ActiveRecord::Base.default_shard
) || raise(::ActiveRecord::ConnectionNotEstablished)
......
......@@ -38,8 +38,8 @@ module Gitlab
def unstick_or_continue_sticking(env)
namespaces_and_ids = sticking_namespaces(env)
namespaces_and_ids.each do |(model, namespace, id)|
model.sticking.unstick_or_continue_sticking(namespace, id)
namespaces_and_ids.each do |(sticking, namespace, id)|
sticking.unstick_or_continue_sticking(namespace, id)
end
end
......@@ -47,8 +47,8 @@ module Gitlab
def stick_if_necessary(env)
namespaces_and_ids = sticking_namespaces(env)
namespaces_and_ids.each do |model, namespace, id|
model.sticking.stick_if_necessary(namespace, id)
namespaces_and_ids.each do |sticking, namespace, id|
sticking.stick_if_necessary(namespace, id)
end
end
......@@ -74,7 +74,7 @@ module Gitlab
# models that support load balancing. In the future (if we
# determined this to be OK) we may be able to relax this.
::Gitlab::Database::LoadBalancing.base_models.map do |model|
[model, :user, warden.user.id]
[model.sticking, :user, warden.user.id]
end
elsif env[STICK_OBJECT].present?
env[STICK_OBJECT].to_a
......
......@@ -12,7 +12,6 @@ module Gitlab
def initialize(load_balancer)
@load_balancer = load_balancer
@model = load_balancer.configuration.model
end
# Unsticks or continues sticking the current request.
......@@ -28,7 +27,7 @@ module Gitlab
unstick_or_continue_sticking(namespace, id)
env[RackMiddleware::STICK_OBJECT] ||= Set.new
env[RackMiddleware::STICK_OBJECT] << [@model, namespace, id]
env[RackMiddleware::STICK_OBJECT] << [self, namespace, id]
end
# Sticks to the primary if a write was performed.
......
......@@ -3,17 +3,12 @@
require 'spec_helper'
RSpec.describe Gitlab::Database::LoadBalancing::Configuration do
let(:model) do
config = ActiveRecord::DatabaseConfigurations::HashConfig
.new('main', 'test', configuration_hash)
double(:model, connection_db_config: config)
end
let(:configuration_hash) { {} }
let(:db_config) { ActiveRecord::DatabaseConfigurations::HashConfig.new('test', 'ci', configuration_hash) }
let(:model) { double(:model, connection_db_config: db_config) }
describe '.for_model' do
context 'when load balancing is not configured' do
let(:configuration_hash) { {} }
it 'uses the default settings' do
config = described_class.for_model(model)
......@@ -105,6 +100,14 @@ RSpec.describe Gitlab::Database::LoadBalancing::Configuration do
expect(config.pool_size).to eq(4)
end
end
it 'calls reuse_primary_connection!' do
expect_next_instance_of(described_class) do |subject|
expect(subject).to receive(:reuse_primary_connection!).and_call_original
end
described_class.for_model(model)
end
end
describe '#load_balancing_enabled?' do
......@@ -180,4 +183,58 @@ RSpec.describe Gitlab::Database::LoadBalancing::Configuration do
end
end
end
describe '#db_config_name' do
let(:config) { described_class.new(model) }
subject { config.db_config_name }
it 'returns connection name as symbol' do
is_expected.to eq(:ci)
end
end
describe '#replica_db_config' do
let(:model) { double(:model, connection_db_config: db_config, connection_specification_name: 'Ci::CiDatabaseRecord') }
let(:config) { described_class.for_model(model) }
it 'returns exactly db_config' do
expect(config.replica_db_config).to eq(db_config)
end
context 'when GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci=main' do
it 'does not change replica_db_config' do
stub_env('GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci', 'main')
expect(config.replica_db_config).to eq(db_config)
end
end
end
describe 'reuse_primary_connection!' do
let(:model) { double(:model, connection_db_config: db_config, connection_specification_name: 'Ci::CiDatabaseRecord') }
let(:config) { described_class.for_model(model) }
context 'when GITLAB_LOAD_BALANCING_REUSE_PRIMARY_* not configured' do
it 'the primary connection uses default specification' do
expect(config.primary_connection_specification_name).to eq('Ci::CiDatabaseRecord')
end
end
context 'when GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci=main' do
it 'the primary connection uses main connection' do
stub_env('GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci', 'main')
expect(config.primary_connection_specification_name).to eq('ActiveRecord::Base')
end
end
context 'when GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci=unknown' do
it 'raises exception' do
stub_env('GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci', 'unknown')
expect { config.reuse_primary_connection! }.to raise_error /Invalid value for/
end
end
end
end
......@@ -4,10 +4,11 @@ require 'spec_helper'
RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
let(:conflict_error) { Class.new(RuntimeError) }
let(:db_host) { ActiveRecord::Base.connection_pool.db_config.host }
let(:model) { ActiveRecord::Base }
let(:db_host) { model.connection_pool.db_config.host }
let(:config) do
Gitlab::Database::LoadBalancing::Configuration
.new(ActiveRecord::Base, [db_host, db_host])
.new(model, [db_host, db_host])
end
let(:lb) { described_class.new(config) }
......@@ -452,4 +453,51 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
expect(lb.send(:get_write_location, double(select_all: []))).to be_nil
end
end
describe 'primary connection re-use', :reestablished_active_record_base do
let(:model) { Ci::CiDatabaseRecord }
before do
# fake additional Database
model.establish_connection(
ActiveRecord::DatabaseConfigurations::HashConfig.new(Rails.env, 'ci', ActiveRecord::Base.connection_db_config.configuration_hash)
)
end
describe '#read' do
it 'returns ci replica connection' do
expect { |b| lb.read(&b) }.to yield_with_args do |args|
expect(args.pool.db_config.name).to eq('ci_replica')
end
end
context 'when GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci=main' do
it 'returns ci replica connection' do
stub_env('GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci', 'main')
expect { |b| lb.read(&b) }.to yield_with_args do |args|
expect(args.pool.db_config.name).to eq('ci_replica')
end
end
end
end
describe '#read_write' do
it 'returns Ci::CiDatabaseRecord connection' do
expect { |b| lb.read_write(&b) }.to yield_with_args do |args|
expect(args.pool.db_config.name).to eq('ci')
end
end
context 'when GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci=main' do
it 'returns ActiveRecord::Base connection' do
stub_env('GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci', 'main')
expect { |b| lb.read_write(&b) }.to yield_with_args do |args|
expect(args.pool.db_config.name).to eq('main')
end
end
end
end
end
end
......@@ -6,12 +6,12 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
let(:app) { double(:app) }
let(:middleware) { described_class.new(app) }
let(:warden_user) { double(:warden, user: double(:user, id: 42)) }
let(:single_sticking_object) { Set.new([[ActiveRecord::Base, :user, 42]]) }
let(:single_sticking_object) { Set.new([[ActiveRecord::Base.sticking, :user, 42]]) }
let(:multiple_sticking_objects) do
Set.new([
[ActiveRecord::Base, :user, 42],
[ActiveRecord::Base, :runner, '123456789'],
[ActiveRecord::Base, :runner, '1234']
[ActiveRecord::Base.sticking, :user, 42],
[ActiveRecord::Base.sticking, :runner, '123456789'],
[ActiveRecord::Base.sticking, :runner, '1234']
])
end
......@@ -162,7 +162,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
it 'returns the warden user if present' do
env = { 'warden' => warden_user }
ids = Gitlab::Database::LoadBalancing.base_models.map do |model|
[model, :user, 42]
[model.sticking, :user, 42]
end
expect(middleware.sticking_namespaces(env)).to eq(ids)
......@@ -181,9 +181,9 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
env = { described_class::STICK_OBJECT => multiple_sticking_objects }
expect(middleware.sticking_namespaces(env)).to eq([
[ActiveRecord::Base, :user, 42],
[ActiveRecord::Base, :runner, '123456789'],
[ActiveRecord::Base, :runner, '1234']
[ActiveRecord::Base.sticking, :user, 42],
[ActiveRecord::Base.sticking, :runner, '123456789'],
[ActiveRecord::Base.sticking, :runner, '1234']
])
end
end
......
......@@ -22,7 +22,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
sticking.stick_or_unstick_request(env, :user, 42)
expect(env[Gitlab::Database::LoadBalancing::RackMiddleware::STICK_OBJECT].to_a)
.to eq([[ActiveRecord::Base, :user, 42]])
.to eq([[sticking, :user, 42]])
end
it 'sticks or unsticks multiple objects and updates the Rack environment' do
......@@ -42,8 +42,8 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
sticking.stick_or_unstick_request(env, :runner, '123456789')
expect(env[Gitlab::Database::LoadBalancing::RackMiddleware::STICK_OBJECT].to_a).to eq([
[ActiveRecord::Base, :user, 42],
[ActiveRecord::Base, :runner, '123456789']
[sticking, :user, 42],
[sticking, :runner, '123456789']
])
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