diff options
Diffstat (limited to 'spec/lib/gitlab/database/load_balancing')
13 files changed, 3149 insertions, 0 deletions
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 |