Commit 45bd7c92 authored by Steve Abrams's avatar Steve Abrams

Merge branch 'qmnguyen0711/implement-use-replicas-for-all-reads' into 'master'

Implement DB `use_replicas_for_read_queries`

See merge request gitlab-org/gitlab!59167
parents 586ca7a5 00dba8ce
---
title: Implement use_replicas_for_read_queries
merge_request: 59167
author:
type: added
......@@ -60,7 +60,7 @@ module Gitlab
end
def transaction(*args, &block)
if ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries?
if current_session.fallback_to_replicas_for_ambiguous_queries?
track_read_only_transaction!
read_using_load_balancer(:transaction, args, &block)
else
......@@ -73,7 +73,7 @@ module Gitlab
# Delegates all unknown messages to a read-write connection.
def method_missing(name, *args, &block)
if ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries?
if current_session.fallback_to_replicas_for_ambiguous_queries?
read_using_load_balancer(name, args, &block)
else
write_using_load_balancer(name, args, &block)
......@@ -84,10 +84,15 @@ module Gitlab
#
# name - The name of the method to call on a connection object.
def read_using_load_balancer(name, args, &block)
method = ::Gitlab::Database::LoadBalancing::Session.current.use_primary? ? :read_write : :read
@load_balancer.send(method) do |connection|
connection.send(name, *args, &block)
if current_session.use_primary? &&
!current_session.use_replicas_for_read_queries?
@load_balancer.read_write do |connection|
connection.send(name, *args, &block)
end
else
@load_balancer.read do |connection|
connection.send(name, *args, &block)
end
end
end
......@@ -105,7 +110,7 @@ module Gitlab
# Sticking has to be enabled before calling the method. Not doing so
# could lead to methods called in a block still being performed on a
# secondary instead of on a primary (when necessary).
::Gitlab::Database::LoadBalancing::Session.current.write! if sticky
current_session.write! if sticky
connection.send(name, *args, &block)
end
......@@ -115,6 +120,10 @@ module Gitlab
private
def current_session
::Gitlab::Database::LoadBalancing::Session.current
end
def track_read_only_transaction!
Thread.current[READ_ONLY_TRANSACTION_KEY] = true
end
......
......@@ -27,6 +27,8 @@ module Gitlab
@use_primary = false
@performed_write = false
@ignore_writes = false
@fallback_to_replicas_for_ambiguous_queries = false
@use_replicas_for_read_queries = false
end
def use_primary?
......@@ -55,6 +57,27 @@ module Gitlab
@ignore_writes = false
end
# Indicates that the read SQL statements from anywhere inside this
# blocks should use a replica, regardless of the current primary
# stickiness or whether a write query is already performed in the
# current session. This interface is reserved mostly for performance
# purpose. This is a good tool to push expensive queries, which can
# tolerate the replica lags, to the replicas.
#
# Write and ambiguous queries inside this block are still handled by
# the primary.
def use_replicas_for_read_queries(&blk)
previous_flag = @use_replicas_for_read_queries
@use_replicas_for_read_queries = true
yield
ensure
@use_replicas_for_read_queries = previous_flag
end
def use_replicas_for_read_queries?
@use_replicas_for_read_queries == true
end
# Indicate that the ambiguous SQL statements from anywhere inside this
# block should use a replica. The ambiguous statements include:
# - Transactions.
......
......@@ -153,6 +153,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
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)
......@@ -236,9 +237,34 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
.and_return(session)
end
describe 'with a regular session' do
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)
......@@ -250,6 +276,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
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')
......
......@@ -139,6 +139,78 @@ RSpec.describe Gitlab::Database::LoadBalancing::Session do
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 }
......
......@@ -561,6 +561,92 @@ RSpec.describe Gitlab::Database::LoadBalancing do
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
[
-> {
......@@ -613,7 +699,7 @@ RSpec.describe Gitlab::Database::LoadBalancing do
-> {
model.create!(name: 'Test1')
::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
model.first
model.connection.exec_query("SELECT 1")
end
},
false, [:primary, :primary]
......@@ -624,7 +710,7 @@ RSpec.describe Gitlab::Database::LoadBalancing do
-> {
::Gitlab::Database::LoadBalancing::Session.current.use_primary!
::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
model.first
model.connection.exec_query("SELECT 1")
end
},
false, [:primary]
......@@ -635,7 +721,7 @@ RSpec.describe Gitlab::Database::LoadBalancing do
-> {
::Gitlab::Database::LoadBalancing::Session.current.use_primary do
::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
model.first
model.connection.exec_query("SELECT 1")
end
end
},
......@@ -647,7 +733,7 @@ RSpec.describe Gitlab::Database::LoadBalancing do
-> {
::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
::Gitlab::Database::LoadBalancing::Session.current.use_primary do
model.first
model.connection.exec_query("SELECT 1")
end
end
},
......@@ -658,12 +744,25 @@ RSpec.describe Gitlab::Database::LoadBalancing do
[
-> {
::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
model.first
model.connection.exec_query("SELECT 1")
model.delete_all
model.where(name: 'test1').to_a
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
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment