697 lines
22 KiB
Ruby
697 lines
22 KiB
Ruby
# frozen_string_literal: true
|
|
|
|
require 'spec_helper'
|
|
|
|
RSpec.describe Gitlab::Database::LoadBalancing do
|
|
describe '.proxy' do
|
|
before do
|
|
@previous_proxy = ActiveRecord::Base.load_balancing_proxy
|
|
|
|
ActiveRecord::Base.load_balancing_proxy = connection_proxy
|
|
end
|
|
|
|
after do
|
|
ActiveRecord::Base.load_balancing_proxy = @previous_proxy
|
|
end
|
|
|
|
context 'when configured' do
|
|
let(:connection_proxy) { double(:connection_proxy) }
|
|
|
|
it 'returns the connection proxy' do
|
|
expect(subject.proxy).to eq(connection_proxy)
|
|
end
|
|
end
|
|
|
|
context 'when not configured' do
|
|
let(:connection_proxy) { nil }
|
|
|
|
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 the configuration for the load balancer' do
|
|
raw = ActiveRecord::Base.connection_db_config.configuration_hash
|
|
cfg = described_class.configuration
|
|
|
|
# There isn't much to test here as the load balancing settings might not
|
|
# (and likely aren't) set when running tests.
|
|
expect(cfg.pool_size).to eq(raw[:pool])
|
|
end
|
|
end
|
|
|
|
describe '.enable?' do
|
|
before do
|
|
allow(described_class.configuration)
|
|
.to receive(:hosts)
|
|
.and_return(%w(foo))
|
|
end
|
|
|
|
it 'returns false when no hosts are specified' do
|
|
allow(described_class.configuration).to receive(:hosts).and_return([])
|
|
|
|
expect(described_class.enable?).to eq(false)
|
|
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
|
|
|
|
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.configuration).to receive(:hosts).and_return([])
|
|
allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(false)
|
|
|
|
allow(described_class.configuration)
|
|
.to receive(:service_discovery_enabled?)
|
|
.and_return(true)
|
|
|
|
expect(described_class.enable?).to eq(true)
|
|
end
|
|
end
|
|
|
|
describe '.configured?' do
|
|
it 'returns true when hosts are configured' do
|
|
allow(described_class.configuration)
|
|
.to receive(:hosts)
|
|
.and_return(%w[foo])
|
|
|
|
expect(described_class.configured?).to eq(true)
|
|
end
|
|
|
|
it 'returns true when service discovery is enabled' do
|
|
allow(described_class.configuration).to receive(:hosts).and_return([])
|
|
allow(described_class.configuration)
|
|
.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.configuration).to receive(:hosts).and_return([])
|
|
allow(described_class.configuration)
|
|
.to receive(:service_discovery_enabled?)
|
|
.and_return(false)
|
|
|
|
expect(described_class.configured?).to eq(false)
|
|
end
|
|
end
|
|
|
|
describe '.configure_proxy' do
|
|
before do
|
|
allow(ActiveRecord::Base).to receive(:load_balancing_proxy=)
|
|
end
|
|
|
|
it 'configures the connection proxy' do
|
|
described_class.configure_proxy
|
|
|
|
expect(ActiveRecord::Base).to have_received(:load_balancing_proxy=)
|
|
.with(Gitlab::Database::LoadBalancing::ConnectionProxy)
|
|
end
|
|
|
|
context 'when service discovery is enabled' do
|
|
it 'runs initial service discovery when configuring the connection proxy' do
|
|
discover = instance_spy(Gitlab::Database::LoadBalancing::ServiceDiscovery)
|
|
|
|
allow(described_class.configuration)
|
|
.to receive(:service_discovery)
|
|
.and_return({ record: 'foo' })
|
|
|
|
expect(Gitlab::Database::LoadBalancing::ServiceDiscovery)
|
|
.to receive(:new)
|
|
.with(
|
|
an_instance_of(Gitlab::Database::LoadBalancing::LoadBalancer),
|
|
an_instance_of(Hash)
|
|
)
|
|
.and_return(discover)
|
|
|
|
expect(discover).to receive(:perform_service_discovery)
|
|
|
|
described_class.configure_proxy
|
|
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.configuration)
|
|
.to receive(:service_discovery_enabled?)
|
|
.and_return(true)
|
|
|
|
instance = double(:instance)
|
|
config = Gitlab::Database::LoadBalancing::Configuration
|
|
.new(ActiveRecord::Base)
|
|
lb = Gitlab::Database::LoadBalancing::LoadBalancer.new(config)
|
|
proxy = Gitlab::Database::LoadBalancing::ConnectionProxy.new(lb)
|
|
|
|
allow(described_class)
|
|
.to receive(:proxy)
|
|
.and_return(proxy)
|
|
|
|
expect(Gitlab::Database::LoadBalancing::ServiceDiscovery)
|
|
.to receive(:new)
|
|
.with(lb, 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
|
|
context 'when the load balancing is not configured' do
|
|
let(:connection) { ActiveRecord::Base.connection }
|
|
|
|
it 'returns primary' do
|
|
expect(described_class.db_role_for_connection(connection)).to eq(:primary)
|
|
end
|
|
end
|
|
|
|
context 'when the NullPool is used for connection' do
|
|
let(:pool) { ActiveRecord::ConnectionAdapters::NullPool.new }
|
|
let(:connection) { double(:connection, pool: pool) }
|
|
|
|
it 'returns unknown' do
|
|
expect(described_class.db_role_for_connection(connection)).to eq(:unknown)
|
|
end
|
|
end
|
|
|
|
context 'when the load balancing is configured' do
|
|
let(:db_host) { ActiveRecord::Base.connection_pool.db_config.host }
|
|
let(:config) do
|
|
Gitlab::Database::LoadBalancing::Configuration
|
|
.new(ActiveRecord::Base, [db_host])
|
|
end
|
|
|
|
let(:load_balancer) { described_class::LoadBalancer.new(config) }
|
|
let(:proxy) { described_class::ConnectionProxy.new(load_balancer) }
|
|
|
|
context 'when a proxy connection is used' do
|
|
it 'returns :unknown' do
|
|
expect(described_class.db_role_for_connection(proxy)).to eq(:unknown)
|
|
end
|
|
end
|
|
|
|
context 'when a read connection is used' do
|
|
it 'returns :replica' do
|
|
proxy.load_balancer.read do |connection|
|
|
expect(described_class.db_role_for_connection(connection)).to eq(:replica)
|
|
end
|
|
end
|
|
end
|
|
|
|
context 'when a read_write connection is used' do
|
|
it 'returns :primary' do
|
|
proxy.load_balancer.read_write do |connection|
|
|
expect(described_class.db_role_for_connection(connection)).to eq(:primary)
|
|
end
|
|
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', :db_load_balancing, :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
|
|
|
|
let(:model) do
|
|
Class.new(ApplicationRecord) do
|
|
self.table_name = "load_balancing_test"
|
|
end
|
|
end
|
|
|
|
before do
|
|
model.singleton_class.prepend ::Gitlab::Database::LoadBalancing::ActiveRecordProxy
|
|
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
|
|
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
|
|
it 'redirects queries to the right roles' do
|
|
roles = []
|
|
|
|
# If we don't run any queries, the pool may be a NullPool. This can
|
|
# result in some tests reporting a role as `:unknown`, even though the
|
|
# tests themselves are correct.
|
|
#
|
|
# To prevent this from happening we simply run a simple query to
|
|
# ensure the proper pool type is put in place. The exact query doesn't
|
|
# matter, provided it actually runs a query and thus creates a proper
|
|
# connection pool.
|
|
model.count
|
|
|
|
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
|
|
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
|