debian-mirror-gitlab/spec/lib/gitlab/database/load_balancing_spec.rb
2021-12-11 22:18:48 +05:30

542 lines
17 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 '.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