561 lines
18 KiB
Ruby
561 lines
18 KiB
Ruby
# frozen_string_literal: true
|
|
|
|
require 'spec_helper'
|
|
|
|
RSpec.describe Gitlab::Database::LoadBalancing do
|
|
describe '.base_models' do
|
|
it 'returns the models to apply load balancing to' do
|
|
models = described_class.base_models
|
|
|
|
expect(models).to include(ActiveRecord::Base)
|
|
|
|
if Gitlab::Database.has_config?(:ci)
|
|
expect(models).to include(Ci::ApplicationRecord)
|
|
end
|
|
end
|
|
|
|
it 'returns the models as a frozen array' do
|
|
expect(described_class.base_models).to be_frozen
|
|
end
|
|
end
|
|
|
|
describe '.each_load_balancer' do
|
|
it 'yields every load balancer to the supplied block' do
|
|
lbs = []
|
|
|
|
described_class.each_load_balancer do |lb|
|
|
lbs << lb
|
|
end
|
|
|
|
expect(lbs.length).to eq(described_class.base_models.length)
|
|
end
|
|
|
|
it 'returns an Enumerator when no block is given' do
|
|
res = described_class.each_load_balancer
|
|
|
|
expect(res.next)
|
|
.to be_an_instance_of(Gitlab::Database::LoadBalancing::LoadBalancer)
|
|
end
|
|
end
|
|
|
|
describe '.primary_only?' do
|
|
it 'returns true if all load balancers have no replicas' do
|
|
described_class.each_load_balancer do |lb|
|
|
allow(lb).to receive(:primary_only?).and_return(true)
|
|
end
|
|
|
|
expect(described_class.primary_only?).to eq(true)
|
|
end
|
|
|
|
it 'returns false if at least one has replicas' do
|
|
described_class.each_load_balancer.with_index do |lb, index|
|
|
allow(lb).to receive(:primary_only?).and_return(index != 0)
|
|
end
|
|
|
|
expect(described_class.primary_only?).to eq(false)
|
|
end
|
|
end
|
|
|
|
describe '.release_hosts' do
|
|
it 'releases the host of every load balancer' do
|
|
described_class.each_load_balancer do |lb|
|
|
expect(lb).to receive(:release_host)
|
|
end
|
|
|
|
described_class.release_hosts
|
|
end
|
|
end
|
|
|
|
describe '.db_role_for_connection' do
|
|
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
|
|
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
|
|
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', :database_replica, :delete do
|
|
before(:all) do
|
|
ActiveRecord::Schema.define do
|
|
create_table :_test_load_balancing_test, force: true do |t|
|
|
t.string :name, null: true
|
|
end
|
|
end
|
|
end
|
|
|
|
after(:all) do
|
|
ActiveRecord::Schema.define do
|
|
drop_table :_test_load_balancing_test, force: true
|
|
end
|
|
end
|
|
|
|
let(:model) do
|
|
Class.new(ApplicationRecord) do
|
|
self.table_name = "_test_load_balancing_test"
|
|
end
|
|
end
|
|
|
|
where(:queries, :include_transaction, :expected_results) do
|
|
[
|
|
# Read methods
|
|
[-> { model.first }, false, [:replica]],
|
|
[-> { model.find_by(id: 123) }, false, [:replica]],
|
|
[-> { model.where(name: 'hello').to_a }, false, [:replica]],
|
|
|
|
# Write methods
|
|
[-> { model.create!(name: 'test1') }, false, [:primary]],
|
|
[
|
|
-> {
|
|
instance = model.create!(name: 'test1')
|
|
instance.update!(name: 'test2')
|
|
},
|
|
false, [:primary, :primary]
|
|
],
|
|
[-> { model.update_all(name: 'test2') }, false, [:primary]],
|
|
[
|
|
-> {
|
|
instance = model.create!(name: 'test1')
|
|
instance.destroy!
|
|
},
|
|
false, [:primary, :primary]
|
|
],
|
|
[-> { model.delete_all }, false, [:primary]],
|
|
|
|
# Custom query
|
|
[-> { model.connection.exec_query('SELECT 1').to_a }, false, [:primary]],
|
|
|
|
# Reads after a write
|
|
[
|
|
-> {
|
|
model.first
|
|
model.create!(name: 'test1')
|
|
model.first
|
|
model.find_by(name: 'test1')
|
|
},
|
|
false, [:replica, :primary, :primary, :primary]
|
|
],
|
|
|
|
# Inside a transaction
|
|
[
|
|
-> {
|
|
model.transaction do
|
|
model.find_by(name: 'test1')
|
|
model.create!(name: 'test1')
|
|
instance = model.find_by(name: 'test1')
|
|
instance.update!(name: 'test2')
|
|
end
|
|
model.find_by(name: 'test1')
|
|
},
|
|
true, [:primary, :primary, :primary, :primary, :primary, :primary, :primary]
|
|
],
|
|
|
|
# Nested transaction
|
|
[
|
|
-> {
|
|
model.transaction do
|
|
model.transaction do
|
|
model.create!(name: 'test1')
|
|
end
|
|
model.update_all(name: 'test2')
|
|
end
|
|
model.find_by(name: 'test1')
|
|
},
|
|
true, [:primary, :primary, :primary, :primary, :primary]
|
|
],
|
|
|
|
# Read-only transaction
|
|
[
|
|
-> {
|
|
model.transaction do
|
|
model.first
|
|
model.where(name: 'test1').to_a
|
|
end
|
|
},
|
|
true, [:primary, :primary, :primary, :primary]
|
|
],
|
|
|
|
# use_primary
|
|
[
|
|
-> {
|
|
::Gitlab::Database::LoadBalancing::Session.current.use_primary do
|
|
model.first
|
|
model.where(name: 'test1').to_a
|
|
end
|
|
model.first
|
|
},
|
|
false, [:primary, :primary, :replica]
|
|
],
|
|
|
|
# use_primary!
|
|
[
|
|
-> {
|
|
model.first
|
|
::Gitlab::Database::LoadBalancing::Session.current.use_primary!
|
|
model.where(name: 'test1').to_a
|
|
},
|
|
false, [:replica, :primary]
|
|
],
|
|
|
|
# use_replicas_for_read_queries does not affect read queries
|
|
[
|
|
-> {
|
|
::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
|
|
model.where(name: 'test1').to_a
|
|
end
|
|
},
|
|
false, [:replica]
|
|
],
|
|
|
|
# use_replicas_for_read_queries does not affect write queries
|
|
[
|
|
-> {
|
|
::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
|
|
model.create!(name: 'test1')
|
|
end
|
|
},
|
|
false, [:primary]
|
|
],
|
|
|
|
# use_replicas_for_read_queries does not affect ambiguous queries
|
|
[
|
|
-> {
|
|
::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
|
|
model.connection.exec_query("SELECT 1")
|
|
end
|
|
},
|
|
false, [:primary]
|
|
],
|
|
|
|
# use_replicas_for_read_queries ignores use_primary! for read queries
|
|
[
|
|
-> {
|
|
::Gitlab::Database::LoadBalancing::Session.current.use_primary!
|
|
::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
|
|
model.where(name: 'test1').to_a
|
|
end
|
|
},
|
|
false, [:replica]
|
|
],
|
|
|
|
# use_replicas_for_read_queries adheres use_primary! for write queries
|
|
[
|
|
-> {
|
|
::Gitlab::Database::LoadBalancing::Session.current.use_primary!
|
|
::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
|
|
model.create!(name: 'test1')
|
|
end
|
|
},
|
|
false, [:primary]
|
|
],
|
|
|
|
# use_replicas_for_read_queries adheres use_primary! for ambiguous queries
|
|
[
|
|
-> {
|
|
::Gitlab::Database::LoadBalancing::Session.current.use_primary!
|
|
::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
|
|
model.connection.exec_query('SELECT 1')
|
|
end
|
|
},
|
|
false, [:primary]
|
|
],
|
|
|
|
# use_replicas_for_read_queries ignores use_primary blocks
|
|
[
|
|
-> {
|
|
::Gitlab::Database::LoadBalancing::Session.current.use_primary do
|
|
::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
|
|
model.where(name: 'test1').to_a
|
|
end
|
|
end
|
|
},
|
|
false, [:replica]
|
|
],
|
|
|
|
# use_replicas_for_read_queries ignores a session already performed write
|
|
[
|
|
-> {
|
|
::Gitlab::Database::LoadBalancing::Session.current.write!
|
|
::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
|
|
model.where(name: 'test1').to_a
|
|
end
|
|
},
|
|
false, [:replica]
|
|
],
|
|
|
|
# fallback_to_replicas_for_ambiguous_queries
|
|
[
|
|
-> {
|
|
::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
|
|
model.first
|
|
model.where(name: 'test1').to_a
|
|
end
|
|
},
|
|
false, [:replica, :replica]
|
|
],
|
|
|
|
# fallback_to_replicas_for_ambiguous_queries for read-only transaction
|
|
[
|
|
-> {
|
|
::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
|
|
model.transaction do
|
|
model.first
|
|
model.where(name: 'test1').to_a
|
|
end
|
|
end
|
|
},
|
|
false, [:replica, :replica]
|
|
],
|
|
|
|
# A custom read query inside fallback_to_replicas_for_ambiguous_queries
|
|
[
|
|
-> {
|
|
::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
|
|
model.connection.exec_query("SELECT 1")
|
|
end
|
|
},
|
|
false, [:replica]
|
|
],
|
|
|
|
# A custom read query inside a transaction fallback_to_replicas_for_ambiguous_queries
|
|
[
|
|
-> {
|
|
::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
|
|
model.transaction do
|
|
model.connection.exec_query("SET LOCAL statement_timeout = 5000")
|
|
model.count
|
|
end
|
|
end
|
|
},
|
|
true, [:replica, :replica, :replica, :replica]
|
|
],
|
|
|
|
# fallback_to_replicas_for_ambiguous_queries after a write
|
|
[
|
|
-> {
|
|
model.create!(name: 'Test1')
|
|
::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
|
|
model.connection.exec_query("SELECT 1")
|
|
end
|
|
},
|
|
false, [:primary, :primary]
|
|
],
|
|
|
|
# fallback_to_replicas_for_ambiguous_queries after use_primary!
|
|
[
|
|
-> {
|
|
::Gitlab::Database::LoadBalancing::Session.current.use_primary!
|
|
::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
|
|
model.connection.exec_query("SELECT 1")
|
|
end
|
|
},
|
|
false, [:primary]
|
|
],
|
|
|
|
# fallback_to_replicas_for_ambiguous_queries inside use_primary
|
|
[
|
|
-> {
|
|
::Gitlab::Database::LoadBalancing::Session.current.use_primary do
|
|
::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
|
|
model.connection.exec_query("SELECT 1")
|
|
end
|
|
end
|
|
},
|
|
false, [:primary]
|
|
],
|
|
|
|
# use_primary inside fallback_to_replicas_for_ambiguous_queries
|
|
[
|
|
-> {
|
|
::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
|
|
::Gitlab::Database::LoadBalancing::Session.current.use_primary do
|
|
model.connection.exec_query("SELECT 1")
|
|
end
|
|
end
|
|
},
|
|
false, [:primary]
|
|
],
|
|
|
|
# A write query inside fallback_to_replicas_for_ambiguous_queries
|
|
[
|
|
-> {
|
|
::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
|
|
model.connection.exec_query("SELECT 1")
|
|
model.delete_all
|
|
model.connection.exec_query("SELECT 1")
|
|
end
|
|
},
|
|
false, [:replica, :primary, :primary]
|
|
],
|
|
|
|
# use_replicas_for_read_queries incorporates with fallback_to_replicas_for_ambiguous_queries
|
|
[
|
|
-> {
|
|
::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
|
|
::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
|
|
model.connection.exec_query('SELECT 1')
|
|
model.where(name: 'test1').to_a
|
|
end
|
|
end
|
|
},
|
|
false, [:replica, :replica]
|
|
]
|
|
]
|
|
end
|
|
|
|
with_them do
|
|
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[_test_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
|