107 lines
3.8 KiB
Ruby
107 lines
3.8 KiB
Ruby
# frozen_string_literal: true
|
|
|
|
module Gitlab
|
|
module Database
|
|
module LoadBalancing
|
|
# Module used for handling sticking connections to a primary, if
|
|
# necessary.
|
|
class Sticking
|
|
# The number of seconds after which a session should stop reading from
|
|
# the primary.
|
|
EXPIRATION = 30
|
|
|
|
def initialize(load_balancer)
|
|
@load_balancer = load_balancer
|
|
end
|
|
|
|
# Returns true if any caught up replica is found. This does not mean all replicas are caught up but the found
|
|
# caught up replica will be stored in the SafeRequestStore available as LoadBalancer#host for future queries.
|
|
# With use_primary_on_empty_location: true we will assume you need the primary if we can't find a matching
|
|
# location for the namespace, id pair. You should only use use_primary_on_empty_location in rare cases because
|
|
# we unstick once we find all replicas are caught up one time so it can be wasteful on the primary.
|
|
def find_caught_up_replica(namespace, id, use_primary_on_failure: true, use_primary_on_empty_location: false)
|
|
location = last_write_location_for(namespace, id)
|
|
|
|
result = if location
|
|
up_to_date_result = @load_balancer.select_up_to_date_host(location)
|
|
|
|
unstick(namespace, id) if up_to_date_result == LoadBalancer::ALL_CAUGHT_UP
|
|
|
|
up_to_date_result != LoadBalancer::NONE_CAUGHT_UP
|
|
else
|
|
# Some callers want to err on the side of caution and be really sure that a caught up replica was
|
|
# found. If we did not have any location to check then we must force `use_primary!` if they they
|
|
# use_primary_on_empty_location
|
|
!use_primary_on_empty_location
|
|
end
|
|
|
|
::Gitlab::Database::LoadBalancing::Session.current.use_primary! if !result && use_primary_on_failure
|
|
|
|
result
|
|
end
|
|
|
|
# Starts sticking to the primary for the given namespace and id, using
|
|
# the latest WAL pointer from the primary.
|
|
def stick(namespace, id)
|
|
with_primary_write_location do |location|
|
|
set_write_location_for(namespace, id, location)
|
|
end
|
|
::Gitlab::Database::LoadBalancing::Session.current.use_primary!
|
|
end
|
|
|
|
def bulk_stick(namespace, ids)
|
|
with_primary_write_location do |location|
|
|
ids.each do |id|
|
|
set_write_location_for(namespace, id, location)
|
|
end
|
|
end
|
|
|
|
::Gitlab::Database::LoadBalancing::Session.current.use_primary!
|
|
end
|
|
|
|
private
|
|
|
|
def with_primary_write_location
|
|
# When only using the primary, there's no point in getting write
|
|
# locations, as the primary is always in sync with itself.
|
|
return if @load_balancer.primary_only?
|
|
|
|
location = @load_balancer.primary_write_location
|
|
|
|
return if location.blank?
|
|
|
|
yield(location)
|
|
end
|
|
|
|
def unstick(namespace, id)
|
|
with_redis do |redis|
|
|
redis.del(redis_key_for(namespace, id))
|
|
end
|
|
end
|
|
|
|
def set_write_location_for(namespace, id, location)
|
|
with_redis do |redis|
|
|
redis.set(redis_key_for(namespace, id), location, ex: EXPIRATION)
|
|
end
|
|
end
|
|
|
|
def last_write_location_for(namespace, id)
|
|
with_redis do |redis|
|
|
redis.get(redis_key_for(namespace, id))
|
|
end
|
|
end
|
|
|
|
def redis_key_for(namespace, id)
|
|
name = @load_balancer.name
|
|
|
|
"database-load-balancing/write-location/#{name}/#{namespace}/#{id}"
|
|
end
|
|
|
|
def with_redis(&block)
|
|
Gitlab::Redis::DbLoadBalancing.with(&block)
|
|
end
|
|
end
|
|
end
|
|
end
|
|
end
|