diff options
author | GitLab Bot <gitlab-bot@gitlab.com> | 2021-06-16 21:25:58 +0300 |
---|---|---|
committer | GitLab Bot <gitlab-bot@gitlab.com> | 2021-06-16 21:25:58 +0300 |
commit | a5f4bba440d7f9ea47046a0a561d49adf0a1e6d4 (patch) | |
tree | fb69158581673816a8cd895f9d352dcb3c678b1e /spec/lib/gitlab/database | |
parent | d16b2e8639e99961de6ddc93909f3bb5c1445ba1 (diff) |
Add latest changes from gitlab-org/gitlab@14-0-stable-eev14.0.0-rc42
Diffstat (limited to 'spec/lib/gitlab/database')
24 files changed, 4388 insertions, 18 deletions
diff --git a/spec/lib/gitlab/database/background_migration/batched_job_spec.rb b/spec/lib/gitlab/database/background_migration/batched_job_spec.rb index 78e0b7627e9..2de784d3e16 100644 --- a/spec/lib/gitlab/database/background_migration/batched_job_spec.rb +++ b/spec/lib/gitlab/database/background_migration/batched_job_spec.rb @@ -49,16 +49,6 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedJob, type: :model d let(:batched_job) { build(:batched_background_migration_job) } let(:batched_migration) { batched_job.batched_migration } - describe '#migration_aborted?' do - before do - batched_migration.status = :aborted - end - - it 'returns the migration aborted?' do - expect(batched_job.migration_aborted?).to eq(batched_migration.aborted?) - end - end - describe '#migration_job_class' do it 'returns the migration job_class' do expect(batched_job.migration_job_class).to eq(batched_migration.job_class) diff --git a/spec/lib/gitlab/database/background_migration/batched_migration_spec.rb b/spec/lib/gitlab/database/background_migration/batched_migration_spec.rb index 43e34325419..d881390cd52 100644 --- a/spec/lib/gitlab/database/background_migration/batched_migration_spec.rb +++ b/spec/lib/gitlab/database/background_migration/batched_migration_spec.rb @@ -19,6 +19,12 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m end end + describe 'validations' do + subject { build(:batched_background_migration) } + + it { is_expected.to validate_uniqueness_of(:job_arguments).scoped_to(:job_class_name, :table_name, :column_name) } + end + describe '.queue_order' do let!(:migration1) { create(:batched_background_migration) } let!(:migration2) { create(:batched_background_migration) } @@ -36,6 +42,38 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m it 'returns the first active migration according to queue order' do expect(described_class.active_migration).to eq(migration2) + create(:batched_background_migration_job, batched_migration: migration1, batch_size: 1000, status: :succeeded) + end + end + + describe '.queued' do + let!(:migration1) { create(:batched_background_migration, :finished) } + let!(:migration2) { create(:batched_background_migration, :paused) } + let!(:migration3) { create(:batched_background_migration, :active) } + + it 'returns active and paused migrations' do + expect(described_class.queued).to contain_exactly(migration2, migration3) + end + end + + describe '.successful_rows_counts' do + let!(:migration1) { create(:batched_background_migration) } + let!(:migration2) { create(:batched_background_migration) } + let!(:migration_without_jobs) { create(:batched_background_migration) } + + before do + create(:batched_background_migration_job, batched_migration: migration1, batch_size: 1000, status: :succeeded) + create(:batched_background_migration_job, batched_migration: migration1, batch_size: 200, status: :failed) + create(:batched_background_migration_job, batched_migration: migration2, batch_size: 500, status: :succeeded) + create(:batched_background_migration_job, batched_migration: migration2, batch_size: 200, status: :running) + end + + it 'returns totals from successful jobs' do + results = described_class.successful_rows_counts([migration1, migration2, migration_without_jobs]) + + expect(results[migration1.id]).to eq(1000) + expect(results[migration2.id]).to eq(500) + expect(results[migration_without_jobs.id]).to eq(nil) end end @@ -324,4 +362,29 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m subject end end + + describe '.for_configuration' do + let!(:migration) do + create( + :batched_background_migration, + job_class_name: 'MyJobClass', + table_name: :projects, + column_name: :id, + job_arguments: [[:id], [:id_convert_to_bigint]] + ) + end + + before do + create(:batched_background_migration, job_class_name: 'OtherClass') + create(:batched_background_migration, table_name: 'other_table') + create(:batched_background_migration, column_name: 'other_column') + create(:batched_background_migration, job_arguments: %w[other arguments]) + end + + it 'finds the migration matching the given configuration parameters' do + actual = described_class.for_configuration('MyJobClass', :projects, :id, [[:id], [:id_convert_to_bigint]]) + + expect(actual).to contain_exactly(migration) + end + end end diff --git a/spec/lib/gitlab/database/consistency_spec.rb b/spec/lib/gitlab/database/consistency_spec.rb new file mode 100644 index 00000000000..35fa65512ae --- /dev/null +++ b/spec/lib/gitlab/database/consistency_spec.rb @@ -0,0 +1,25 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe Gitlab::Database::Consistency do + let(:session) do + Gitlab::Database::LoadBalancing::Session.current + end + + describe '.with_read_consistency' do + it 'sticks to primary database' do + expect(session).not_to be_using_primary + + block = -> (&control) do + described_class.with_read_consistency do + expect(session).to be_using_primary + + control.call + end + end + + expect { |probe| block.call(&probe) }.to yield_control + end + end +end diff --git a/spec/lib/gitlab/database/count/reltuples_count_strategy_spec.rb b/spec/lib/gitlab/database/count/reltuples_count_strategy_spec.rb index 324ed498abc..cdcc862c376 100644 --- a/spec/lib/gitlab/database/count/reltuples_count_strategy_spec.rb +++ b/spec/lib/gitlab/database/count/reltuples_count_strategy_spec.rb @@ -27,7 +27,7 @@ RSpec.describe Gitlab::Database::Count::ReltuplesCountStrategy do end context 'when models using single-type inheritance are used' do - let(:models) { [Group, CiService, Namespace] } + let(:models) { [Group, Integrations::BaseCi, Namespace] } before do models.each do |model| diff --git a/spec/lib/gitlab/database/dynamic_model_helpers_spec.rb b/spec/lib/gitlab/database/dynamic_model_helpers_spec.rb index 23ad621d0ee..0844616ee1c 100644 --- a/spec/lib/gitlab/database/dynamic_model_helpers_spec.rb +++ b/spec/lib/gitlab/database/dynamic_model_helpers_spec.rb @@ -3,12 +3,12 @@ require 'spec_helper' RSpec.describe Gitlab::Database::DynamicModelHelpers do + let(:including_class) { Class.new.include(described_class) } + let(:table_name) { 'projects' } + describe '#define_batchable_model' do subject { including_class.new.define_batchable_model(table_name) } - let(:including_class) { Class.new.include(described_class) } - let(:table_name) { 'projects' } - it 'is an ActiveRecord model' do expect(subject.ancestors).to include(ActiveRecord::Base) end @@ -25,4 +25,86 @@ RSpec.describe Gitlab::Database::DynamicModelHelpers do expect(subject.inheritance_column).to eq('_type_disabled') end end + + describe '#each_batch' do + subject { including_class.new } + + before do + create_list(:project, 2) + end + + context 'when no transaction is open' do + before do + allow(subject).to receive(:transaction_open?).and_return(false) + end + + it 'iterates table in batches' do + each_batch_size = ->(&block) do + subject.each_batch(table_name, of: 1) do |batch| + block.call(batch.size) + end + end + + expect { |b| each_batch_size.call(&b) } + .to yield_successive_args(1, 1) + end + end + + context 'when transaction is open' do + before do + allow(subject).to receive(:transaction_open?).and_return(true) + end + + it 'raises an error' do + expect { subject.each_batch(table_name, of: 1) { |batch| batch.size } } + .to raise_error(RuntimeError, /each_batch should not run inside a transaction/) + end + end + end + + describe '#each_batch_range' do + subject { including_class.new } + + let(:first_project) { create(:project) } + let(:second_project) { create(:project) } + + context 'when no transaction is open' do + before do + allow(subject).to receive(:transaction_open?).and_return(false) + end + + it 'iterates table in batch ranges' do + expect { |b| subject.each_batch_range(table_name, of: 1, &b) } + .to yield_successive_args( + [first_project.id, first_project.id], + [second_project.id, second_project.id] + ) + end + + it 'yields only one batch if bigger than the table size' do + expect { |b| subject.each_batch_range(table_name, of: 2, &b) } + .to yield_successive_args([first_project.id, second_project.id]) + end + + it 'makes it possible to apply a scope' do + each_batch_limited = ->(&b) do + subject.each_batch_range(table_name, scope: ->(table) { table.limit(1) }, of: 1, &b) + end + + expect { |b| each_batch_limited.call(&b) } + .to yield_successive_args([first_project.id, first_project.id]) + end + end + + context 'when transaction is open' do + before do + allow(subject).to receive(:transaction_open?).and_return(true) + end + + it 'raises an error' do + expect { subject.each_batch_range(table_name, of: 1) { 1 } } + .to raise_error(RuntimeError, /each_batch should not run inside a transaction/) + end + end + end end diff --git a/spec/lib/gitlab/database/load_balancing/active_record_proxy_spec.rb b/spec/lib/gitlab/database/load_balancing/active_record_proxy_spec.rb new file mode 100644 index 00000000000..8886ce9756d --- /dev/null +++ b/spec/lib/gitlab/database/load_balancing/active_record_proxy_spec.rb @@ -0,0 +1,20 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe Gitlab::Database::LoadBalancing::ActiveRecordProxy do + describe '#connection' do + it 'returns a connection proxy' do + dummy = Class.new do + include Gitlab::Database::LoadBalancing::ActiveRecordProxy + end + + proxy = double(:proxy) + + expect(Gitlab::Database::LoadBalancing).to receive(:proxy) + .and_return(proxy) + + expect(dummy.new.connection).to eq(proxy) + 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 new file mode 100644 index 00000000000..015dd2ba8d2 --- /dev/null +++ b/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb @@ -0,0 +1,316 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do + let(:proxy) { described_class.new } + + describe '#select' do + it 'performs a read' do + expect(proxy).to receive(:read_using_load_balancer).with(:select, ['foo']) + + proxy.select('foo') + end + end + + describe '#select_all' do + let(:override_proxy) { ActiveRecord::Base.connection.class } + + # We can't use :Gitlab::Utils::Override because this method is dynamically prepended + it 'method signatures match' do + expect(proxy.method(:select_all).parameters).to eq(override_proxy.instance_method(:select_all).parameters) + end + + describe 'using a SELECT query' do + it 'runs the query on a secondary' do + arel = double(:arel) + + expect(proxy).to receive(:read_using_load_balancer) + .with(:select_all, [arel, 'foo', []]) + + proxy.select_all(arel, 'foo') + end + end + + describe 'using a SELECT FOR UPDATE query' do + it 'runs the query on the primary and sticks to it' do + arel = double(:arel, locked: true) + + expect(proxy).to receive(:write_using_load_balancer) + .with(:select_all, [arel, 'foo', []], sticky: true) + + proxy.select_all(arel, 'foo') + end + end + end + + Gitlab::Database::LoadBalancing::ConnectionProxy::NON_STICKY_READS.each do |name| + describe "#{name}" do + it 'runs the query on the replica' do + expect(proxy).to receive(:read_using_load_balancer) + .with(name, ['foo']) + + proxy.send(name, 'foo') + end + end + end + + Gitlab::Database::LoadBalancing::ConnectionProxy::STICKY_WRITES.each do |name| + describe "#{name}" do + it 'runs the query on the primary and sticks to it' do + expect(proxy).to receive(:write_using_load_balancer) + .with(name, ['foo'], sticky: true) + + proxy.send(name, 'foo') + end + end + end + + describe '.insert_all!' do + before do + ActiveRecord::Schema.define do + create_table :connection_proxy_bulk_insert, force: true do |t| + t.string :name, null: true + end + end + end + + after do + ActiveRecord::Schema.define do + drop_table :connection_proxy_bulk_insert, force: true + end + end + + let(:model_class) do + Class.new(ApplicationRecord) do + self.table_name = "connection_proxy_bulk_insert" + end + end + + it 'inserts data in bulk' do + expect(model_class).to receive(:connection) + .at_least(:once) + .and_return(proxy) + + expect(proxy).to receive(:write_using_load_balancer) + .at_least(:once) + .and_call_original + + expect do + model_class.insert_all! [ + { name: "item1" }, + { name: "item2" } + ] + end.to change { model_class.count }.by(2) + end + end + + # We have an extra test for #transaction here to make sure that nested queries + # are also sent to a primary. + describe '#transaction' do + let(:session) { double(:session) } + + before do + allow(Gitlab::Database::LoadBalancing::Session).to receive(:current) + .and_return(session) + end + + context 'session fallbacks ambiguous queries to replicas' do + let(:replica) { double(:connection) } + + before do + allow(session).to receive(:fallback_to_replicas_for_ambiguous_queries?).and_return(true) + allow(session).to receive(:use_primary?).and_return(false) + allow(replica).to receive(:transaction).and_yield + allow(replica).to receive(:select) + end + + 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) + .twice.and_yield(replica) + expect(proxy.load_balancer).not_to receive(:read_write) + expect(session).not_to receive(:write!) + + proxy.transaction { proxy.select('true') } + end + end + + 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) + + expect do + proxy.transaction { proxy.insert('something') } + end.to raise_error(Gitlab::Database::LoadBalancing::ConnectionProxy::WriteInsideReadOnlyTransactionError) + end + end + end + + context 'session does not fallback to replicas for ambiguous queries' do + let(:primary) { double(:connection) } + + before do + allow(session).to receive(:fallback_to_replicas_for_ambiguous_queries?).and_return(false) + allow(session).to receive(:use_replicas_for_read_queries?).and_return(false) + allow(session).to receive(:use_primary?).and_return(true) + allow(primary).to receive(:transaction).and_yield + allow(primary).to receive(:select) + allow(primary).to receive(:insert) + end + + 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) + .twice.and_yield(primary) + expect(proxy.load_balancer).not_to receive(:read) + expect(session).to receive(:write!) + + proxy.transaction { proxy.select('true') } + end + end + + 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) + .twice.and_yield(primary) + expect(proxy.load_balancer).not_to receive(:read) + expect(session).to receive(:write!).twice + + proxy.transaction { proxy.insert('something') } + end + end + end + end + + describe '#method_missing' do + it 'runs the query on the primary without sticking to it' do + expect(proxy).to receive(:write_using_load_balancer) + .with(:foo, ['foo']) + + proxy.foo('foo') + end + + it 'properly forwards trailing hash arguments' do + allow(proxy.load_balancer).to receive(:read_write) + + expect(proxy).to receive(:write_using_load_balancer).and_call_original + + expect { proxy.case_sensitive_comparison(:table, :attribute, :column, { value: :value, format: :format }) } + .not_to raise_error + end + + context 'current session prefers to fallback ambiguous queries to replicas' do + let(:session) { double(:session) } + + before do + allow(Gitlab::Database::LoadBalancing::Session).to receive(:current) + .and_return(session) + allow(session).to receive(:fallback_to_replicas_for_ambiguous_queries?).and_return(true) + allow(session).to receive(:use_primary?).and_return(false) + end + + it 'runs the query on the replica' do + expect(proxy).to receive(:read_using_load_balancer).with(:foo, ['foo']) + + proxy.foo('foo') + end + + it 'properly forwards trailing hash arguments' do + allow(proxy.load_balancer).to receive(:read) + + expect(proxy).to receive(:read_using_load_balancer).and_call_original + + expect { proxy.case_sensitive_comparison(:table, :attribute, :column, { value: :value, format: :format }) } + .not_to raise_error + end + end + end + + describe '#read_using_load_balancer' do + let(:session) { double(:session) } + let(:connection) { double(:connection) } + + before do + allow(Gitlab::Database::LoadBalancing::Session).to receive(:current) + .and_return(session) + end + + context 'with a regular session' do + it 'uses a secondary' do + allow(session).to receive(:use_primary?).and_return(false) + 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) + + proxy.read_using_load_balancer(:foo, ['foo']) + end + end + + context 'with a regular session and forcing all reads to replicas' do + it 'uses a secondary' do + allow(session).to receive(:use_primary?).and_return(false) + 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) + + proxy.read_using_load_balancer(:foo, ['foo']) + end + end + + context 'with a session using the primary but forcing all reads to replicas' do + it 'uses a secondary' do + allow(session).to receive(:use_primary?).and_return(true) + 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) + + proxy.read_using_load_balancer(:foo, ['foo']) + end + end + + describe 'with a session using the primary' do + it 'uses the primary' do + allow(session).to receive(:use_primary?).and_return(true) + 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_write) + .and_yield(connection) + + proxy.read_using_load_balancer(:foo, ['foo']) + end + end + end + + describe '#write_using_load_balancer' do + let(:session) { double(:session) } + let(:connection) { double(:connection) } + + before do + allow(Gitlab::Database::LoadBalancing::Session).to receive(:current) + .and_return(session) + end + + it 'uses but does not stick to the primary when sticking is disabled' do + expect(proxy.load_balancer).to receive(:read_write).and_yield(connection) + expect(connection).to receive(:foo).with('foo') + expect(session).not_to receive(:write!) + + proxy.write_using_load_balancer(:foo, ['foo']) + end + + it 'sticks to the primary when sticking is enabled' do + expect(proxy.load_balancer).to receive(:read_write).and_yield(connection) + expect(connection).to receive(:foo).with('foo') + expect(session).to receive(:write!) + + proxy.write_using_load_balancer(:foo, ['foo'], sticky: true) + end + end +end diff --git a/spec/lib/gitlab/database/load_balancing/host_list_spec.rb b/spec/lib/gitlab/database/load_balancing/host_list_spec.rb new file mode 100644 index 00000000000..873b599f84d --- /dev/null +++ b/spec/lib/gitlab/database/load_balancing/host_list_spec.rb @@ -0,0 +1,188 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe Gitlab::Database::LoadBalancing::HostList do + def expect_metrics(hosts) + expect(Gitlab::Metrics.registry.get(:db_load_balancing_hosts).get({})).to eq(hosts) + end + + before do + allow(Gitlab::Database) + .to receive(:create_connection_pool) + .and_return(ActiveRecord::Base.connection_pool) + end + + let(:load_balancer) { double(:load_balancer) } + let(:host_count) { 2 } + + let(:host_list) do + hosts = Array.new(host_count) do + Gitlab::Database::LoadBalancing::Host.new('localhost', load_balancer, port: 5432) + end + + described_class.new(hosts) + end + + describe '#initialize' do + it 'sets metrics for current number of hosts and current index' do + host_list + + expect_metrics(2) + end + end + + describe '#length' do + it 'returns the number of hosts in the list' do + expect(host_list.length).to eq(2) + end + end + + describe '#host_names_and_ports' do + context 'with ports' do + it 'returns the host names of all hosts' do + hosts = [ + ['localhost', 5432], + ['localhost', 5432] + ] + + expect(host_list.host_names_and_ports).to eq(hosts) + end + end + + context 'without ports' do + let(:host_list) do + hosts = Array.new(2) do + Gitlab::Database::LoadBalancing::Host.new('localhost', load_balancer) + end + + described_class.new(hosts) + end + + it 'returns the host names of all hosts' do + hosts = [ + ['localhost', nil], + ['localhost', nil] + ] + + expect(host_list.host_names_and_ports).to eq(hosts) + end + end + end + + describe '#manage_pool?' do + before do + allow(Gitlab::Database).to receive(:create_connection_pool) { double(:connection) } + end + + 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 + + expect(host_list.hosts).to eq(first) + expect(host_list.hosts.object_id).not_to eq(first.object_id) + end + end + + describe '#hosts=' do + it 'updates the list of hosts to use' do + host_list.hosts = [ + Gitlab::Database::LoadBalancing::Host.new('foo', load_balancer) + ] + + expect(host_list.length).to eq(1) + expect(host_list.hosts[0].host).to eq('foo') + expect_metrics(1) + end + end + + describe '#next' do + it 'returns a host' do + expect(host_list.next) + .to be_an_instance_of(Gitlab::Database::LoadBalancing::Host) + end + + it 'cycles through all available hosts' do + expect(host_list.next).to eq(host_list.hosts[0]) + expect_metrics(2) + + expect(host_list.next).to eq(host_list.hosts[1]) + expect_metrics(2) + + expect(host_list.next).to eq(host_list.hosts[0]) + expect_metrics(2) + end + + it 'skips hosts that are offline' do + allow(host_list.hosts[0]).to receive(:online?).and_return(false) + + expect(host_list.next).to eq(host_list.hosts[1]) + expect_metrics(2) + end + + it 'returns nil if no hosts are online' do + host_list.hosts.each do |host| + allow(host).to receive(:online?).and_return(false) + end + + expect(host_list.next).to be_nil + expect_metrics(2) + end + + it 'returns nil if no hosts are available' do + expect(described_class.new.next).to be_nil + end + end + + describe '#shuffle' do + let(:host_count) { 3 } + + it 'randomizes the list' do + 2.times do + all_hosts = host_list.hosts + + host_list.shuffle + + expect(host_list.length).to eq(host_count) + expect(host_list.hosts).to contain_exactly(*all_hosts) + end + end + end +end diff --git a/spec/lib/gitlab/database/load_balancing/host_spec.rb b/spec/lib/gitlab/database/load_balancing/host_spec.rb new file mode 100644 index 00000000000..4dfddef68c8 --- /dev/null +++ b/spec/lib/gitlab/database/load_balancing/host_spec.rb @@ -0,0 +1,445 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe Gitlab::Database::LoadBalancing::Host do + let(:load_balancer) do + Gitlab::Database::LoadBalancing::LoadBalancer.new(%w[localhost]) + end + + let(:host) { load_balancer.host_list.hosts.first } + + before do + allow(Gitlab::Database).to receive(:create_connection_pool) + .and_return(ActiveRecord::Base.connection_pool) + end + + def raise_and_wrap(wrapper, original) + raise original + rescue original.class + raise wrapper, 'boom' + end + + def wrapped_exception(wrapper, original) + raise_and_wrap(wrapper, original.new) + rescue wrapper => error + error + end + + describe '#connection' do + it 'returns a connection from the pool' do + expect(host.pool).to receive(:connection) + + host.connection + end + end + + describe '#disconnect!' do + it 'disconnects the pool' do + connection = double(:connection, in_use?: false) + pool = double(:pool, connections: [connection]) + + allow(host) + .to receive(:pool) + .and_return(pool) + + expect(host) + .not_to receive(:sleep) + + expect(host.pool) + .to receive(:disconnect!) + + host.disconnect! + end + + it 'disconnects the pool when waiting for connections takes too long' do + connection = double(:connection, in_use?: true) + pool = double(:pool, connections: [connection]) + + allow(host) + .to receive(:pool) + .and_return(pool) + + expect(host.pool) + .to receive(:disconnect!) + + host.disconnect!(1) + end + end + + describe '#release_connection' do + it 'releases the current connection from the pool' do + expect(host.pool).to receive(:release_connection) + + host.release_connection + end + end + + describe '#offline!' do + it 'marks the host as offline' do + expect(host.pool).to receive(:disconnect!) + + expect(Gitlab::Database::LoadBalancing::Logger).to receive(:warn) + .with(hash_including(event: :host_offline)) + .and_call_original + + host.offline! + end + end + + describe '#online?' do + context 'when the replica status is recent enough' do + before do + expect(host).to receive(:check_replica_status?).and_return(false) + end + + it 'returns the latest status' do + expect(host).not_to receive(:refresh_status) + expect(Gitlab::Database::LoadBalancing::Logger).not_to receive(:info) + expect(Gitlab::Database::LoadBalancing::Logger).not_to receive(:warn) + + expect(host).to be_online + end + + it 'returns an offline status' do + host.offline! + + expect(host).not_to receive(:refresh_status) + expect(Gitlab::Database::LoadBalancing::Logger).not_to receive(:info) + expect(Gitlab::Database::LoadBalancing::Logger).not_to receive(:warn) + + expect(host).not_to be_online + end + end + + context 'when the replica status is outdated' do + before do + expect(host) + .to receive(:check_replica_status?) + .and_return(true) + end + + it 'refreshes the status' do + expect(Gitlab::Database::LoadBalancing::Logger).to receive(:info) + .with(hash_including(event: :host_online)) + .and_call_original + + expect(host).to be_online + end + + context 'and replica is not up to date' do + before do + expect(host).to receive(:replica_is_up_to_date?).and_return(false) + end + + it 'marks the host offline' do + expect(Gitlab::Database::LoadBalancing::Logger).to receive(:warn) + .with(hash_including(event: :host_offline)) + .and_call_original + + expect(host).not_to be_online + end + end + end + + context 'when the replica is not online' do + it 'returns false when ActionView::Template::Error is raised' do + wrapped_error = wrapped_exception(ActionView::Template::Error, StandardError) + + allow(host) + .to receive(:check_replica_status?) + .and_raise(wrapped_error) + + expect(host).not_to be_online + end + + it 'returns false when ActiveRecord::StatementInvalid is raised' do + allow(host) + .to receive(:check_replica_status?) + .and_raise(ActiveRecord::StatementInvalid.new('foo')) + + expect(host).not_to be_online + end + + it 'returns false when PG::Error is raised' do + allow(host) + .to receive(:check_replica_status?) + .and_raise(PG::Error) + + expect(host).not_to be_online + end + end + end + + describe '#refresh_status' do + it 'refreshes the status' do + host.offline! + + expect(host) + .to receive(:replica_is_up_to_date?) + .and_call_original + + host.refresh_status + + expect(host).to be_online + end + end + + describe '#check_replica_status?' do + it 'returns true when we need to check the replica status' do + allow(host) + .to receive(:last_checked_at) + .and_return(1.year.ago) + + expect(host.check_replica_status?).to eq(true) + end + + it 'returns false when we do not need to check the replica status' do + freeze_time do + allow(host) + .to receive(:last_checked_at) + .and_return(Time.zone.now) + + expect(host.check_replica_status?).to eq(false) + end + end + end + + describe '#replica_is_up_to_date?' do + context 'when the lag time is below the threshold' do + it 'returns true' do + expect(host) + .to receive(:replication_lag_below_threshold?) + .and_return(true) + + expect(host.replica_is_up_to_date?).to eq(true) + end + end + + context 'when the lag time exceeds the threshold' do + before do + allow(host) + .to receive(:replication_lag_below_threshold?) + .and_return(false) + end + + it 'returns true if the data is recent enough' do + expect(host) + .to receive(:data_is_recent_enough?) + .and_return(true) + + expect(host.replica_is_up_to_date?).to eq(true) + end + + it 'returns false when the data is not recent enough' do + expect(host) + .to receive(:data_is_recent_enough?) + .and_return(false) + + expect(host.replica_is_up_to_date?).to eq(false) + end + end + end + + describe '#replication_lag_below_threshold' do + it 'returns true when the lag time is below the threshold' do + expect(host) + .to receive(:replication_lag_time) + .and_return(1) + + expect(host.replication_lag_below_threshold?).to eq(true) + end + + it 'returns false when the lag time exceeds the threshold' do + expect(host) + .to receive(:replication_lag_time) + .and_return(9000) + + expect(host.replication_lag_below_threshold?).to eq(false) + end + + it 'returns false when no lag time could be calculated' do + expect(host) + .to receive(:replication_lag_time) + .and_return(nil) + + expect(host.replication_lag_below_threshold?).to eq(false) + end + end + + describe '#data_is_recent_enough?' do + it 'returns true when the data is recent enough' do + expect(host.data_is_recent_enough?).to eq(true) + end + + it 'returns false when the data is not recent enough' do + diff = Gitlab::Database::LoadBalancing.max_replication_difference * 2 + + expect(host) + .to receive(:query_and_release) + .and_return({ 'diff' => diff }) + + expect(host.data_is_recent_enough?).to eq(false) + end + + it 'returns false when no lag size could be calculated' do + expect(host) + .to receive(:replication_lag_size) + .and_return(nil) + + expect(host.data_is_recent_enough?).to eq(false) + end + end + + describe '#replication_lag_time' do + it 'returns the lag time as a Float' do + expect(host.replication_lag_time).to be_an_instance_of(Float) + end + + it 'returns nil when the database query returned no rows' do + expect(host) + .to receive(:query_and_release) + .and_return({}) + + expect(host.replication_lag_time).to be_nil + end + end + + describe '#replication_lag_size' do + it 'returns the lag size as an Integer' do + expect(host.replication_lag_size).to be_an_instance_of(Integer) + end + + it 'returns nil when the database query returned no rows' do + expect(host) + .to receive(:query_and_release) + .and_return({}) + + expect(host.replication_lag_size).to be_nil + end + + it 'returns nil when the database connection fails' do + wrapped_error = wrapped_exception(ActionView::Template::Error, StandardError) + + allow(host) + .to receive(:connection) + .and_raise(wrapped_error) + + expect(host.replication_lag_size).to be_nil + end + end + + describe '#primary_write_location' do + it 'returns the write location of the primary' do + expect(host.primary_write_location).to be_an_instance_of(String) + expect(host.primary_write_location).not_to be_empty + end + end + + describe '#caught_up?' do + let(:connection) { double(:connection) } + + before do + allow(connection).to receive(:quote).and_return('foo') + end + + it 'returns true when a host has caught up' do + allow(host).to receive(:connection).and_return(connection) + expect(connection).to receive(:select_all).and_return([{ 'result' => 't' }]) + + expect(host.caught_up?('foo')).to eq(true) + end + + it 'returns true when a host has caught up' do + allow(host).to receive(:connection).and_return(connection) + expect(connection).to receive(:select_all).and_return([{ 'result' => true }]) + + expect(host.caught_up?('foo')).to eq(true) + end + + it 'returns false when a host has not caught up' do + allow(host).to receive(:connection).and_return(connection) + expect(connection).to receive(:select_all).and_return([{ 'result' => 'f' }]) + + expect(host.caught_up?('foo')).to eq(false) + end + + it 'returns false when a host has not caught up' do + allow(host).to receive(:connection).and_return(connection) + expect(connection).to receive(:select_all).and_return([{ 'result' => false }]) + + expect(host.caught_up?('foo')).to eq(false) + end + + it 'returns false when the connection fails' do + wrapped_error = wrapped_exception(ActionView::Template::Error, StandardError) + + allow(host) + .to receive(:connection) + .and_raise(wrapped_error) + + expect(host.caught_up?('foo')).to eq(false) + end + end + + describe '#database_replica_location' do + let(:connection) { double(:connection) } + + it 'returns the write ahead location of the replica', :aggregate_failures do + expect(host) + .to receive(:query_and_release) + .and_return({ 'location' => '0/D525E3A8' }) + + expect(host.database_replica_location).to be_an_instance_of(String) + end + + it 'returns nil when the database query returned no rows' do + expect(host) + .to receive(:query_and_release) + .and_return({}) + + expect(host.database_replica_location).to be_nil + end + + it 'returns nil when the database connection fails' do + wrapped_error = wrapped_exception(ActionView::Template::Error, StandardError) + + allow(host) + .to receive(:connection) + .and_raise(wrapped_error) + + expect(host.database_replica_location).to be_nil + end + end + + describe '#query_and_release' do + it 'executes a SQL query' do + results = host.query_and_release('SELECT 10 AS number') + + expect(results).to be_an_instance_of(Hash) + expect(results['number'].to_i).to eq(10) + end + + it 'releases the connection after running the query' do + expect(host) + .to receive(:release_connection) + .once + + host.query_and_release('SELECT 10 AS number') + end + + it 'returns an empty Hash in the event of an error' do + expect(host.connection) + .to receive(:select_all) + .and_raise(RuntimeError, 'kittens') + + expect(host.query_and_release('SELECT 10 AS number')).to eq({}) + end + end + + describe '#host' do + it 'returns the hostname' do + expect(host.host).to eq('localhost') + end + end +end diff --git a/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb b/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb new file mode 100644 index 00000000000..4705bb23885 --- /dev/null +++ b/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb @@ -0,0 +1,522 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do + let(:pool) { Gitlab::Database.create_connection_pool(2) } + let(:conflict_error) { Class.new(RuntimeError) } + + let(:lb) { described_class.new(%w(localhost localhost)) } + + before do + allow(Gitlab::Database).to receive(:create_connection_pool) + .and_return(pool) + stub_const( + 'Gitlab::Database::LoadBalancing::LoadBalancer::PG::TRSerializationFailure', + conflict_error + ) + end + + def raise_and_wrap(wrapper, original) + raise original + rescue original.class + raise wrapper, 'boop' + end + + def wrapped_exception(wrapper, original) + raise_and_wrap(wrapper, original.new) + rescue wrapper => error + error + end + + def twice_wrapped_exception(top, middle, original) + begin + raise_and_wrap(middle, original.new) + rescue middle => middle_error + raise_and_wrap(top, middle_error) + end + rescue top => top_error + top_error + end + + describe '#read' do + it 'yields a connection for a read' do + connection = double(:connection) + host = double(:host) + + allow(lb).to receive(:host).and_return(host) + allow(host).to receive(:query_cache_enabled).and_return(true) + + expect(host).to receive(:connection).and_return(connection) + + expect { |b| lb.read(&b) }.to yield_with_args(connection) + end + + it 'ensures that query cache is enabled' do + connection = double(:connection) + host = double(:host) + + allow(lb).to receive(:host).and_return(host) + allow(host).to receive(:query_cache_enabled).and_return(false) + allow(host).to receive(:connection).and_return(connection) + + expect(host).to receive(:enable_query_cache!).once + + lb.read { 10 } + end + + it 'marks hosts that are offline' do + allow(lb).to receive(:connection_error?).and_return(true) + + expect(lb.host_list.hosts[0]).to receive(:offline!) + expect(lb).to receive(:release_host) + + raised = false + + returned = lb.read do + unless raised + raised = true + raise + end + + 10 + end + + expect(returned).to eq(10) + end + + it 'retries a query in the event of a serialization failure' do + raised = false + + expect(lb).to receive(:release_host) + + returned = lb.read do + unless raised + raised = true + raise conflict_error + end + + 10 + end + + expect(returned).to eq(10) + end + + it 'retries every host at most 3 times when a query conflict is raised' do + expect(lb).to receive(:release_host).exactly(6).times + expect(lb).to receive(:read_write) + + lb.read { raise conflict_error } + end + + it 'uses the primary if no secondaries are available' do + allow(lb).to receive(:connection_error?).and_return(true) + + expect(lb.host_list.hosts).to all(receive(:online?).and_return(false)) + + expect(lb).to receive(:read_write).and_call_original + + expect { |b| lb.read(&b) } + .to yield_with_args(ActiveRecord::Base.retrieve_connection) + end + end + + describe '#read_write' do + it 'yields a connection for a write' do + expect { |b| lb.read_write(&b) } + .to yield_with_args(ActiveRecord::Base.retrieve_connection) + end + + it 'uses a retry with exponential backoffs' do + expect(lb).to receive(:retry_with_backoff).and_yield + + lb.read_write { 10 } + 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(ActiveRecord::Base.connection_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) + end + + it 'stores the host in a thread-local variable' do + RequestStore.delete(described_class::CACHE_KEY) + RequestStore.delete(described_class::VALID_HOSTS_CACHE_KEY) + + expect(lb.host_list).to receive(:next).once.and_call_original + + lb.host + lb.host + end + end + + describe '#release_host' do + it 'releases the host and its connection' do + host = lb.host + + expect(host).to receive(:disable_query_cache!) + + lb.release_host + + expect(RequestStore[described_class::CACHE_KEY]).to be_nil + expect(RequestStore[described_class::VALID_HOSTS_CACHE_KEY]).to be_nil + end + end + + describe '#release_primary_connection' do + it 'releases the connection to the primary' do + expect(ActiveRecord::Base.connection_pool).to receive(:release_connection) + + lb.release_primary_connection + end + end + + describe '#primary_write_location' do + it 'returns a String in the right format' do + expect(lb.primary_write_location).to match(%r{[A-F0-9]{1,8}/[A-F0-9]{1,8}}) + end + + it 'raises an error if the write location could not be retrieved' do + connection = double(:connection) + + allow(lb).to receive(:read_write).and_yield(connection) + allow(connection).to receive(:select_all).and_return([]) + + expect { lb.primary_write_location }.to raise_error(RuntimeError) + end + end + + describe '#all_caught_up?' do + it 'returns true if all hosts caught up to the write location' do + expect(lb.host_list.hosts).to all(receive(:caught_up?).with('foo').and_return(true)) + + expect(lb.all_caught_up?('foo')).to eq(true) + end + + it 'returns false if a host has not yet caught up' do + expect(lb.host_list.hosts[0]).to receive(:caught_up?) + .with('foo') + .and_return(true) + + expect(lb.host_list.hosts[1]).to receive(:caught_up?) + .with('foo') + .and_return(false) + + expect(lb.all_caught_up?('foo')).to eq(false) + end + end + + describe '#retry_with_backoff' do + it 'returns the value returned by the block' do + value = lb.retry_with_backoff { 10 } + + expect(value).to eq(10) + end + + it 're-raises errors not related to database connections' do + expect(lb).not_to receive(:sleep) # to make sure we're not retrying + + expect { lb.retry_with_backoff { raise 'boop' } } + .to raise_error(RuntimeError) + end + + it 'retries the block when a connection error is raised' do + allow(lb).to receive(:connection_error?).and_return(true) + expect(lb).to receive(:sleep).with(2) + expect(lb).to receive(:release_primary_connection) + + raised = false + returned = lb.retry_with_backoff do + unless raised + raised = true + raise + end + + 10 + end + + expect(returned).to eq(10) + end + + it 're-raises the connection error if the retries did not succeed' do + allow(lb).to receive(:connection_error?).and_return(true) + expect(lb).to receive(:sleep).with(2).ordered + expect(lb).to receive(:sleep).with(4).ordered + expect(lb).to receive(:sleep).with(16).ordered + + expect(lb).to receive(:release_primary_connection).exactly(3).times + + expect { lb.retry_with_backoff { raise } }.to raise_error(RuntimeError) + end + end + + describe '#connection_error?' do + before do + stub_const('Gitlab::Database::LoadBalancing::LoadBalancer::CONNECTION_ERRORS', + [NotImplementedError]) + end + + it 'returns true for a connection error' do + error = NotImplementedError.new + + expect(lb.connection_error?(error)).to eq(true) + end + + it 'returns true for a wrapped connection error' do + wrapped = wrapped_exception(ActiveRecord::StatementInvalid, NotImplementedError) + + expect(lb.connection_error?(wrapped)).to eq(true) + end + + it 'returns true for a wrapped connection error from a view' do + wrapped = wrapped_exception(ActionView::Template::Error, NotImplementedError) + + expect(lb.connection_error?(wrapped)).to eq(true) + end + + it 'returns true for deeply wrapped/nested errors' do + top = twice_wrapped_exception(ActionView::Template::Error, ActiveRecord::StatementInvalid, NotImplementedError) + + expect(lb.connection_error?(top)).to eq(true) + end + + it 'returns true for an invalid encoding error' do + error = RuntimeError.new('invalid encoding name: unicode') + + expect(lb.connection_error?(error)).to eq(true) + end + + it 'returns false for errors not related to database connections' do + error = RuntimeError.new + + expect(lb.connection_error?(error)).to eq(false) + end + end + + describe '#serialization_failure?' do + let(:conflict_error) { Class.new(RuntimeError) } + + before do + stub_const( + 'Gitlab::Database::LoadBalancing::LoadBalancer::PG::TRSerializationFailure', + conflict_error + ) + end + + it 'returns for a serialization error' do + expect(lb.serialization_failure?(conflict_error.new)).to eq(true) + end + + it 'returns true for a wrapped error' do + wrapped = wrapped_exception(ActionView::Template::Error, conflict_error) + + expect(lb.serialization_failure?(wrapped)).to eq(true) + end + end + + describe '#select_caught_up_hosts' do + let(:location) { 'AB/12345'} + let(:hosts) { lb.host_list.hosts } + let(:valid_host_list) { RequestStore[described_class::VALID_HOSTS_CACHE_KEY] } + let(:valid_hosts) { valid_host_list.hosts } + + subject { lb.select_caught_up_hosts(location) } + + context 'when all replicas are caught up' do + before do + expect(hosts).to all(receive(:caught_up?).with(location).and_return(true)) + end + + it 'returns true and sets all hosts to valid' do + expect(subject).to be true + expect(valid_host_list).to be_a(Gitlab::Database::LoadBalancing::HostList) + expect(valid_hosts).to contain_exactly(*hosts) + end + end + + context 'when none of the replicas are caught up' do + before do + expect(hosts).to all(receive(:caught_up?).with(location).and_return(false)) + end + + it 'returns false and does not set the valid hosts' do + expect(subject).to be false + expect(valid_host_list).to be_nil + end + end + + context 'when one of the replicas is caught up' do + before do + expect(hosts[0]).to receive(:caught_up?).with(location).and_return(false) + expect(hosts[1]).to receive(:caught_up?).with(location).and_return(true) + end + + it 'returns true and sets one host to valid' do + expect(subject).to be true + expect(valid_host_list).to be_a(Gitlab::Database::LoadBalancing::HostList) + expect(valid_hosts).to contain_exactly(hosts[1]) + end + + it 'host always returns the caught-up replica' do + subject + + 3.times do + expect(lb.host).to eq(hosts[1]) + RequestStore.delete(described_class::CACHE_KEY) + end + end + end + end + + describe '#select_caught_up_hosts' do + let(:location) { 'AB/12345'} + let(:hosts) { lb.host_list.hosts } + let(:set_host) { RequestStore[described_class::CACHE_KEY] } + + subject { lb.select_up_to_date_host(location) } + + context 'when none of the replicas are caught up' do + before do + expect(hosts).to all(receive(:caught_up?).with(location).and_return(false)) + end + + it 'returns false and does not update the host thread-local variable' do + expect(subject).to be false + expect(set_host).to be_nil + end + end + + context 'when any of the replicas is caught up' do + before do + # `allow` for non-caught up host, because we may not even check it, if will find the caught up one earlier + allow(hosts[0]).to receive(:caught_up?).with(location).and_return(false) + expect(hosts[1]).to receive(:caught_up?).with(location).and_return(true) + end + + it 'returns true and sets host thread-local variable' do + expect(subject).to be true + expect(set_host).to eq(hosts[1]) + end + end + 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 new file mode 100644 index 00000000000..01367716518 --- /dev/null +++ b/spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb @@ -0,0 +1,243 @@ +# frozen_string_literal: true + +require 'spec_helper' + +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([[:user, 42]]) } + let(:multiple_sticking_objects) do + Set.new([ + [:user, 42], + [:runner, '123456789'], + [:runner, '1234'] + ]) + end + + after do + Gitlab::Database::LoadBalancing::Session.clear_session + end + + describe '.stick_or_unstick' do + before do + allow(Gitlab::Database::LoadBalancing).to receive(:enable?) + .and_return(true) + end + + it 'sticks or unsticks a single object and updates the Rack environment' do + expect(Gitlab::Database::LoadBalancing::Sticking) + .to receive(:unstick_or_continue_sticking) + .with(:user, 42) + + env = {} + + described_class.stick_or_unstick(env, :user, 42) + + expect(env[described_class::STICK_OBJECT].to_a).to eq([[:user, 42]]) + end + + it 'sticks or unsticks multiple objects and updates the Rack environment' do + expect(Gitlab::Database::LoadBalancing::Sticking) + .to receive(:unstick_or_continue_sticking) + .with(:user, 42) + .ordered + + expect(Gitlab::Database::LoadBalancing::Sticking) + .to receive(:unstick_or_continue_sticking) + .with(:runner, '123456789') + .ordered + + env = {} + + described_class.stick_or_unstick(env, :user, 42) + described_class.stick_or_unstick(env, :runner, '123456789') + + expect(env[described_class::STICK_OBJECT].to_a).to eq([ + [:user, 42], + [:runner, '123456789'] + ]) + end + end + + describe '#call' do + it 'handles a request' do + env = {} + + expect(middleware).to receive(:clear).twice + + expect(middleware).to receive(:unstick_or_continue_sticking).with(env) + expect(middleware).to receive(:stick_if_necessary).with(env) + + expect(app).to receive(:call).with(env).and_return(10) + + expect(middleware.call(env)).to eq(10) + end + end + + describe '#unstick_or_continue_sticking' do + it 'does not stick if no namespace and identifier could be found' do + expect(Gitlab::Database::LoadBalancing::Sticking) + .not_to receive(:unstick_or_continue_sticking) + + middleware.unstick_or_continue_sticking({}) + end + + it 'sticks to the primary if a warden user is found' do + env = { 'warden' => warden_user } + + expect(Gitlab::Database::LoadBalancing::Sticking) + .to receive(:unstick_or_continue_sticking) + .with(:user, 42) + + middleware.unstick_or_continue_sticking(env) + end + + it 'sticks to the primary if a sticking namespace and identifier is found' do + env = { described_class::STICK_OBJECT => single_sticking_object } + + expect(Gitlab::Database::LoadBalancing::Sticking) + .to receive(:unstick_or_continue_sticking) + .with(:user, 42) + + middleware.unstick_or_continue_sticking(env) + end + + it 'sticks to the primary if multiple sticking namespaces and identifiers were found' do + env = { described_class::STICK_OBJECT => multiple_sticking_objects } + + expect(Gitlab::Database::LoadBalancing::Sticking) + .to receive(:unstick_or_continue_sticking) + .with(:user, 42) + .ordered + + expect(Gitlab::Database::LoadBalancing::Sticking) + .to receive(:unstick_or_continue_sticking) + .with(:runner, '123456789') + .ordered + + expect(Gitlab::Database::LoadBalancing::Sticking) + .to receive(:unstick_or_continue_sticking) + .with(:runner, '1234') + .ordered + + middleware.unstick_or_continue_sticking(env) + end + end + + describe '#stick_if_necessary' do + it 'does not stick to the primary if not necessary' do + expect(Gitlab::Database::LoadBalancing::Sticking) + .not_to receive(:stick_if_necessary) + + middleware.stick_if_necessary({}) + end + + it 'sticks to the primary if a warden user is found' do + env = { 'warden' => warden_user } + + expect(Gitlab::Database::LoadBalancing::Sticking) + .to receive(:stick_if_necessary) + .with(:user, 42) + + middleware.stick_if_necessary(env) + end + + it 'sticks to the primary if a a single sticking object is found' do + env = { described_class::STICK_OBJECT => single_sticking_object } + + expect(Gitlab::Database::LoadBalancing::Sticking) + .to receive(:stick_if_necessary) + .with(:user, 42) + + middleware.stick_if_necessary(env) + end + + it 'sticks to the primary if multiple sticking namespaces and identifiers were found' do + env = { described_class::STICK_OBJECT => multiple_sticking_objects } + + expect(Gitlab::Database::LoadBalancing::Sticking) + .to receive(:stick_if_necessary) + .with(:user, 42) + .ordered + + expect(Gitlab::Database::LoadBalancing::Sticking) + .to receive(:stick_if_necessary) + .with(:runner, '123456789') + .ordered + + expect(Gitlab::Database::LoadBalancing::Sticking) + .to receive(:stick_if_necessary) + .with(:runner, '1234') + .ordered + + middleware.stick_if_necessary(env) + end + end + + describe '#clear' do + it 'clears the currently used host and session' do + lb = double(:lb) + session = double(:session) + + allow(middleware).to receive(:load_balancer).and_return(lb) + + expect(lb).to receive(:release_host) + + stub_const('Gitlab::Database::LoadBalancing::RackMiddleware::Session', + session) + + expect(session).to receive(:clear_session) + + middleware.clear + end + end + + describe '.load_balancer' do + it 'returns a the load balancer' do + proxy = double(:proxy) + + expect(Gitlab::Database::LoadBalancing).to receive(:proxy) + .and_return(proxy) + + expect(proxy).to receive(:load_balancer) + + middleware.load_balancer + end + end + + describe '#sticking_namespaces_and_ids' do + context 'using a Warden request' do + it 'returns the warden user if present' do + env = { 'warden' => warden_user } + + expect(middleware.sticking_namespaces_and_ids(env)).to eq([[:user, 42]]) + end + + it 'returns an empty Array if no user was present' do + warden = double(:warden, user: nil) + env = { 'warden' => warden } + + expect(middleware.sticking_namespaces_and_ids(env)).to eq([]) + end + end + + context 'using a request with a manually set sticking object' do + it 'returns the sticking object' do + env = { described_class::STICK_OBJECT => multiple_sticking_objects } + + expect(middleware.sticking_namespaces_and_ids(env)).to eq([ + [:user, 42], + [:runner, '123456789'], + [:runner, '1234'] + ]) + end + end + + context 'using a regular request' do + it 'returns an empty Array' do + expect(middleware.sticking_namespaces_and_ids({})).to eq([]) + end + end + end +end diff --git a/spec/lib/gitlab/database/load_balancing/resolver_spec.rb b/spec/lib/gitlab/database/load_balancing/resolver_spec.rb new file mode 100644 index 00000000000..0051cf50255 --- /dev/null +++ b/spec/lib/gitlab/database/load_balancing/resolver_spec.rb @@ -0,0 +1,80 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe Gitlab::Database::LoadBalancing::Resolver do + describe '#resolve' do + let(:ip_addr) { IPAddr.new('127.0.0.2') } + + context 'when nameserver is an IP' do + it 'returns an IPAddr object' do + service = described_class.new('127.0.0.2') + + expect(service.resolve).to eq(ip_addr) + end + end + + context 'when nameserver is not an IP' do + subject { described_class.new('localhost').resolve } + + it 'looks the nameserver up in the hosts file' do + allow_next_instance_of(Resolv::Hosts) do |instance| + allow(instance).to receive(:getaddress).with('localhost').and_return('127.0.0.2') + end + + expect(subject).to eq(ip_addr) + end + + context 'when nameserver is not in the hosts file' do + it 'looks the nameserver up in DNS' do + resource = double(:resource, address: ip_addr) + packet = double(:packet, answer: [resource]) + + allow_next_instance_of(Resolv::Hosts) do |instance| + allow(instance).to receive(:getaddress).with('localhost').and_raise(Resolv::ResolvError) + end + + allow(Net::DNS::Resolver).to receive(:start) + .with('localhost', Net::DNS::A) + .and_return(packet) + + expect(subject).to eq(ip_addr) + end + + context 'when nameserver is not in DNS' do + it 'raises an exception' do + allow_next_instance_of(Resolv::Hosts) do |instance| + allow(instance).to receive(:getaddress).with('localhost').and_raise(Resolv::ResolvError) + end + + allow(Net::DNS::Resolver).to receive(:start) + .with('localhost', Net::DNS::A) + .and_return(double(:packet, answer: [])) + + expect { subject }.to raise_exception( + described_class::UnresolvableNameserverError, + 'could not resolve localhost' + ) + end + end + + context 'when DNS does not respond' do + it 'raises an exception' do + allow_next_instance_of(Resolv::Hosts) do |instance| + allow(instance).to receive(:getaddress).with('localhost').and_raise(Resolv::ResolvError) + end + + allow(Net::DNS::Resolver).to receive(:start) + .with('localhost', Net::DNS::A) + .and_raise(Net::DNS::Resolver::NoResponseError) + + expect { subject }.to raise_exception( + described_class::UnresolvableNameserverError, + 'no response from DNS server(s)' + ) + end + end + end + end + end +end diff --git a/spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb b/spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb new file mode 100644 index 00000000000..7fc7b5e8d11 --- /dev/null +++ b/spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb @@ -0,0 +1,252 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do + let(:service) do + described_class.new(nameserver: 'localhost', port: 8600, record: 'foo') + end + + before do + resource = double(:resource, address: IPAddr.new('127.0.0.1')) + packet = double(:packet, answer: [resource]) + + allow(Net::DNS::Resolver).to receive(:start) + .with('localhost', Net::DNS::A) + .and_return(packet) + end + + describe '#initialize' do + describe ':record_type' do + subject { described_class.new(nameserver: 'localhost', port: 8600, record: 'foo', record_type: record_type) } + + context 'with a supported type' do + let(:record_type) { 'SRV' } + + it { expect(subject.record_type).to eq Net::DNS::SRV } + end + + context 'with an unsupported type' do + let(:record_type) { 'AAAA' } + + it 'raises an argument error' do + expect { subject }.to raise_error(ArgumentError, 'Unsupported record type: AAAA') + end + end + end + end + + describe '#start' do + before do + allow(service) + .to receive(:loop) + .and_yield + end + + it 'starts service discovery in a new thread' do + expect(service) + .to receive(:refresh_if_necessary) + .and_return(5) + + expect(service) + .to receive(:rand) + .and_return(2) + + expect(service) + .to receive(:sleep) + .with(7) + + service.start.join + end + + it 'reports exceptions to Sentry' do + error = StandardError.new + + expect(service) + .to receive(:refresh_if_necessary) + .and_raise(error) + + expect(Gitlab::ErrorTracking) + .to receive(:track_exception) + .with(error) + + expect(service) + .to receive(:rand) + .and_return(2) + + expect(service) + .to receive(:sleep) + .with(62) + + service.start.join + end + end + + describe '#refresh_if_necessary' do + let(:address_foo) { described_class::Address.new('foo') } + let(:address_bar) { described_class::Address.new('bar') } + + context 'when a refresh is necessary' do + before do + allow(service) + .to receive(:addresses_from_load_balancer) + .and_return(%w[localhost]) + + allow(service) + .to receive(:addresses_from_dns) + .and_return([10, [address_foo, address_bar]]) + end + + it 'refreshes the load balancer hosts' do + expect(service) + .to receive(:replace_hosts) + .with([address_foo, address_bar]) + + expect(service.refresh_if_necessary).to eq(10) + end + end + + context 'when a refresh is not necessary' do + before do + allow(service) + .to receive(:addresses_from_load_balancer) + .and_return(%w[localhost]) + + allow(service) + .to receive(:addresses_from_dns) + .and_return([10, %w[localhost]]) + end + + it 'does not refresh the load balancer hosts' do + expect(service) + .not_to receive(:replace_hosts) + + expect(service.refresh_if_necessary).to eq(10) + end + end + end + + describe '#replace_hosts' do + let(:address_foo) { described_class::Address.new('foo') } + let(:address_bar) { described_class::Address.new('bar') } + + let(:load_balancer) do + Gitlab::Database::LoadBalancing::LoadBalancer.new([address_foo]) + end + + before do + allow(service) + .to receive(:load_balancer) + .and_return(load_balancer) + end + + it 'replaces the hosts of the load balancer' do + service.replace_hosts([address_bar]) + + expect(load_balancer.host_list.host_names_and_ports).to eq([['bar', nil]]) + end + + it 'disconnects the old connections' do + host = load_balancer.host_list.hosts.first + + allow(service) + .to receive(:disconnect_timeout) + .and_return(2) + + expect(host) + .to receive(:disconnect!) + .with(2) + + service.replace_hosts([address_bar]) + end + end + + describe '#addresses_from_dns' do + let(:service) { described_class.new(nameserver: 'localhost', port: 8600, record: 'foo', record_type: record_type) } + let(:packet) { double(:packet, answer: [res1, res2]) } + + before do + allow(service.resolver) + .to receive(:search) + .with('foo', described_class::RECORD_TYPES[record_type]) + .and_return(packet) + end + + context 'with an A record' do + let(:record_type) { 'A' } + + let(:res1) { double(:resource, address: IPAddr.new('255.255.255.0'), ttl: 90) } + let(:res2) { double(:resource, address: IPAddr.new('127.0.0.1'), ttl: 90) } + + it 'returns a TTL and ordered list of IP addresses' do + addresses = [ + described_class::Address.new('127.0.0.1'), + described_class::Address.new('255.255.255.0') + ] + + expect(service.addresses_from_dns).to eq([90, addresses]) + end + end + + context 'with an SRV record' do + let(:record_type) { 'SRV' } + + let(:res1) { double(:resource, host: 'foo1.service.consul.', port: 5432, weight: 1, priority: 1, ttl: 90) } + let(:res2) { double(:resource, host: 'foo2.service.consul.', port: 5433, weight: 1, priority: 1, ttl: 90) } + let(:res3) { double(:resource, host: 'foo3.service.consul.', port: 5434, weight: 1, priority: 1, ttl: 90) } + let(:packet) { double(:packet, answer: [res1, res2, res3], additional: []) } + + before do + expect_next_instance_of(Gitlab::Database::LoadBalancing::SrvResolver) do |resolver| + allow(resolver).to receive(:address_for).with('foo1.service.consul.').and_return(IPAddr.new('255.255.255.0')) + allow(resolver).to receive(:address_for).with('foo2.service.consul.').and_return(IPAddr.new('127.0.0.1')) + allow(resolver).to receive(:address_for).with('foo3.service.consul.').and_return(nil) + end + end + + it 'returns a TTL and ordered list of hosts' do + addresses = [ + described_class::Address.new('127.0.0.1', 5433), + described_class::Address.new('255.255.255.0', 5432) + ] + + expect(service.addresses_from_dns).to eq([90, addresses]) + end + end + end + + describe '#new_wait_time_for' do + it 'returns the DNS TTL if greater than the default interval' do + res = double(:resource, ttl: 90) + + expect(service.new_wait_time_for([res])).to eq(90) + end + + it 'returns the default interval if greater than the DNS TTL' do + res = double(:resource, ttl: 10) + + expect(service.new_wait_time_for([res])).to eq(60) + end + + it 'returns the default interval if no resources are given' do + expect(service.new_wait_time_for([])).to eq(60) + end + end + + describe '#addresses_from_load_balancer' do + it 'returns the ordered host names of the load balancer' do + load_balancer = Gitlab::Database::LoadBalancing::LoadBalancer.new(%w[b a]) + + allow(service) + .to receive(:load_balancer) + .and_return(load_balancer) + + addresses = [ + described_class::Address.new('a'), + described_class::Address.new('b') + ] + + expect(service.addresses_from_load_balancer).to eq(addresses) + end + end +end diff --git a/spec/lib/gitlab/database/load_balancing/session_spec.rb b/spec/lib/gitlab/database/load_balancing/session_spec.rb new file mode 100644 index 00000000000..74512f76fd4 --- /dev/null +++ b/spec/lib/gitlab/database/load_balancing/session_spec.rb @@ -0,0 +1,353 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe Gitlab::Database::LoadBalancing::Session do + after do + described_class.clear_session + end + + describe '.current' do + it 'returns the current session' do + expect(described_class.current).to be_an_instance_of(described_class) + end + end + + describe '.clear_session' do + it 'clears the current session' do + described_class.current + described_class.clear_session + + expect(RequestStore[described_class::CACHE_KEY]).to be_nil + end + end + + describe '.without_sticky_writes' do + it 'ignores sticky write events sent by a connection proxy' do + described_class.without_sticky_writes do + described_class.current.write! + end + + session = described_class.current + + expect(session).not_to be_using_primary + end + + it 'still is aware of write that happened' do + described_class.without_sticky_writes do + described_class.current.write! + end + + session = described_class.current + + expect(session.performed_write?).to be true + end + end + + describe '#use_primary?' do + it 'returns true when the primary should be used' do + instance = described_class.new + + instance.use_primary! + + expect(instance.use_primary?).to eq(true) + end + + it 'returns false when a secondary should be used' do + expect(described_class.new.use_primary?).to eq(false) + end + + it 'returns true when a write was performed' do + instance = described_class.new + + instance.write! + + expect(instance.use_primary?).to eq(true) + end + end + + describe '#use_primary' do + let(:instance) { described_class.new } + + context 'when primary was used before' do + before do + instance.write! + end + + it 'restores state after use' do + expect { |blk| instance.use_primary(&blk) }.to yield_with_no_args + + expect(instance.use_primary?).to eq(true) + end + end + + context 'when primary was not used' do + it 'restores state after use' do + expect { |blk| instance.use_primary(&blk) }.to yield_with_no_args + + expect(instance.use_primary?).to eq(false) + end + end + + it 'uses primary during block' do + expect do |blk| + instance.use_primary do + expect(instance.use_primary?).to eq(true) + + # call yield probe + blk.to_proc.call + end + end.to yield_control + end + + it 'continues using primary when write was performed' do + instance.use_primary do + instance.write! + end + + expect(instance.use_primary?).to eq(true) + end + end + + describe '#performed_write?' do + it 'returns true if a write was performed' do + instance = described_class.new + + instance.write! + + expect(instance.performed_write?).to eq(true) + end + end + + describe '#ignore_writes' do + it 'ignores write events' do + instance = described_class.new + + instance.ignore_writes { instance.write! } + + expect(instance).not_to be_using_primary + expect(instance.performed_write?).to eq true + end + + it 'does not prevent using primary if an exception is raised' do + instance = described_class.new + + instance.ignore_writes { raise ArgumentError } rescue ArgumentError + instance.write! + + expect(instance).to be_using_primary + end + end + + describe '#use_replicas_for_read_queries' do + let(:instance) { described_class.new } + + it 'sets the flag inside the block' do + expect do |blk| + instance.use_replicas_for_read_queries do + expect(instance.use_replicas_for_read_queries?).to eq(true) + + # call yield probe + blk.to_proc.call + end + end.to yield_control + + expect(instance.use_replicas_for_read_queries?).to eq(false) + end + + it 'restores state after use' do + expect do |blk| + instance.use_replicas_for_read_queries do + instance.use_replicas_for_read_queries do + expect(instance.use_replicas_for_read_queries?).to eq(true) + + # call yield probe + blk.to_proc.call + end + + expect(instance.use_replicas_for_read_queries?).to eq(true) + end + end.to yield_control + + expect(instance.use_replicas_for_read_queries?).to eq(false) + end + + context 'when primary was used before' do + before do + instance.use_primary! + end + + it 'sets the flag inside the block' do + expect do |blk| + instance.use_replicas_for_read_queries do + expect(instance.use_replicas_for_read_queries?).to eq(true) + + # call yield probe + blk.to_proc.call + end + end.to yield_control + + expect(instance.use_replicas_for_read_queries?).to eq(false) + end + end + + context 'when a write query is performed before' do + before do + instance.write! + end + + it 'sets the flag inside the block' do + expect do |blk| + instance.use_replicas_for_read_queries do + expect(instance.use_replicas_for_read_queries?).to eq(true) + + # call yield probe + blk.to_proc.call + end + end.to yield_control + + expect(instance.use_replicas_for_read_queries?).to eq(false) + end + end + end + + describe '#fallback_to_replicas_for_ambiguous_queries' do + let(:instance) { described_class.new } + + it 'sets the flag inside the block' do + expect do |blk| + instance.fallback_to_replicas_for_ambiguous_queries do + expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(true) + + # call yield probe + blk.to_proc.call + end + end.to yield_control + + expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false) + end + + it 'restores state after use' do + expect do |blk| + instance.fallback_to_replicas_for_ambiguous_queries do + instance.fallback_to_replicas_for_ambiguous_queries do + expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(true) + + # call yield probe + blk.to_proc.call + end + + expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(true) + end + end.to yield_control + + expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false) + end + + context 'when primary was used before' do + before do + instance.use_primary! + end + + it 'uses primary during block' do + expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false) + + expect do |blk| + instance.fallback_to_replicas_for_ambiguous_queries do + expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false) + + # call yield probe + blk.to_proc.call + end + end.to yield_control + + expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false) + end + end + + context 'when a write was performed before' do + before do + instance.write! + end + + it 'uses primary during block' do + expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false) + + expect do |blk| + instance.fallback_to_replicas_for_ambiguous_queries do + expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false) + + # call yield probe + blk.to_proc.call + end + end.to yield_control + + expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false) + end + end + + context 'when primary was used inside the block' do + it 'uses primary aterward' do + expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false) + + instance.fallback_to_replicas_for_ambiguous_queries do + expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(true) + + instance.use_primary! + + expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false) + end + + expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false) + end + + it 'restores state after use' do + instance.fallback_to_replicas_for_ambiguous_queries do + instance.fallback_to_replicas_for_ambiguous_queries do + expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(true) + + instance.use_primary! + + expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false) + end + + expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false) + end + + expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false) + end + end + + context 'when a write was performed inside the block' do + it 'uses primary aterward' do + expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false) + + instance.fallback_to_replicas_for_ambiguous_queries do + expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(true) + + instance.write! + + expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false) + end + + expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false) + end + + it 'restores state after use' do + instance.fallback_to_replicas_for_ambiguous_queries do + instance.fallback_to_replicas_for_ambiguous_queries do + expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(true) + + instance.write! + + expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false) + end + + expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false) + end + + expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false) + 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 new file mode 100644 index 00000000000..90051172fca --- /dev/null +++ b/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb @@ -0,0 +1,161 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do + let(:middleware) { described_class.new } + + after do + Gitlab::Database::LoadBalancing::Session.clear_session + end + + describe '#call' do + shared_context 'data consistency worker class' do |data_consistency, feature_flag| + let(:worker_class) do + Class.new do + def self.name + 'TestDataConsistencyWorker' + end + + include ApplicationWorker + + data_consistency data_consistency, feature_flag: feature_flag + + def perform(*args) + end + end + end + + before do + stub_const('TestDataConsistencyWorker', worker_class) + end + end + + shared_examples_for 'does not pass database locations' do + it 'does not pass database locations', :aggregate_failures do + middleware.call(worker_class, job, double(:queue), redis_pool) { 10 } + + expect(job['database_replica_location']).to be_nil + expect(job['database_write_location']).to be_nil + end + end + + shared_examples_for 'mark data consistency location' do |data_consistency| + include_context 'data consistency worker class', data_consistency, :load_balancing_for_test_data_consistency_worker + + let(:location) { '0/D525E3A8' } + + context 'when feature flag load_balancing_for_sidekiq is disabled' do + before do + stub_feature_flags(load_balancing_for_test_data_consistency_worker: false) + end + + include_examples 'does not pass database locations' + end + + context 'when write was not performed' do + before do + allow(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_primary?).and_return(false) + end + + it 'passes database_replica_location' do + expect(middleware).to receive_message_chain(:load_balancer, :host, "database_replica_location").and_return(location) + + middleware.call(worker_class, job, double(:queue), redis_pool) { 10 } + + expect(job['database_replica_location']).to eq(location) + end + end + + context 'when write was performed' do + before do + allow(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_primary?).and_return(true) + end + + it 'passes primary write location', :aggregate_failures do + expect(middleware).to receive_message_chain(:load_balancer, :primary_write_location).and_return(location) + + middleware.call(worker_class, job, double(:queue), redis_pool) { 10 } + + expect(job['database_write_location']).to eq(location) + end + end + end + + shared_examples_for 'database location was already provided' do |provided_database_location, other_location| + shared_examples_for 'does not set database location again' do |use_primary| + before do + allow(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_primary?).and_return(use_primary) + end + + it 'does not set database locations again' do + middleware.call(worker_class, job, double(:queue), redis_pool) { 10 } + + expect(job[provided_database_location]).to eq(old_location) + expect(job[other_location]).to be_nil + end + end + + let(:old_location) { '0/D525E3A8' } + let(:new_location) { 'AB/12345' } + let(:job) { { "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", provided_database_location => old_location } } + + before do + allow(middleware).to receive_message_chain(:load_balancer, :primary_write_location).and_return(new_location) + allow(middleware).to receive_message_chain(:load_balancer, :database_replica_location).and_return(new_location) + end + + context "when write was performed" do + include_examples 'does not set database location again', true + end + + context "when write was not performed" do + include_examples 'does not set database location again', false + end + end + + let(:queue) { 'default' } + let(:redis_pool) { Sidekiq.redis_pool } + let(:worker_class) { 'TestDataConsistencyWorker' } + let(:job) { { "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e" } } + + before do + skip_feature_flags_yaml_validation + skip_default_enabled_yaml_check + end + + context 'when worker cannot be constantized' do + let(:worker_class) { 'ActionMailer::MailDeliveryJob' } + + include_examples 'does not pass database locations' + end + + context 'when worker class does not include ApplicationWorker' do + let(:worker_class) { ActiveJob::QueueAdapters::SidekiqAdapter::JobWrapper } + + include_examples 'does not pass database locations' + end + + context 'database write location was already provided' do + include_examples 'database location was already provided', 'database_write_location', 'database_replica_location' + end + + context 'database replica location was already provided' do + include_examples 'database location was already provided', 'database_replica_location', 'database_write_location' + end + + context 'when worker data consistency is :always' do + include_context 'data consistency worker class', :always, :load_balancing_for_test_data_consistency_worker + + include_examples 'does not pass database locations' + end + + context 'when worker data consistency is :delayed' do + include_examples 'mark data consistency location', :delayed + end + + context 'when worker data consistency is :sticky' do + 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 new file mode 100644 index 00000000000..b7cd0caa922 --- /dev/null +++ b/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb @@ -0,0 +1,201 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware do + let(:middleware) { described_class.new } + + after do + Gitlab::Database::LoadBalancing::Session.clear_session + end + + describe '#call' do + shared_context 'data consistency worker class' do |data_consistency, feature_flag| + let(:worker_class) do + Class.new do + def self.name + 'TestDataConsistencyWorker' + end + + include ApplicationWorker + + data_consistency data_consistency, feature_flag: feature_flag + + def perform(*args) + end + end + end + + before do + stub_const('TestDataConsistencyWorker', worker_class) + end + end + + shared_examples_for 'stick to the primary' do + it 'sticks to the primary' do + middleware.call(worker, job, double(:queue)) do + expect(Gitlab::Database::LoadBalancing::Session.current.use_primary?).to be_truthy + end + end + end + + shared_examples_for 'replica is up to date' do |location, data_consistency| + it 'does not stick to the primary', :aggregate_failures do + expect(middleware).to receive(:replica_caught_up?).with(location).and_return(true) + + middleware.call(worker, job, double(:queue)) do + expect(Gitlab::Database::LoadBalancing::Session.current.use_primary?).not_to be_truthy + end + + expect(job[:database_chosen]).to eq('replica') + end + + it "updates job hash with data_consistency :#{data_consistency}" do + middleware.call(worker, job, double(:queue)) do + expect(job).to include(data_consistency: data_consistency.to_s) + end + end + end + + shared_examples_for 'sticks based on data consistency' do |data_consistency| + include_context 'data consistency worker class', data_consistency, :load_balancing_for_test_data_consistency_worker + + context 'when load_balancing_for_test_data_consistency_worker is disabled' do + before do + stub_feature_flags(load_balancing_for_test_data_consistency_worker: false) + end + + include_examples 'stick to the primary' + end + + context 'when database replica location is set' do + let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e', 'database_replica_location' => '0/D525E3A8' } } + + before do + allow(middleware).to receive(:replica_caught_up?).and_return(true) + end + + it_behaves_like 'replica is up to date', '0/D525E3A8', data_consistency + end + + context 'when database primary location is set' do + let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e', 'database_write_location' => '0/D525E3A8' } } + + before do + allow(middleware).to receive(:replica_caught_up?).and_return(true) + end + + it_behaves_like 'replica is up to date', '0/D525E3A8', data_consistency + end + + context 'when database location is not set' do + let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e' } } + + it_behaves_like 'stick to the primary', nil + end + end + + let(:queue) { 'default' } + let(:redis_pool) { Sidekiq.redis_pool } + let(:worker) { worker_class.new } + let(:job) { { "retry" => 3, "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", 'database_replica_location' => '0/D525E3A8' } } + let(:block) { 10 } + + before do + skip_feature_flags_yaml_validation + skip_default_enabled_yaml_check + allow(middleware).to receive(:clear) + allow(Gitlab::Database::LoadBalancing::Session.current).to receive(:performed_write?).and_return(true) + end + + context 'when worker class does not include ApplicationWorker' do + let(:worker) { ActiveJob::QueueAdapters::SidekiqAdapter::JobWrapper.new } + + include_examples 'stick to the primary' + end + + context 'when worker data consistency is :always' do + include_context 'data consistency worker class', :always, :load_balancing_for_test_data_consistency_worker + + include_examples 'stick to the primary' + end + + context 'when worker data consistency is :delayed' do + include_examples 'sticks based on data consistency', :delayed + + context 'when replica is not up to date' do + before do + allow(::Gitlab::Database::LoadBalancing).to receive_message_chain(:proxy, :load_balancer, :release_host) + allow(::Gitlab::Database::LoadBalancing).to receive_message_chain(:proxy, :load_balancer, :select_up_to_date_host).and_return(false) + end + + around do |example| + with_sidekiq_server_middleware do |chain| + chain.add described_class + Sidekiq::Testing.disable! { example.run } + end + end + + context 'when job is executed first' do + it 'raise an error and retries', :aggregate_failures do + expect do + process_job(job) + end.to raise_error(Sidekiq::JobRetry::Skip) + + expect(job['error_class']).to eq('Gitlab::Database::LoadBalancing::SidekiqServerMiddleware::JobReplicaNotUpToDate') + expect(job[:database_chosen]).to eq('retry') + end + end + + context 'when job is retried' do + it 'stick to the primary', :aggregate_failures do + expect do + process_job(job) + end.to raise_error(Sidekiq::JobRetry::Skip) + + process_job(job) + expect(job[:database_chosen]).to eq('primary') + end + end + + context 'replica selection mechanism feature flag rollout' do + before do + stub_feature_flags(sidekiq_load_balancing_rotate_up_to_date_replica: false) + end + + it 'uses different implmentation' do + expect(::Gitlab::Database::LoadBalancing).to receive_message_chain(:proxy, :load_balancer, :host, :caught_up?).and_return(false) + + expect do + process_job(job) + end.to raise_error(Sidekiq::JobRetry::Skip) + end + end + end + end + + context 'when worker data consistency is :sticky' do + include_examples 'sticks based on data consistency', :sticky + + context 'when replica is not up to date' do + before do + allow(middleware).to receive(:replica_caught_up?).and_return(false) + end + + include_examples 'stick to the primary' + + it 'updates job hash with primary database chosen', :aggregate_failures do + expect { |b| middleware.call(worker, job, double(:queue), &b) }.to yield_control + + expect(job[:database_chosen]).to eq('primary') + end + end + end + end + + def process_job(job) + Sidekiq::JobRetry.new.local(worker_class, job, queue) do + worker_class.process_job(job) + end + end +end diff --git a/spec/lib/gitlab/database/load_balancing/srv_resolver_spec.rb b/spec/lib/gitlab/database/load_balancing/srv_resolver_spec.rb new file mode 100644 index 00000000000..6ac0608d485 --- /dev/null +++ b/spec/lib/gitlab/database/load_balancing/srv_resolver_spec.rb @@ -0,0 +1,61 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe Gitlab::Database::LoadBalancing::SrvResolver do + let(:resolver) { Net::DNS::Resolver.new(nameservers: '127.0.0.1', port: 8600, use_tcp: true) } + let(:additional) { dns_response_packet_from_fixture('srv_with_a_rr_in_additional_section').additional } + + describe '#address_for' do + let(:host) { 'patroni-02-db-gstg.node.east-us-2.consul.' } + + subject { described_class.new(resolver, additional).address_for(host) } + + context 'when additional section contains an A record' do + it 'returns an IP4 address' do + expect(subject).to eq(IPAddr.new('10.224.29.102')) + end + end + + context 'when additional section contains an AAAA record' do + let(:host) { 'a.gtld-servers.net.' } + let(:additional) { dns_response_packet_from_fixture('a_with_aaaa_rr_in_additional_section').additional } + + it 'returns an IP6 address' do + expect(subject).to eq(IPAddr.new('2001:503:a83e::2:30')) + end + end + + context 'when additional section does not contain A nor AAAA records' do + let(:additional) { [] } + + context 'when host resolves to an A record' do + before do + allow(resolver).to receive(:search).with(host, Net::DNS::ANY).and_return(dns_response_packet_from_fixture('a_rr')) + end + + it 'returns an IP4 address' do + expect(subject).to eq(IPAddr.new('10.224.29.102')) + end + end + + context 'when host does resolves to an AAAA record' do + before do + allow(resolver).to receive(:search).with(host, Net::DNS::ANY).and_return(dns_response_packet_from_fixture('aaaa_rr')) + end + + it 'returns an IP6 address' do + expect(subject).to eq(IPAddr.new('2a00:1450:400e:80a::200e')) + end + end + end + end + + def dns_response_packet_from_fixture(fixture_name) + fixture = File.read(Rails.root + "spec/fixtures/dns/#{fixture_name}.json") + encoded_payload = Gitlab::Json.parse(fixture)['payload'] + payload = Base64.decode64(encoded_payload) + + Net::DNS::Packet.parse(payload) + end +end diff --git a/spec/lib/gitlab/database/load_balancing/sticking_spec.rb b/spec/lib/gitlab/database/load_balancing/sticking_spec.rb new file mode 100644 index 00000000000..bf4e3756e0e --- /dev/null +++ b/spec/lib/gitlab/database/load_balancing/sticking_spec.rb @@ -0,0 +1,307 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do + after do + Gitlab::Database::LoadBalancing::Session.clear_session + end + + describe '.stick_if_necessary' do + context 'when sticking is disabled' do + it 'does not perform any sticking' do + expect(described_class).not_to receive(:stick) + + described_class.stick_if_necessary(:user, 42) + end + end + + context 'when sticking is enabled' do + before do + allow(Gitlab::Database::LoadBalancing).to receive(:enable?) + .and_return(true) + end + + it 'does not stick if no write was performed' do + allow(Gitlab::Database::LoadBalancing::Session.current) + .to receive(:performed_write?) + .and_return(false) + + expect(described_class).not_to receive(:stick) + + described_class.stick_if_necessary(:user, 42) + end + + it 'sticks to the primary if a write was performed' do + allow(Gitlab::Database::LoadBalancing::Session.current) + .to receive(:performed_write?) + .and_return(true) + + expect(described_class).to receive(:stick).with(:user, 42) + + described_class.stick_if_necessary(:user, 42) + end + end + end + + describe '.all_caught_up?' do + let(:lb) { double(:lb) } + + before do + allow(described_class).to receive(:load_balancer).and_return(lb) + end + + it 'returns true if no write location could be found' do + allow(described_class).to receive(:last_write_location_for) + .with(:user, 42) + .and_return(nil) + + expect(lb).not_to receive(:all_caught_up?) + + expect(described_class.all_caught_up?(:user, 42)).to eq(true) + end + + it 'returns true, and unsticks if all secondaries have caught up' do + allow(described_class).to receive(:last_write_location_for) + .with(:user, 42) + .and_return('foo') + + allow(lb).to receive(:all_caught_up?).with('foo').and_return(true) + + expect(described_class).to receive(:unstick).with(:user, 42) + + expect(described_class.all_caught_up?(:user, 42)).to eq(true) + end + + it 'return false if the secondaries have not yet caught up' do + allow(described_class).to receive(:last_write_location_for) + .with(:user, 42) + .and_return('foo') + + allow(lb).to receive(:all_caught_up?).with('foo').and_return(false) + + expect(described_class.all_caught_up?(:user, 42)).to eq(false) + end + end + + describe '.unstick_or_continue_sticking' do + let(:lb) { double(:lb) } + + before do + allow(described_class).to receive(:load_balancer).and_return(lb) + end + + it 'simply returns if no write location could be found' do + allow(described_class).to receive(:last_write_location_for) + .with(:user, 42) + .and_return(nil) + + expect(lb).not_to receive(:all_caught_up?) + + described_class.unstick_or_continue_sticking(:user, 42) + end + + it 'unsticks if all secondaries have caught up' do + allow(described_class).to receive(:last_write_location_for) + .with(:user, 42) + .and_return('foo') + + allow(lb).to receive(:all_caught_up?).with('foo').and_return(true) + + expect(described_class).to receive(:unstick).with(:user, 42) + + described_class.unstick_or_continue_sticking(:user, 42) + end + + it 'continues using the primary if the secondaries have not yet caught up' do + allow(described_class).to receive(:last_write_location_for) + .with(:user, 42) + .and_return('foo') + + allow(lb).to receive(:all_caught_up?).with('foo').and_return(false) + + expect(Gitlab::Database::LoadBalancing::Session.current) + .to receive(:use_primary!) + + described_class.unstick_or_continue_sticking(:user, 42) + end + end + + RSpec.shared_examples 'sticking' do + context 'when sticking is disabled' do + it 'does not perform any sticking', :aggregate_failures do + expect(described_class).not_to receive(:set_write_location_for) + expect(Gitlab::Database::LoadBalancing::Session.current).not_to receive(:use_primary!) + + described_class.bulk_stick(:user, ids) + end + end + + context 'when sticking is enabled' do + before do + allow(Gitlab::Database::LoadBalancing).to receive(:configured?).and_return(true) + + lb = double(:lb, primary_write_location: 'foo') + + allow(described_class).to receive(:load_balancer).and_return(lb) + end + + it 'sticks an entity to the primary', :aggregate_failures do + ids.each do |id| + expect(described_class).to receive(:set_write_location_for) + .with(:user, id, 'foo') + end + + expect(Gitlab::Database::LoadBalancing::Session.current) + .to receive(:use_primary!) + + subject + end + end + end + + describe '.stick' do + it_behaves_like 'sticking' do + let(:ids) { [42] } + subject { described_class.stick(:user, ids.first) } + end + end + + describe '.bulk_stick' do + it_behaves_like 'sticking' do + let(:ids) { [42, 43] } + subject { described_class.bulk_stick(:user, ids) } + end + end + + describe '.mark_primary_write_location' do + context 'when enabled' do + before do + allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true) + allow(Gitlab::Database::LoadBalancing).to receive(:configured?).and_return(true) + end + + it 'updates the write location with the load balancer' do + lb = double(:lb, primary_write_location: 'foo') + + allow(described_class).to receive(:load_balancer).and_return(lb) + + expect(described_class).to receive(:set_write_location_for) + .with(:user, 42, 'foo') + + described_class.mark_primary_write_location(:user, 42) + end + end + + context 'when load balancing is configured but not enabled' do + before do + allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false) + allow(Gitlab::Database::LoadBalancing).to receive(:configured?).and_return(true) + end + + it 'updates the write location with the main ActiveRecord connection' do + allow(described_class).to receive(:load_balancer).and_return(nil) + expect(ActiveRecord::Base).to receive(:connection).and_call_original + expect(described_class).to receive(:set_write_location_for) + .with(:user, 42, anything) + + described_class.mark_primary_write_location(:user, 42) + end + + context 'when write location is nil' do + before do + allow(Gitlab::Database).to receive(:get_write_location).and_return(nil) + end + + it 'does not update the write location' do + expect(described_class).not_to receive(:set_write_location_for) + + described_class.mark_primary_write_location(:user, 42) + end + end + end + + context 'when load balancing is disabled' do + before do + allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false) + allow(Gitlab::Database::LoadBalancing).to receive(:configured?).and_return(false) + end + + it 'updates the write location with the main ActiveRecord connection' do + expect(described_class).not_to receive(:set_write_location_for) + + described_class.mark_primary_write_location(:user, 42) + end + end + end + + describe '.unstick' do + it 'removes the sticking data from Redis' do + described_class.set_write_location_for(:user, 4, 'foo') + described_class.unstick(:user, 4) + + expect(described_class.last_write_location_for(:user, 4)).to be_nil + end + end + + describe '.last_write_location_for' do + it 'returns the last WAL write location for a user' do + described_class.set_write_location_for(:user, 4, 'foo') + + expect(described_class.last_write_location_for(:user, 4)).to eq('foo') + end + end + + describe '.redis_key_for' do + it 'returns a String' do + expect(described_class.redis_key_for(:user, 42)) + .to eq('database-load-balancing/write-location/user/42') + end + end + + describe '.load_balancer' do + it 'returns a the load balancer' do + proxy = double(:proxy) + + expect(Gitlab::Database::LoadBalancing).to receive(:proxy) + .and_return(proxy) + + expect(proxy).to receive(:load_balancer) + + described_class.load_balancer + end + end + + describe '.select_caught_up_replicas' do + let(:lb) { double(:lb) } + + before do + allow(described_class).to receive(:load_balancer).and_return(lb) + end + + context 'with no write location' do + before do + allow(described_class).to receive(:last_write_location_for) + .with(:project, 42).and_return(nil) + end + + it 'returns false and does not try to find caught up hosts' do + expect(described_class).not_to receive(:select_caught_up_hosts) + expect(described_class.select_caught_up_replicas(:project, 42)).to be false + end + end + + context 'with write location' do + before do + allow(described_class).to receive(:last_write_location_for) + .with(:project, 42).and_return('foo') + end + + it 'returns true, selects hosts, and unsticks if any secondary has caught up' do + expect(lb).to receive(:select_caught_up_hosts).and_return(true) + expect(described_class).to receive(:unstick).with(:project, 42) + expect(described_class.select_caught_up_replicas(:project, 42)).to be true + end + end + end +end diff --git a/spec/lib/gitlab/database/load_balancing_spec.rb b/spec/lib/gitlab/database/load_balancing_spec.rb new file mode 100644 index 00000000000..e7de7f2b43b --- /dev/null +++ b/spec/lib/gitlab/database/load_balancing_spec.rb @@ -0,0 +1,834 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe Gitlab::Database::LoadBalancing do + include_context 'clear DB Load Balancing configuration' + + before do + stub_env('ENABLE_LOAD_BALANCING_FOR_FOSS', 'true') + end + + describe '.proxy' do + context 'when configured' do + before do + allow(ActiveRecord::Base.singleton_class).to receive(:prepend) + subject.configure_proxy + end + + it 'returns the connection proxy' do + expect(subject.proxy).to be_an_instance_of(subject::ConnectionProxy) + end + end + + context 'when not configured' do + it 'returns nil' do + expect(subject.proxy).to be_nil + end + + it 'tracks an error to sentry' do + expect(Gitlab::ErrorTracking).to receive(:track_exception).with( + an_instance_of(subject::ProxyNotConfiguredError) + ) + + subject.proxy + end + end + end + + describe '.configuration' do + it 'returns a Hash' do + lb_config = { 'hosts' => %w(foo) } + + original_db_config = Gitlab::Database.config + modified_db_config = original_db_config.merge(load_balancing: lb_config) + expect(Gitlab::Database).to receive(:config).and_return(modified_db_config) + + expect(described_class.configuration).to eq(lb_config) + end + end + + describe '.max_replication_difference' do + context 'without an explicitly configured value' do + it 'returns the default value' do + allow(described_class) + .to receive(:configuration) + .and_return({}) + + expect(described_class.max_replication_difference).to eq(8.megabytes) + end + end + + context 'with an explicitly configured value' do + it 'returns the configured value' do + allow(described_class) + .to receive(:configuration) + .and_return({ 'max_replication_difference' => 4 }) + + expect(described_class.max_replication_difference).to eq(4) + end + end + end + + describe '.max_replication_lag_time' do + context 'without an explicitly configured value' do + it 'returns the default value' do + allow(described_class) + .to receive(:configuration) + .and_return({}) + + expect(described_class.max_replication_lag_time).to eq(60) + end + end + + context 'with an explicitly configured value' do + it 'returns the configured value' do + allow(described_class) + .to receive(:configuration) + .and_return({ 'max_replication_lag_time' => 4 }) + + expect(described_class.max_replication_lag_time).to eq(4) + end + end + end + + describe '.replica_check_interval' do + context 'without an explicitly configured value' do + it 'returns the default value' do + allow(described_class) + .to receive(:configuration) + .and_return({}) + + expect(described_class.replica_check_interval).to eq(60) + end + end + + context 'with an explicitly configured value' do + it 'returns the configured value' do + allow(described_class) + .to receive(:configuration) + .and_return({ 'replica_check_interval' => 4 }) + + expect(described_class.replica_check_interval).to eq(4) + end + end + end + + describe '.hosts' do + it 'returns a list of hosts' do + allow(described_class) + .to receive(:configuration) + .and_return({ 'hosts' => %w(foo bar baz) }) + + expect(described_class.hosts).to eq(%w(foo bar baz)) + end + end + + describe '.pool_size' do + it 'returns a Fixnum' do + expect(described_class.pool_size).to be_a_kind_of(Integer) + end + end + + describe '.enable?' do + before do + clear_load_balancing_configuration + allow(described_class).to receive(:hosts).and_return(%w(foo)) + end + + it 'returns false when no hosts are specified' do + allow(described_class).to receive(:hosts).and_return([]) + + expect(described_class.enable?).to eq(false) + end + + it 'returns false when Sidekiq is being used' do + allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(true) + + expect(described_class.enable?).to eq(false) + end + + it 'returns false when running inside a Rake task' do + allow(Gitlab::Runtime).to receive(:rake?).and_return(true) + + expect(described_class.enable?).to eq(false) + end + + it 'returns true when load balancing should be enabled' do + allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(false) + + expect(described_class.enable?).to eq(true) + end + + it 'returns true when service discovery is enabled' do + allow(described_class).to receive(:hosts).and_return([]) + allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(false) + + allow(described_class) + .to receive(:service_discovery_enabled?) + .and_return(true) + + expect(described_class.enable?).to eq(true) + end + + context 'when ENABLE_LOAD_BALANCING_FOR_SIDEKIQ environment variable is set' do + before do + stub_env('ENABLE_LOAD_BALANCING_FOR_SIDEKIQ', 'true') + end + + it 'returns true when Sidekiq is being used' do + allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(true) + + expect(described_class.enable?).to eq(true) + end + end + end + + describe '.configured?' do + before do + clear_load_balancing_configuration + end + + it 'returns true when Sidekiq is being used' do + allow(described_class).to receive(:hosts).and_return(%w(foo)) + allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(true) + expect(described_class.configured?).to eq(true) + end + + it 'returns true when service discovery is enabled in Sidekiq' do + allow(described_class).to receive(:hosts).and_return([]) + allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(true) + + allow(described_class) + .to receive(:service_discovery_enabled?) + .and_return(true) + + expect(described_class.configured?).to eq(true) + end + + it 'returns false when neither service discovery nor hosts are configured' do + allow(described_class).to receive(:hosts).and_return([]) + + allow(described_class) + .to receive(:service_discovery_enabled?) + .and_return(false) + + expect(described_class.configured?).to eq(false) + end + end + + describe '.configure_proxy' do + it 'configures the connection proxy' do + allow(ActiveRecord::Base.singleton_class).to receive(:prepend) + + described_class.configure_proxy + + expect(ActiveRecord::Base.singleton_class).to have_received(:prepend) + .with(Gitlab::Database::LoadBalancing::ActiveRecordProxy) + end + end + + describe '.active_record_models' do + it 'returns an Array' do + expect(described_class.active_record_models).to be_an_instance_of(Array) + end + end + + describe '.service_discovery_enabled?' do + it 'returns true if service discovery is enabled' do + allow(described_class) + .to receive(:configuration) + .and_return('discover' => { 'record' => 'foo' }) + + expect(described_class.service_discovery_enabled?).to eq(true) + end + + it 'returns false if service discovery is disabled' do + expect(described_class.service_discovery_enabled?).to eq(false) + end + end + + describe '.service_discovery_configuration' do + context 'when no configuration is provided' do + it 'returns a default configuration Hash' do + expect(described_class.service_discovery_configuration).to eq( + nameserver: 'localhost', + port: 8600, + record: nil, + record_type: 'A', + interval: 60, + disconnect_timeout: 120, + use_tcp: false + ) + end + end + + context 'when configuration is provided' do + it 'returns a Hash including the custom configuration' do + allow(described_class) + .to receive(:configuration) + .and_return('discover' => { 'record' => 'foo', 'record_type' => 'SRV' }) + + expect(described_class.service_discovery_configuration).to eq( + nameserver: 'localhost', + port: 8600, + record: 'foo', + record_type: 'SRV', + interval: 60, + disconnect_timeout: 120, + use_tcp: false + ) + end + end + end + + describe '.start_service_discovery' do + it 'does not start if service discovery is disabled' do + expect(Gitlab::Database::LoadBalancing::ServiceDiscovery) + .not_to receive(:new) + + described_class.start_service_discovery + end + + it 'starts service discovery if enabled' do + allow(described_class) + .to receive(:service_discovery_enabled?) + .and_return(true) + + instance = double(:instance) + + expect(Gitlab::Database::LoadBalancing::ServiceDiscovery) + .to receive(:new) + .with(an_instance_of(Hash)) + .and_return(instance) + + expect(instance) + .to receive(:start) + + described_class.start_service_discovery + end + 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 + + it 'returns primary' do + expect(described_class.db_role_for_connection(connection)).to be(: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)) } + + before do + allow(ActiveRecord::Base.singleton_class).to receive(:prepend) + + 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) + + subject.configure_proxy(proxy) + end + + context 'when the load balancer returns :replica' 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) + end + end + + context 'when the load balancer returns :primary' 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) + end + end + end + end + + # For such an important module like LoadBalancing, full mocking is not + # enough. This section implements some integration tests to test a full flow + # of the load balancer. + # - A real model with a table backed behind is defined + # - The load balancing module is set up for this module only, as to prevent + # breaking other tests. The replica configuration is cloned from the test + # configuraiton. + # - In each test, we listen to the SQL queries (via sql.active_record + # instrumentation) while triggering real queries from the defined model. + # - We assert the desinations (replica/primary) of the queries in order. + describe 'LoadBalancing integration tests', :delete do + before(:all) do + ActiveRecord::Schema.define do + create_table :load_balancing_test, force: true do |t| + t.string :name, null: true + end + end + end + + after(:all) do + ActiveRecord::Schema.define do + drop_table :load_balancing_test, force: true + end + end + + shared_context 'LoadBalancing setup' do + let(:development_db_config) { ActiveRecord::Base.configurations.configs_for(env_name: 'development').first.configuration_hash } + let(:hosts) { [development_db_config[:host]] } + let(:model) do + Class.new(ApplicationRecord) do + self.table_name = "load_balancing_test" + end + end + + before do + # Preloading testing class + model.singleton_class.prepend ::Gitlab::Database::LoadBalancing::ActiveRecordProxy + + # Setup load balancing + clear_load_balancing_configuration + allow(ActiveRecord::Base.singleton_class).to receive(:prepend) + subject.configure_proxy(::Gitlab::Database::LoadBalancing::ConnectionProxy.new(hosts)) + + original_db_config = Gitlab::Database.config + modified_db_config = original_db_config.merge(load_balancing: { hosts: hosts }) + allow(Gitlab::Database).to receive(:config).and_return(modified_db_config) + + ::Gitlab::Database::LoadBalancing::Session.clear_session + end + end + + where(:queries, :include_transaction, :expected_results) do + [ + # Read methods + [-> { model.first }, false, [:replica]], + [-> { model.find_by(id: 123) }, false, [:replica]], + [-> { model.where(name: 'hello').to_a }, false, [:replica]], + + # Write methods + [-> { model.create!(name: 'test1') }, false, [:primary]], + [ + -> { + instance = model.create!(name: 'test1') + instance.update!(name: 'test2') + }, + false, [:primary, :primary] + ], + [-> { model.update_all(name: 'test2') }, false, [:primary]], + [ + -> { + instance = model.create!(name: 'test1') + instance.destroy! + }, + false, [:primary, :primary] + ], + [-> { model.delete_all }, false, [:primary]], + + # Custom query + [-> { model.connection.exec_query('SELECT 1').to_a }, false, [:primary]], + + # Reads after a write + [ + -> { + model.first + model.create!(name: 'test1') + model.first + model.find_by(name: 'test1') + }, + false, [:replica, :primary, :primary, :primary] + ], + + # Inside a transaction + [ + -> { + model.transaction do + model.find_by(name: 'test1') + model.create!(name: 'test1') + instance = model.find_by(name: 'test1') + instance.update!(name: 'test2') + end + model.find_by(name: 'test1') + }, + true, [:primary, :primary, :primary, :primary, :primary, :primary, :primary] + ], + + # Nested transaction + [ + -> { + model.transaction do + model.transaction do + model.create!(name: 'test1') + end + model.update_all(name: 'test2') + end + model.find_by(name: 'test1') + }, + true, [:primary, :primary, :primary, :primary, :primary] + ], + + # Read-only transaction + [ + -> { + model.transaction do + model.first + model.where(name: 'test1').to_a + end + }, + true, [:primary, :primary, :primary, :primary] + ], + + # use_primary + [ + -> { + ::Gitlab::Database::LoadBalancing::Session.current.use_primary do + model.first + model.where(name: 'test1').to_a + end + model.first + }, + false, [:primary, :primary, :replica] + ], + + # use_primary! + [ + -> { + model.first + ::Gitlab::Database::LoadBalancing::Session.current.use_primary! + model.where(name: 'test1').to_a + }, + false, [:replica, :primary] + ], + + # use_replicas_for_read_queries does not affect read queries + [ + -> { + ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do + model.where(name: 'test1').to_a + end + }, + false, [:replica] + ], + + # use_replicas_for_read_queries does not affect write queries + [ + -> { + ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do + model.create!(name: 'test1') + end + }, + false, [:primary] + ], + + # use_replicas_for_read_queries does not affect ambiguous queries + [ + -> { + ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do + model.connection.exec_query("SELECT 1") + end + }, + false, [:primary] + ], + + # use_replicas_for_read_queries ignores use_primary! for read queries + [ + -> { + ::Gitlab::Database::LoadBalancing::Session.current.use_primary! + ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do + model.where(name: 'test1').to_a + end + }, + false, [:replica] + ], + + # use_replicas_for_read_queries adheres use_primary! for write queries + [ + -> { + ::Gitlab::Database::LoadBalancing::Session.current.use_primary! + ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do + model.create!(name: 'test1') + end + }, + false, [:primary] + ], + + # use_replicas_for_read_queries adheres use_primary! for ambiguous queries + [ + -> { + ::Gitlab::Database::LoadBalancing::Session.current.use_primary! + ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do + model.connection.exec_query('SELECT 1') + end + }, + false, [:primary] + ], + + # use_replicas_for_read_queries ignores use_primary blocks + [ + -> { + ::Gitlab::Database::LoadBalancing::Session.current.use_primary do + ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do + model.where(name: 'test1').to_a + end + end + }, + false, [:replica] + ], + + # use_replicas_for_read_queries ignores a session already performed write + [ + -> { + ::Gitlab::Database::LoadBalancing::Session.current.write! + ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do + model.where(name: 'test1').to_a + end + }, + false, [:replica] + ], + + # fallback_to_replicas_for_ambiguous_queries + [ + -> { + ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do + model.first + model.where(name: 'test1').to_a + end + }, + false, [:replica, :replica] + ], + + # fallback_to_replicas_for_ambiguous_queries for read-only transaction + [ + -> { + ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do + model.transaction do + model.first + model.where(name: 'test1').to_a + end + end + }, + false, [:replica, :replica] + ], + + # A custom read query inside fallback_to_replicas_for_ambiguous_queries + [ + -> { + ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do + model.connection.exec_query("SELECT 1") + end + }, + false, [:replica] + ], + + # A custom read query inside a transaction fallback_to_replicas_for_ambiguous_queries + [ + -> { + ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do + model.transaction do + model.connection.exec_query("SET LOCAL statement_timeout = 5000") + model.count + end + end + }, + true, [:replica, :replica, :replica, :replica] + ], + + # fallback_to_replicas_for_ambiguous_queries after a write + [ + -> { + model.create!(name: 'Test1') + ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do + model.connection.exec_query("SELECT 1") + end + }, + false, [:primary, :primary] + ], + + # fallback_to_replicas_for_ambiguous_queries after use_primary! + [ + -> { + ::Gitlab::Database::LoadBalancing::Session.current.use_primary! + ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do + model.connection.exec_query("SELECT 1") + end + }, + false, [:primary] + ], + + # fallback_to_replicas_for_ambiguous_queries inside use_primary + [ + -> { + ::Gitlab::Database::LoadBalancing::Session.current.use_primary do + ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do + model.connection.exec_query("SELECT 1") + end + end + }, + false, [:primary] + ], + + # use_primary inside fallback_to_replicas_for_ambiguous_queries + [ + -> { + ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do + ::Gitlab::Database::LoadBalancing::Session.current.use_primary do + model.connection.exec_query("SELECT 1") + end + end + }, + false, [:primary] + ], + + # A write query inside fallback_to_replicas_for_ambiguous_queries + [ + -> { + ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do + model.connection.exec_query("SELECT 1") + model.delete_all + model.connection.exec_query("SELECT 1") + end + }, + false, [:replica, :primary, :primary] + ], + + # use_replicas_for_read_queries incorporates with fallback_to_replicas_for_ambiguous_queries + [ + -> { + ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do + ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do + model.connection.exec_query('SELECT 1') + model.where(name: 'test1').to_a + end + end + }, + false, [:replica, :replica] + ] + ] + end + + with_them do + include_context 'LoadBalancing setup' + + it 'redirects queries to the right roles' do + roles = [] + + subscriber = ActiveSupport::Notifications.subscribe('sql.active_record') do |event| + payload = event.payload + + assert = + if payload[:name] == 'SCHEMA' + false + elsif payload[:name] == 'SQL' # Custom query + true + else + keywords = %w[load_balancing_test] + keywords += %w[begin commit] if include_transaction + keywords.any? { |keyword| payload[:sql].downcase.include?(keyword) } + end + + if assert + db_role = ::Gitlab::Database::LoadBalancing.db_role_for_connection(payload[:connection]) + roles << db_role + end + end + + self.instance_exec(&queries) + + expect(roles).to eql(expected_results) + ensure + ActiveSupport::Notifications.unsubscribe(subscriber) if subscriber + end + end + + context 'custom connection handling' do + where(:queries, :expected_role) do + [ + # Reload cache. The schema loading queries should be handled by + # primary. + [ + -> { + model.connection.clear_cache! + model.connection.schema_cache.add('users') + model.connection.pool.release_connection + }, + :primary + ], + + # Call model's connection method + [ + -> { + connection = model.connection + connection.select_one('SELECT 1') + connection.pool.release_connection + }, + :replica + ], + + # Retrieve connection via #retrieve_connection + [ + -> { + connection = model.retrieve_connection + connection.select_one('SELECT 1') + connection.pool.release_connection + }, + :primary + ] + ] + end + + with_them do + include_context 'LoadBalancing setup' + + it 'redirects queries to the right roles' do + roles = [] + + subscriber = ActiveSupport::Notifications.subscribe('sql.active_record') do |event| + role = ::Gitlab::Database::LoadBalancing.db_role_for_connection(event.payload[:connection]) + roles << role if role.present? + end + + self.instance_exec(&queries) + + expect(roles).to all(eql(expected_role)) + ensure + ActiveSupport::Notifications.unsubscribe(subscriber) if subscriber + end + end + end + + context 'a write inside a transaction inside fallback_to_replicas_for_ambiguous_queries block' do + include_context 'LoadBalancing setup' + + it 'raises an exception' do + expect do + ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do + model.transaction do + model.first + model.create!(name: 'hello') + end + end + end.to raise_error(Gitlab::Database::LoadBalancing::ConnectionProxy::WriteInsideReadOnlyTransactionError) + end + end + end +end diff --git a/spec/lib/gitlab/database/migration_helpers_spec.rb b/spec/lib/gitlab/database/migration_helpers_spec.rb index 40720628a89..f0ea07646fb 100644 --- a/spec/lib/gitlab/database/migration_helpers_spec.rb +++ b/spec/lib/gitlab/database/migration_helpers_spec.rb @@ -2001,6 +2001,41 @@ RSpec.describe Gitlab::Database::MigrationHelpers do end end + describe '#ensure_batched_background_migration_is_finished' do + let(:configuration) do + { + job_class_name: 'CopyColumnUsingBackgroundMigrationJob', + table_name: :events, + column_name: :id, + job_arguments: [[:id], [:id_convert_to_bigint]] + } + end + + subject(:ensure_batched_background_migration_is_finished) { model.ensure_batched_background_migration_is_finished(**configuration) } + + it 'raises an error when migration exists and is not marked as finished' do + create(:batched_background_migration, configuration.merge(status: :active)) + + expect { ensure_batched_background_migration_is_finished } + .to raise_error "Expected batched background migration for the given configuration to be marked as 'finished', but it is 'active': #{configuration}" + end + + it 'does not raise error when migration exists and is marked as finished' do + create(:batched_background_migration, configuration.merge(status: :finished)) + + expect { ensure_batched_background_migration_is_finished } + .not_to raise_error + end + + it 'logs a warning when migration does not exist' do + expect(Gitlab::AppLogger).to receive(:warn) + .with("Could not find batched background migration for the given configuration: #{configuration}") + + expect { ensure_batched_background_migration_is_finished } + .not_to raise_error + end + end + describe '#index_exists_by_name?' do it 'returns true if an index exists' do ActiveRecord::Base.connection.execute( diff --git a/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb b/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb index c6d456964cf..e096e7f6e91 100644 --- a/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb +++ b/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb @@ -242,6 +242,98 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do end end + describe '#requeue_background_migration_jobs_by_range_at_intervals' do + let!(:job_class_name) { 'TestJob' } + let!(:pending_job_1) { create(:background_migration_job, class_name: job_class_name, status: :pending, arguments: [1, 2]) } + let!(:pending_job_2) { create(:background_migration_job, class_name: job_class_name, status: :pending, arguments: [3, 4]) } + let!(:successful_job_1) { create(:background_migration_job, class_name: job_class_name, status: :succeeded, arguments: [5, 6]) } + let!(:successful_job_2) { create(:background_migration_job, class_name: job_class_name, status: :succeeded, arguments: [7, 8]) } + + around do |example| + freeze_time do + Sidekiq::Testing.fake! do + example.run + end + end + end + + subject { model.requeue_background_migration_jobs_by_range_at_intervals(job_class_name, 10.minutes) } + + it 'returns the expected duration' do + expect(subject).to eq(20.minutes) + end + + context 'when nothing is queued' do + subject { model.requeue_background_migration_jobs_by_range_at_intervals('FakeJob', 10.minutes) } + + it 'returns expected duration of zero when nothing gets queued' do + expect(subject).to eq(0) + end + end + + it 'queues pending jobs' do + subject + + expect(BackgroundMigrationWorker.jobs[0]['args']).to eq([job_class_name, [1, 2]]) + expect(BackgroundMigrationWorker.jobs[0]['at']).to be_nil + expect(BackgroundMigrationWorker.jobs[1]['args']).to eq([job_class_name, [3, 4]]) + expect(BackgroundMigrationWorker.jobs[1]['at']).to eq(10.minutes.from_now.to_f) + end + + context 'with batch_size option' do + subject { model.requeue_background_migration_jobs_by_range_at_intervals(job_class_name, 10.minutes, batch_size: 1) } + + it 'returns the expected duration' do + expect(subject).to eq(20.minutes) + end + + it 'queues pending jobs' do + subject + + expect(BackgroundMigrationWorker.jobs[0]['args']).to eq([job_class_name, [1, 2]]) + expect(BackgroundMigrationWorker.jobs[0]['at']).to be_nil + expect(BackgroundMigrationWorker.jobs[1]['args']).to eq([job_class_name, [3, 4]]) + expect(BackgroundMigrationWorker.jobs[1]['at']).to eq(10.minutes.from_now.to_f) + end + + it 'retrieve jobs in batches' do + jobs = double('jobs') + expect(Gitlab::Database::BackgroundMigrationJob).to receive(:pending) { jobs } + allow(jobs).to receive(:where).with(class_name: job_class_name) { jobs } + expect(jobs).to receive(:each_batch).with(of: 1) + + subject + end + end + + context 'with initial_delay option' do + let_it_be(:initial_delay) { 3.minutes } + + subject { model.requeue_background_migration_jobs_by_range_at_intervals(job_class_name, 10.minutes, initial_delay: initial_delay) } + + it 'returns the expected duration' do + expect(subject).to eq(23.minutes) + end + + it 'queues pending jobs' do + subject + + expect(BackgroundMigrationWorker.jobs[0]['args']).to eq([job_class_name, [1, 2]]) + expect(BackgroundMigrationWorker.jobs[0]['at']).to eq(3.minutes.from_now.to_f) + expect(BackgroundMigrationWorker.jobs[1]['args']).to eq([job_class_name, [3, 4]]) + expect(BackgroundMigrationWorker.jobs[1]['at']).to eq(13.minutes.from_now.to_f) + end + + context 'when nothing is queued' do + subject { model.requeue_background_migration_jobs_by_range_at_intervals('FakeJob', 10.minutes) } + + it 'returns expected duration of zero when nothing gets queued' do + expect(subject).to eq(0) + end + end + end + end + describe '#perform_background_migration_inline?' do it 'returns true in a test environment' do stub_rails_env('test') @@ -269,6 +361,38 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do allow(Gitlab::Database::PgClass).to receive(:for_table).and_call_original end + context 'when such migration already exists' do + it 'does not create duplicate migration' do + create( + :batched_background_migration, + job_class_name: 'MyJobClass', + table_name: :projects, + column_name: :id, + interval: 10.minutes, + min_value: 5, + max_value: 1005, + batch_class_name: 'MyBatchClass', + batch_size: 200, + sub_batch_size: 20, + job_arguments: [[:id], [:id_convert_to_bigint]] + ) + + expect do + model.queue_batched_background_migration( + 'MyJobClass', + :projects, + :id, + [:id], [:id_convert_to_bigint], + job_interval: 5.minutes, + batch_min_value: 5, + batch_max_value: 1000, + batch_class_name: 'MyBatchClass', + batch_size: 100, + sub_batch_size: 10) + end.not_to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count } + end + end + it 'creates the database record for the migration' do expect(Gitlab::Database::PgClass).to receive(:for_table).with(:projects).and_return(pgclass_info) diff --git a/spec/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers_spec.rb b/spec/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers_spec.rb index 79ddb450d7a..4f1d6302331 100644 --- a/spec/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers_spec.rb +++ b/spec/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers_spec.rb @@ -580,7 +580,7 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::TableManagementHe it 'idempotently cleans up after failed background migrations' do expect(partitioned_model.count).to eq(0) - partitioned_model.insert!(record2.attributes) + partitioned_model.insert(record2.attributes, unique_by: [:id, :created_at]) expect_next_instance_of(Gitlab::Database::PartitioningMigrationHelpers::BackfillPartitionedTable) do |backfill| allow(backfill).to receive(:transaction_open?).and_return(false) diff --git a/spec/lib/gitlab/database/postgresql_adapter/type_map_cache_spec.rb b/spec/lib/gitlab/database/postgresql_adapter/type_map_cache_spec.rb new file mode 100644 index 00000000000..e9c512f94bb --- /dev/null +++ b/spec/lib/gitlab/database/postgresql_adapter/type_map_cache_spec.rb @@ -0,0 +1,68 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe Gitlab::Database::PostgresqlAdapter::TypeMapCache do + let(:db_config) { ActiveRecord::Base.configurations.configs_for(env_name: 'test', name: 'primary').configuration_hash } + let(:adapter_class) { ActiveRecord::ConnectionAdapters::PostgreSQLAdapter } + + before do + adapter_class.type_map_cache.clear + end + + describe '#initialize_type_map' do + it 'caches loading of types in memory' do + recorder_without_cache = ActiveRecord::QueryRecorder.new(skip_schema_queries: false) { initialize_connection.disconnect! } + expect(recorder_without_cache.log).to include(a_string_matching(/FROM pg_type/)).twice + + recorder_with_cache = ActiveRecord::QueryRecorder.new(skip_schema_queries: false) { initialize_connection.disconnect! } + + expect(recorder_with_cache.count).to be < recorder_without_cache.count + + # There's still one pg_type query left here because `#add_pg_decoders` executes another pg_type query + # in https://github.com/rails/rails/blob/v6.1.3.2/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb#L912. + # This query is much cheaper because it only returns very few records. + expect(recorder_with_cache.log).to include(a_string_matching(/FROM pg_type/)).once + end + + it 'only reuses the cache if the connection parameters are exactly the same' do + initialize_connection.disconnect! + + other_config = db_config.dup + other_config[:connect_timeout] = db_config[:connect_timeout].to_i + 10 + + recorder = ActiveRecord::QueryRecorder.new(skip_schema_queries: false) { initialize_connection(other_config).disconnect! } + + expect(recorder.log).to include(a_string_matching(/FROM pg_type/)).twice + end + end + + describe '#reload_type_map' do + it 'clears the cache and executes the type map query again' do + initialize_connection.disconnect! + + connection = initialize_connection + recorder = ActiveRecord::QueryRecorder.new(skip_schema_queries: false) { connection.reload_type_map } + + expect(recorder.log).to include(a_string_matching(/FROM pg_type/)).once + end + end + + # Based on https://github.com/rails/rails/blob/v6.1.3.2/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb#L36-L41 + def initialize_connection(config = db_config) + conn_params = config.symbolize_keys.compact + + conn_params[:user] = conn_params.delete(:username) if conn_params[:username] + conn_params[:dbname] = conn_params.delete(:database) if conn_params[:database] + + valid_conn_param_keys = PG::Connection.conndefaults_hash.keys + [:requiressl] + conn_params.slice!(*valid_conn_param_keys) + + adapter_class.new( + adapter_class.new_client(conn_params), + ActiveRecord::Base.logger, + conn_params, + config + ) + end +end diff --git a/spec/lib/gitlab/database/with_lock_retries_spec.rb b/spec/lib/gitlab/database/with_lock_retries_spec.rb index b08f39fc92a..df2c506e163 100644 --- a/spec/lib/gitlab/database/with_lock_retries_spec.rb +++ b/spec/lib/gitlab/database/with_lock_retries_spec.rb @@ -242,10 +242,10 @@ RSpec.describe Gitlab::Database::WithLockRetries do let(:timing_configuration) { [[0.015.seconds, 0.025.seconds], [0.015.seconds, 0.025.seconds]] } # 15ms, 25ms it 'executes `SET LOCAL lock_timeout` using the configured timeout value in milliseconds' do - expect(ActiveRecord::Base.connection).to receive(:execute).with("SAVEPOINT active_record_1").and_call_original - expect(ActiveRecord::Base.connection).to receive(:execute).with('RESET idle_in_transaction_session_timeout; RESET lock_timeout').and_call_original + expect(ActiveRecord::Base.connection).to receive(:execute).with("RESET idle_in_transaction_session_timeout; RESET lock_timeout").and_call_original + expect(ActiveRecord::Base.connection).to receive(:execute).with("SAVEPOINT active_record_1", "TRANSACTION").and_call_original expect(ActiveRecord::Base.connection).to receive(:execute).with("SET LOCAL lock_timeout TO '15ms'").and_call_original - expect(ActiveRecord::Base.connection).to receive(:execute).with("RELEASE SAVEPOINT active_record_1").and_call_original + expect(ActiveRecord::Base.connection).to receive(:execute).with("RELEASE SAVEPOINT active_record_1", "TRANSACTION").and_call_original subject.run { } end |