diff options
Diffstat (limited to 'spec/lib/gitlab/database/load_balancing')
9 files changed, 410 insertions, 74 deletions
diff --git a/spec/lib/gitlab/database/load_balancing/configuration_spec.rb b/spec/lib/gitlab/database/load_balancing/configuration_spec.rb index 3e5249a3dea..eef248afdf2 100644 --- a/spec/lib/gitlab/database/load_balancing/configuration_spec.rb +++ b/spec/lib/gitlab/database/load_balancing/configuration_spec.rb @@ -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,60 @@ 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::ApplicationRecord') } + 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::ApplicationRecord') } + 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 + stub_env('GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci', nil) + + expect(config.primary_connection_specification_name).to eq('Ci::ApplicationRecord') + 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 diff --git a/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb b/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb index ba2f9485066..ee2718171c0 100644 --- a/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb +++ b/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb @@ -3,12 +3,9 @@ require 'spec_helper' RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do - let(:proxy) do - config = Gitlab::Database::LoadBalancing::Configuration - .new(ActiveRecord::Base) - - described_class.new(Gitlab::Database::LoadBalancing::LoadBalancer.new(config)) - end + let(:config) { Gitlab::Database::LoadBalancing::Configuration.new(ActiveRecord::Base) } + let(:load_balancer) { Gitlab::Database::LoadBalancing::LoadBalancer.new(config) } + let(:proxy) { described_class.new(load_balancer) } describe '#select' do it 'performs a read' do @@ -85,7 +82,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do describe '.insert_all!' do before do ActiveRecord::Schema.define do - create_table :connection_proxy_bulk_insert, force: true do |t| + create_table :_test_connection_proxy_bulk_insert, force: true do |t| t.string :name, null: true end end @@ -93,13 +90,13 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do after do ActiveRecord::Schema.define do - drop_table :connection_proxy_bulk_insert, force: true + drop_table :_test_connection_proxy_bulk_insert, force: true end end let(:model_class) do Class.new(ApplicationRecord) do - self.table_name = "connection_proxy_bulk_insert" + self.table_name = "_test_connection_proxy_bulk_insert" end end @@ -143,9 +140,9 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do context 'with a read query' do it 'runs the transaction and any nested queries on the replica' do - expect(proxy.load_balancer).to receive(:read) + expect(load_balancer).to receive(:read) .twice.and_yield(replica) - expect(proxy.load_balancer).not_to receive(:read_write) + expect(load_balancer).not_to receive(:read_write) expect(session).not_to receive(:write!) proxy.transaction { proxy.select('true') } @@ -154,8 +151,8 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do context 'with a write query' do it 'raises an exception' do - allow(proxy.load_balancer).to receive(:read).and_yield(replica) - allow(proxy.load_balancer).to receive(:read_write).and_yield(replica) + allow(load_balancer).to receive(:read).and_yield(replica) + allow(load_balancer).to receive(:read_write).and_yield(replica) expect do proxy.transaction { proxy.insert('something') } @@ -178,9 +175,9 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do context 'with a read query' do it 'runs the transaction and any nested queries on the primary and stick to it' do - expect(proxy.load_balancer).to receive(:read_write) + expect(load_balancer).to receive(:read_write) .twice.and_yield(primary) - expect(proxy.load_balancer).not_to receive(:read) + expect(load_balancer).not_to receive(:read) expect(session).to receive(:write!) proxy.transaction { proxy.select('true') } @@ -189,9 +186,9 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do context 'with a write query' do it 'runs the transaction and any nested queries on the primary and stick to it' do - expect(proxy.load_balancer).to receive(:read_write) + expect(load_balancer).to receive(:read_write) .twice.and_yield(primary) - expect(proxy.load_balancer).not_to receive(:read) + expect(load_balancer).not_to receive(:read) expect(session).to receive(:write!).twice proxy.transaction { proxy.insert('something') } @@ -209,7 +206,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do end it 'properly forwards keyword arguments' do - allow(proxy.load_balancer).to receive(:read_write) + allow(load_balancer).to receive(:read_write) expect(proxy).to receive(:write_using_load_balancer).and_call_original @@ -234,7 +231,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do end it 'properly forwards keyword arguments' do - allow(proxy.load_balancer).to receive(:read) + allow(load_balancer).to receive(:read) expect(proxy).to receive(:read_using_load_balancer).and_call_original @@ -259,7 +256,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do allow(session).to receive(:use_replicas_for_read_queries?).and_return(false) expect(connection).to receive(:foo).with('foo') - expect(proxy.load_balancer).to receive(:read).and_yield(connection) + expect(load_balancer).to receive(:read).and_yield(connection) proxy.read_using_load_balancer(:foo, 'foo') end @@ -271,7 +268,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do allow(session).to receive(:use_replicas_for_read_queries?).and_return(true) expect(connection).to receive(:foo).with('foo') - expect(proxy.load_balancer).to receive(:read).and_yield(connection) + expect(load_balancer).to receive(:read).and_yield(connection) proxy.read_using_load_balancer(:foo, 'foo') end @@ -283,7 +280,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do allow(session).to receive(:use_replicas_for_read_queries?).and_return(true) expect(connection).to receive(:foo).with('foo') - expect(proxy.load_balancer).to receive(:read).and_yield(connection) + expect(load_balancer).to receive(:read).and_yield(connection) proxy.read_using_load_balancer(:foo, 'foo') end @@ -296,7 +293,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do expect(connection).to receive(:foo).with('foo') - expect(proxy.load_balancer).to receive(:read_write) + expect(load_balancer).to receive(:read_write) .and_yield(connection) proxy.read_using_load_balancer(:foo, 'foo') @@ -314,7 +311,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do end it 'uses but does not stick to the primary' do - expect(proxy.load_balancer).to receive(:read_write).and_yield(connection) + expect(load_balancer).to receive(:read_write).and_yield(connection) expect(connection).to receive(:foo).with('foo') expect(session).not_to receive(:write!) diff --git a/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb b/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb index f824d4cefdf..37b83729125 100644 --- a/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb +++ b/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb @@ -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) } @@ -88,6 +89,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do host = double(:host) allow(lb).to receive(:host).and_return(host) + allow(Rails.application.executor).to receive(:active?).and_return(true) allow(host).to receive(:query_cache_enabled).and_return(false) allow(host).to receive(:connection).and_return(connection) @@ -96,6 +98,20 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do lb.read { 10 } end + it 'does not enable query cache when outside Rails executor context' do + connection = double(:connection) + host = double(:host) + + allow(lb).to receive(:host).and_return(host) + allow(Rails.application.executor).to receive(:active?).and_return(false) + allow(host).to receive(:query_cache_enabled).and_return(false) + allow(host).to receive(:connection).and_return(connection) + + expect(host).not_to receive(:enable_query_cache!) + + lb.read { 10 } + end + it 'marks hosts that are offline' do allow(lb).to receive(:connection_error?).and_return(true) @@ -216,7 +232,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do it 'does not create conflicts with other load balancers when caching hosts' do ci_config = Gitlab::Database::LoadBalancing::Configuration - .new(Ci::CiDatabaseRecord, [db_host, db_host]) + .new(Ci::ApplicationRecord, [db_host, db_host]) lb1 = described_class.new(config) lb2 = described_class.new(ci_config) @@ -459,4 +475,84 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do lb.disconnect!(timeout: 30) end end + + describe '#get_write_location' do + it 'returns a string' do + expect(lb.send(:get_write_location, lb.pool.connection)) + .to be_a(String) + end + + it 'returns nil if there are no results' 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::ApplicationRecord } + + around do |example| + if Gitlab::Database.has_config?(:ci) + example.run + else + # fake additional Database + model.establish_connection( + ActiveRecord::DatabaseConfigurations::HashConfig.new(Rails.env, 'ci', ActiveRecord::Base.connection_db_config.configuration_hash) + ) + + example.run + + # Cleanup connection_specification_name for Ci::ApplicationRecord + model.remove_connection + end + 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::ApplicationRecord 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 + + describe '#wal_diff' do + it 'returns the diff between two write locations' do + loc1 = lb.send(:get_write_location, lb.pool.connection) + + create(:user) # This ensures we get a new WAL location + + loc2 = lb.send(:get_write_location, lb.pool.connection) + diff = lb.wal_diff(loc2, loc1) + + expect(diff).to be_positive + end + end end diff --git a/spec/lib/gitlab/database/load_balancing/primary_host_spec.rb b/spec/lib/gitlab/database/load_balancing/primary_host_spec.rb index 45d81808971..02c9499bedb 100644 --- a/spec/lib/gitlab/database/load_balancing/primary_host_spec.rb +++ b/spec/lib/gitlab/database/load_balancing/primary_host_spec.rb @@ -51,7 +51,11 @@ RSpec.describe Gitlab::Database::LoadBalancing::PrimaryHost do end describe '#offline!' do - it 'does nothing' do + it 'logs the event but does nothing else' do + expect(Gitlab::Database::LoadBalancing::Logger).to receive(:warn) + .with(hash_including(event: :host_offline)) + .and_call_original + expect(host.offline!).to be_nil end end diff --git a/spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb b/spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb index af7e2a4b167..b768d4ecea3 100644 --- a/spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb +++ b/spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb @@ -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 diff --git a/spec/lib/gitlab/database/load_balancing/setup_spec.rb b/spec/lib/gitlab/database/load_balancing/setup_spec.rb index 01646bc76ef..953d83d3b48 100644 --- a/spec/lib/gitlab/database/load_balancing/setup_spec.rb +++ b/spec/lib/gitlab/database/load_balancing/setup_spec.rb @@ -7,19 +7,20 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do it 'sets up the load balancer' do setup = described_class.new(ActiveRecord::Base) - expect(setup).to receive(:disable_prepared_statements) - expect(setup).to receive(:setup_load_balancer) + expect(setup).to receive(:configure_connection) + expect(setup).to receive(:setup_connection_proxy) expect(setup).to receive(:setup_service_discovery) + expect(setup).to receive(:setup_feature_flag_to_model_load_balancing) setup.setup end end - describe '#disable_prepared_statements' do - it 'disables prepared statements and reconnects to the database' do + describe '#configure_connection' do + it 'configures pool, prepared statements and reconnects to the database' do config = double( :config, - configuration_hash: { host: 'localhost' }, + configuration_hash: { host: 'localhost', pool: 2, prepared_statements: true }, env_name: 'test', name: 'main' ) @@ -27,7 +28,11 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do expect(ActiveRecord::DatabaseConfigurations::HashConfig) .to receive(:new) - .with('test', 'main', { host: 'localhost', prepared_statements: false }) + .with('test', 'main', { + host: 'localhost', + prepared_statements: false, + pool: Gitlab::Database.default_pool_size + }) .and_call_original # HashConfig doesn't implement its own #==, so we can't directly compare @@ -36,11 +41,11 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do .to receive(:establish_connection) .with(an_instance_of(ActiveRecord::DatabaseConfigurations::HashConfig)) - described_class.new(model).disable_prepared_statements + described_class.new(model).configure_connection end end - describe '#setup_load_balancer' do + describe '#setup_connection_proxy' do it 'sets up the load balancer' do model = Class.new(ActiveRecord::Base) setup = described_class.new(model) @@ -54,9 +59,9 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do .with(setup.configuration) .and_return(lb) - setup.setup_load_balancer + setup.setup_connection_proxy - expect(model.connection.load_balancer).to eq(lb) + expect(model.load_balancer).to eq(lb) expect(model.sticking) .to be_an_instance_of(Gitlab::Database::LoadBalancing::Sticking) end @@ -77,7 +82,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do model = ActiveRecord::Base setup = described_class.new(model) sv = instance_spy(Gitlab::Database::LoadBalancing::ServiceDiscovery) - lb = model.connection.load_balancer allow(setup.configuration) .to receive(:service_discovery_enabled?) @@ -85,7 +89,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do allow(Gitlab::Database::LoadBalancing::ServiceDiscovery) .to receive(:new) - .with(lb, setup.configuration.service_discovery) + .with(setup.load_balancer, setup.configuration.service_discovery) .and_return(sv) expect(sv).to receive(:perform_service_discovery) @@ -98,7 +102,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do model = ActiveRecord::Base setup = described_class.new(model, start_service_discovery: true) sv = instance_spy(Gitlab::Database::LoadBalancing::ServiceDiscovery) - lb = model.connection.load_balancer allow(setup.configuration) .to receive(:service_discovery_enabled?) @@ -106,7 +109,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do allow(Gitlab::Database::LoadBalancing::ServiceDiscovery) .to receive(:new) - .with(lb, setup.configuration.service_discovery) + .with(setup.load_balancer, setup.configuration.service_discovery) .and_return(sv) expect(sv).to receive(:perform_service_discovery) @@ -116,4 +119,181 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do end end end + + describe '#setup_feature_flag_to_model_load_balancing', :reestablished_active_record_base do + using RSpec::Parameterized::TableSyntax + + where do + { + "with model LB enabled it picks a dedicated CI connection" => { + env_GITLAB_USE_MODEL_LOAD_BALANCING: 'true', + env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: nil, + request_store_active: false, + ff_use_model_load_balancing: nil, + expectations: { + main: { read: 'main_replica', write: 'main' }, + ci: { read: 'ci_replica', write: 'ci' } + } + }, + "with model LB enabled and re-use of primary connection it uses CI connection for reads" => { + env_GITLAB_USE_MODEL_LOAD_BALANCING: 'true', + env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: 'main', + request_store_active: false, + ff_use_model_load_balancing: nil, + expectations: { + main: { read: 'main_replica', write: 'main' }, + ci: { read: 'ci_replica', write: 'main' } + } + }, + "with model LB disabled it fallbacks to use main" => { + env_GITLAB_USE_MODEL_LOAD_BALANCING: 'false', + env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: nil, + request_store_active: false, + ff_use_model_load_balancing: nil, + expectations: { + main: { read: 'main_replica', write: 'main' }, + ci: { read: 'main_replica', write: 'main' } + } + }, + "with model LB disabled, but re-use configured it fallbacks to use main" => { + env_GITLAB_USE_MODEL_LOAD_BALANCING: 'false', + env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: 'main', + request_store_active: false, + ff_use_model_load_balancing: nil, + expectations: { + main: { read: 'main_replica', write: 'main' }, + ci: { read: 'main_replica', write: 'main' } + } + }, + "with FF disabled without RequestStore it uses main" => { + env_GITLAB_USE_MODEL_LOAD_BALANCING: nil, + env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: nil, + request_store_active: false, + ff_use_model_load_balancing: false, + expectations: { + main: { read: 'main_replica', write: 'main' }, + ci: { read: 'main_replica', write: 'main' } + } + }, + "with FF enabled without RequestStore sticking of FF does not work, so it fallbacks to use main" => { + env_GITLAB_USE_MODEL_LOAD_BALANCING: nil, + env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: nil, + request_store_active: false, + ff_use_model_load_balancing: true, + expectations: { + main: { read: 'main_replica', write: 'main' }, + ci: { read: 'main_replica', write: 'main' } + } + }, + "with FF disabled with RequestStore it uses main" => { + env_GITLAB_USE_MODEL_LOAD_BALANCING: nil, + env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: nil, + request_store_active: true, + ff_use_model_load_balancing: false, + expectations: { + main: { read: 'main_replica', write: 'main' }, + ci: { read: 'main_replica', write: 'main' } + } + }, + "with FF enabled with RequestStore it sticks FF and uses CI connection" => { + env_GITLAB_USE_MODEL_LOAD_BALANCING: nil, + env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: nil, + request_store_active: true, + ff_use_model_load_balancing: true, + expectations: { + main: { read: 'main_replica', write: 'main' }, + ci: { read: 'ci_replica', write: 'ci' } + } + }, + "with re-use and FF enabled with RequestStore it sticks FF and uses CI connection for reads" => { + env_GITLAB_USE_MODEL_LOAD_BALANCING: nil, + env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: 'main', + request_store_active: true, + ff_use_model_load_balancing: true, + expectations: { + main: { read: 'main_replica', write: 'main' }, + ci: { read: 'ci_replica', write: 'main' } + } + } + } + end + + with_them do + let(:ci_class) do + Class.new(ActiveRecord::Base) do + def self.name + 'Ci::ApplicationRecordTemporary' + end + + establish_connection ActiveRecord::DatabaseConfigurations::HashConfig.new( + Rails.env, + 'ci', + ActiveRecord::Base.connection_db_config.configuration_hash + ) + end + end + + let(:models) do + { + main: ActiveRecord::Base, + ci: ci_class + } + end + + around do |example| + if request_store_active + Gitlab::WithRequestStore.with_request_store do + RequestStore.clear! + + example.run + end + else + example.run + end + end + + before do + # Rewrite `class_attribute` to use rspec mocking and prevent modifying the objects + allow_next_instance_of(described_class) do |setup| + allow(setup).to receive(:configure_connection) + + allow(setup).to receive(:setup_class_attribute) do |attribute, value| + allow(setup.model).to receive(attribute) { value } + end + end + + stub_env('GITLAB_USE_MODEL_LOAD_BALANCING', env_GITLAB_USE_MODEL_LOAD_BALANCING) + stub_env('GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci', env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci) + stub_feature_flags(use_model_load_balancing: ff_use_model_load_balancing) + + # Make load balancer to force init with a dedicated replicas connections + models.each do |_, model| + described_class.new(model).tap do |subject| + subject.configuration.hosts = [subject.configuration.replica_db_config.host] + subject.setup + end + end + end + + it 'results match expectations' do + result = models.transform_values do |model| + load_balancer = model.connection.instance_variable_get(:@load_balancer) + + { + read: load_balancer.read { |connection| connection.pool.db_config.name }, + write: load_balancer.read_write { |connection| connection.pool.db_config.name } + } + end + + expect(result).to eq(expectations) + end + + it 'does return load_balancer assigned to a given connection' do + models.each do |name, model| + expect(model.load_balancer.name).to eq(name) + expect(model.sticking.instance_variable_get(:@load_balancer)).to eq(model.load_balancer) + end + end + end + end end diff --git a/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb b/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb index 08dd6a0a788..9acf80e684f 100644 --- a/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb +++ b/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb @@ -181,11 +181,11 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do end context 'when worker data consistency is :delayed' do - include_examples 'mark data consistency location', :delayed + include_examples 'mark data consistency location', :delayed end context 'when worker data consistency is :sticky' do - include_examples 'mark data consistency location', :sticky + include_examples 'mark data consistency location', :sticky end end end diff --git a/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb b/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb index 06efdcd8f99..de2ad662d16 100644 --- a/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb +++ b/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb @@ -64,7 +64,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware, :clean_ let(:wal_locations) { { Gitlab::Database::MAIN_DATABASE_NAME.to_sym => location } } it 'does not stick to the primary', :aggregate_failures do - expect(ActiveRecord::Base.connection.load_balancer) + expect(ActiveRecord::Base.load_balancer) .to receive(:select_up_to_date_host) .with(location) .and_return(true) @@ -107,7 +107,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware, :clean_ let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e', 'dedup_wal_locations' => wal_locations } } before do - allow(ActiveRecord::Base.connection.load_balancer) + allow(ActiveRecord::Base.load_balancer) .to receive(:select_up_to_date_host) .with(wal_locations[:main]) .and_return(true) @@ -120,7 +120,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware, :clean_ let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e', 'database_write_location' => '0/D525E3A8' } } before do - allow(ActiveRecord::Base.connection.load_balancer) + allow(ActiveRecord::Base.load_balancer) .to receive(:select_up_to_date_host) .with('0/D525E3A8') .and_return(true) diff --git a/spec/lib/gitlab/database/load_balancing/sticking_spec.rb b/spec/lib/gitlab/database/load_balancing/sticking_spec.rb index 8ceda52ee85..d88554614cf 100644 --- a/spec/lib/gitlab/database/load_balancing/sticking_spec.rb +++ b/spec/lib/gitlab/database/load_balancing/sticking_spec.rb @@ -4,7 +4,7 @@ require 'spec_helper' RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do let(:sticking) do - described_class.new(ActiveRecord::Base.connection.load_balancer) + described_class.new(ActiveRecord::Base.load_balancer) end after do @@ -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 @@ -73,7 +73,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do end describe '#all_caught_up?' do - let(:lb) { ActiveRecord::Base.connection.load_balancer } + let(:lb) { ActiveRecord::Base.load_balancer } let(:last_write_location) { 'foo' } before do @@ -137,7 +137,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do end describe '#unstick_or_continue_sticking' do - let(:lb) { ActiveRecord::Base.connection.load_balancer } + let(:lb) { ActiveRecord::Base.load_balancer } it 'simply returns if no write location could be found' do allow(sticking) @@ -182,13 +182,13 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do RSpec.shared_examples 'sticking' do before do - allow(ActiveRecord::Base.connection.load_balancer) + allow(ActiveRecord::Base.load_balancer) .to receive(:primary_write_location) .and_return('foo') end it 'sticks an entity to the primary', :aggregate_failures do - allow(ActiveRecord::Base.connection.load_balancer) + allow(ActiveRecord::Base.load_balancer) .to receive(:primary_only?) .and_return(false) @@ -227,11 +227,11 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do describe '#mark_primary_write_location' do it 'updates the write location with the load balancer' do - allow(ActiveRecord::Base.connection.load_balancer) + allow(ActiveRecord::Base.load_balancer) .to receive(:primary_write_location) .and_return('foo') - allow(ActiveRecord::Base.connection.load_balancer) + allow(ActiveRecord::Base.load_balancer) .to receive(:primary_only?) .and_return(false) @@ -291,7 +291,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do end describe '#select_caught_up_replicas' do - let(:lb) { ActiveRecord::Base.connection.load_balancer } + let(:lb) { ActiveRecord::Base.load_balancer } context 'with no write location' do before do |