diff options
Diffstat (limited to 'lib/gitlab/database/load_balancing/sticking.rb')
-rw-r--r-- | lib/gitlab/database/load_balancing/sticking.rb | 147 |
1 files changed, 147 insertions, 0 deletions
diff --git a/lib/gitlab/database/load_balancing/sticking.rb b/lib/gitlab/database/load_balancing/sticking.rb new file mode 100644 index 00000000000..efbd7099300 --- /dev/null +++ b/lib/gitlab/database/load_balancing/sticking.rb @@ -0,0 +1,147 @@ +# frozen_string_literal: true + +module Gitlab + module Database + module LoadBalancing + # Module used for handling sticking connections to a primary, if + # necessary. + # + # ## Examples + # + # Sticking a user to the primary: + # + # Sticking.stick_if_necessary(:user, current_user.id) + # + # To unstick if possible, or continue using the primary otherwise: + # + # Sticking.unstick_or_continue_sticking(:user, current_user.id) + module Sticking + # The number of seconds after which a session should stop reading from + # the primary. + EXPIRATION = 30 + + # Sticks to the primary if a write was performed. + def self.stick_if_necessary(namespace, id) + return unless LoadBalancing.enable? + + stick(namespace, id) if Session.current.performed_write? + end + + # Checks if we are caught-up with all the work + def self.all_caught_up?(namespace, id) + location = last_write_location_for(namespace, id) + + return true unless location + + load_balancer.all_caught_up?(location).tap do |caught_up| + unstick(namespace, id) if caught_up + end + end + + # Selects hosts that have caught up with the primary. This ensures + # atomic selection of the host to prevent the host list changing + # in another thread. + # + # Returns true if one host was selected. + def self.select_caught_up_replicas(namespace, id) + location = last_write_location_for(namespace, id) + + # Unlike all_caught_up?, we return false if no write location exists. + # We want to be sure we talk to a replica that has caught up for a specific + # write location. If no such location exists, err on the side of caution. + return false unless location + + load_balancer.select_caught_up_hosts(location).tap do |selected| + unstick(namespace, id) if selected + end + end + + # Sticks to the primary if necessary, otherwise unsticks an object (if + # it was previously stuck to the primary). + def self.unstick_or_continue_sticking(namespace, id) + Session.current.use_primary! unless all_caught_up?(namespace, id) + end + + # Select a replica that has caught up with the primary. If one has not been + # found, stick to the primary. + def self.select_valid_host(namespace, id) + replica_selected = select_caught_up_replicas(namespace, id) + + Session.current.use_primary! unless replica_selected + end + + # Starts sticking to the primary for the given namespace and id, using + # the latest WAL pointer from the primary. + def self.stick(namespace, id) + return unless LoadBalancing.enable? + + mark_primary_write_location(namespace, id) + Session.current.use_primary! + end + + def self.bulk_stick(namespace, ids) + return unless LoadBalancing.enable? + + with_primary_write_location do |location| + ids.each do |id| + set_write_location_for(namespace, id, location) + end + end + + Session.current.use_primary! + end + + def self.with_primary_write_location + return unless LoadBalancing.configured? + + # Load balancing could be enabled for the Web application server, + # but it's not activated for Sidekiq. We should update Redis with + # the write location just in case load balancing is being used. + location = + if LoadBalancing.enable? + load_balancer.primary_write_location + else + Gitlab::Database.get_write_location(ActiveRecord::Base.connection) + end + + return if location.blank? + + yield(location) + end + + def self.mark_primary_write_location(namespace, id) + with_primary_write_location do |location| + set_write_location_for(namespace, id, location) + end + end + + # Stops sticking to the primary. + def self.unstick(namespace, id) + Gitlab::Redis::SharedState.with do |redis| + redis.del(redis_key_for(namespace, id)) + end + end + + def self.set_write_location_for(namespace, id, location) + Gitlab::Redis::SharedState.with do |redis| + redis.set(redis_key_for(namespace, id), location, ex: EXPIRATION) + end + end + + def self.last_write_location_for(namespace, id) + Gitlab::Redis::SharedState.with do |redis| + redis.get(redis_key_for(namespace, id)) + end + end + + def self.redis_key_for(namespace, id) + "database-load-balancing/write-location/#{namespace}/#{id}" + end + + def self.load_balancer + LoadBalancing.proxy.load_balancer + end + end + end + end +end |