debian-mirror-gitlab/lib/gitlab/database/load_balancing/connection_proxy.rb

144 lines
4.3 KiB
Ruby
Raw Normal View History

2021-09-04 01:27:46 +05:30
# frozen_string_literal: true
# rubocop:disable GitlabSecurity/PublicSend
module Gitlab
module Database
module LoadBalancing
# Redirecting of ActiveRecord connections.
#
# The ConnectionProxy class redirects ActiveRecord connection requests to
# the right load balancer pool, depending on the type of query.
class ConnectionProxy
WriteInsideReadOnlyTransactionError = Class.new(StandardError)
READ_ONLY_TRANSACTION_KEY = :load_balacing_read_only_transaction
attr_reader :load_balancer
# These methods perform writes after which we need to stick to the
# primary.
STICKY_WRITES = %i(
delete
delete_all
insert
update
update_all
).freeze
NON_STICKY_READS = %i(
sanitize_limit
select
select_one
select_rows
quote_column_name
).freeze
# hosts - The hosts to use for load balancing.
2021-11-11 11:23:49 +05:30
def initialize(load_balancer)
@load_balancer = load_balancer
2021-09-04 01:27:46 +05:30
end
def select_all(arel, name = nil, binds = [], preparable: nil)
if arel.respond_to?(:locked) && arel.locked
# SELECT ... FOR UPDATE queries should be sent to the primary.
2021-11-11 11:23:49 +05:30
current_session.write!
write_using_load_balancer(:select_all, arel, name, binds)
2021-09-04 01:27:46 +05:30
else
2021-10-27 15:23:28 +05:30
read_using_load_balancer(:select_all, arel, name, binds)
2021-09-04 01:27:46 +05:30
end
end
NON_STICKY_READS.each do |name|
2021-10-27 15:23:28 +05:30
define_method(name) do |*args, **kwargs, &block|
read_using_load_balancer(name, *args, **kwargs, &block)
2021-09-04 01:27:46 +05:30
end
end
STICKY_WRITES.each do |name|
2021-10-27 15:23:28 +05:30
define_method(name) do |*args, **kwargs, &block|
2021-11-11 11:23:49 +05:30
current_session.write!
write_using_load_balancer(name, *args, **kwargs, &block)
2021-09-04 01:27:46 +05:30
end
end
2021-10-27 15:23:28 +05:30
def transaction(*args, **kwargs, &block)
2021-09-04 01:27:46 +05:30
if current_session.fallback_to_replicas_for_ambiguous_queries?
track_read_only_transaction!
2021-10-27 15:23:28 +05:30
read_using_load_balancer(:transaction, *args, **kwargs, &block)
2021-09-04 01:27:46 +05:30
else
2021-11-11 11:23:49 +05:30
current_session.write!
write_using_load_balancer(:transaction, *args, **kwargs, &block)
2021-09-04 01:27:46 +05:30
end
ensure
untrack_read_only_transaction!
end
2021-11-11 11:23:49 +05:30
def respond_to_missing?(name, include_private = false)
@load_balancer.read_write do |connection|
connection.respond_to?(name, include_private)
end
end
2021-09-04 01:27:46 +05:30
# Delegates all unknown messages to a read-write connection.
2021-10-27 15:23:28 +05:30
def method_missing(...)
2021-09-04 01:27:46 +05:30
if current_session.fallback_to_replicas_for_ambiguous_queries?
2021-10-27 15:23:28 +05:30
read_using_load_balancer(...)
2021-09-04 01:27:46 +05:30
else
2021-10-27 15:23:28 +05:30
write_using_load_balancer(...)
2021-09-04 01:27:46 +05:30
end
end
# Performs a read using the load balancer.
#
# name - The name of the method to call on a connection object.
2021-10-27 15:23:28 +05:30
def read_using_load_balancer(...)
2021-09-04 01:27:46 +05:30
if current_session.use_primary? &&
!current_session.use_replicas_for_read_queries?
@load_balancer.read_write do |connection|
2021-10-27 15:23:28 +05:30
connection.send(...)
2021-09-04 01:27:46 +05:30
end
else
@load_balancer.read do |connection|
2021-10-27 15:23:28 +05:30
connection.send(...)
2021-09-04 01:27:46 +05:30
end
end
end
# Performs a write using the load balancer.
#
# name - The name of the method to call on a connection object.
# sticky - If set to true the session will stick to the master after
# the write.
2021-11-11 11:23:49 +05:30
def write_using_load_balancer(...)
2021-09-04 01:27:46 +05:30
if read_only_transaction?
raise WriteInsideReadOnlyTransactionError, 'A write query is performed inside a read-only transaction'
end
@load_balancer.read_write do |connection|
2021-11-11 11:23:49 +05:30
connection.send(...)
2021-09-04 01:27:46 +05:30
end
end
private
def current_session
::Gitlab::Database::LoadBalancing::Session.current
end
def track_read_only_transaction!
Thread.current[READ_ONLY_TRANSACTION_KEY] = true
end
def untrack_read_only_transaction!
Thread.current[READ_ONLY_TRANSACTION_KEY] = nil
end
def read_only_transaction?
Thread.current[READ_ONLY_TRANSACTION_KEY] == true
end
end
end
end
end