1 files changed, 63 insertions, 55 deletions

diff --git a/lib/gitlab/database/load_balancing/sticking.rb b/lib/gitlab/database/load_balancing/sticking.rb
index 20d42b9a694..df4ad18581f 100644
--- a/lib/gitlab/database/load_balancing/sticking.rb
+++ b/lib/gitlab/database/load_balancing/sticking.rb
@@ -5,36 +5,47 @@ module Gitlab
     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
+      class 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?
+        def initialize(load_balancer)
+          @load_balancer = load_balancer
+          @model = load_balancer.configuration.model
+        end
 
-          stick(namespace, id) if Session.current.performed_write?
+        # Unsticks or continues sticking the current request.
+        #
+        # This method also updates the Rack environment so #call can later
+        # determine if we still need to stick or not.
+        #
+        # env - The Rack environment.
+        # namespace - The namespace to use for sticking.
+        # id - The identifier to use for sticking.
+        # model - The ActiveRecord model to scope sticking to.
+        def stick_or_unstick_request(env, namespace, id)
+          unstick_or_continue_sticking(namespace, id)
+
+          env[RackMiddleware::STICK_OBJECT] ||= Set.new
+          env[RackMiddleware::STICK_OBJECT] << [@model, namespace, id]
+        end
+
+        # Sticks to the primary if a write was performed.
+        def stick_if_necessary(namespace, id)
+          stick(namespace, id) if ::Gitlab::Database::LoadBalancing::Session.current.performed_write?
         end
 
-        # Checks if we are caught-up with all the work
-        def self.all_caught_up?(namespace, id)
+        def all_caught_up?(namespace, id)
           location = last_write_location_for(namespace, id)
 
           return true unless location
 
-          load_balancer.select_up_to_date_host(location).tap do |found|
-            ActiveSupport::Notifications.instrument('caught_up_replica_pick.load_balancing', { result: found } )
+          @load_balancer.select_up_to_date_host(location).tap do |found|
+            ActiveSupport::Notifications.instrument(
+              'caught_up_replica_pick.load_balancing',
+              { result: found }
+            )
 
             unstick(namespace, id) if found
           end
@@ -45,7 +56,7 @@ module Gitlab
         # in another thread.
         #
         # Returns true if one host was selected.
-        def self.select_caught_up_replicas(namespace, id)
+        def 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.
@@ -53,95 +64,92 @@ module Gitlab
           # write location. If no such location exists, err on the side of caution.
           return false unless location
 
-          load_balancer.select_up_to_date_host(location).tap do |selected|
+          @load_balancer.select_up_to_date_host(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)
+        def unstick_or_continue_sticking(namespace, id)
+          return if all_caught_up?(namespace, id)
+
+          ::Gitlab::Database::LoadBalancing::Session.current.use_primary!
         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)
+        def select_valid_host(namespace, id)
+          replica_selected =
+            select_caught_up_replicas(namespace, id)
 
-          Session.current.use_primary! unless replica_selected
+          ::Gitlab::Database::LoadBalancing::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?
-
+        def stick(namespace, id)
           mark_primary_write_location(namespace, id)
-          Session.current.use_primary!
+          ::Gitlab::Database::LoadBalancing::Session.current.use_primary!
         end
 
-        def self.bulk_stick(namespace, ids)
-          return unless LoadBalancing.enable?
-
+        def bulk_stick(namespace, ids)
           with_primary_write_location do |location|
             ids.each do |id|
               set_write_location_for(namespace, id, location)
             end
           end
 
-          Session.current.use_primary!
+          ::Gitlab::Database::LoadBalancing::Session.current.use_primary!
         end
 
-        def self.with_primary_write_location
-          return unless LoadBalancing.configured?
+        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?
 
-          # 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.main.get_write_location(ActiveRecord::Base.connection)
-            end
+          location = @load_balancer.primary_write_location
 
           return if location.blank?
 
           yield(location)
         end
 
-        def self.mark_primary_write_location(namespace, id)
+        def 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)
+        def unstick(namespace, id)
           Gitlab::Redis::SharedState.with do |redis|
             redis.del(redis_key_for(namespace, id))
+            redis.del(old_redis_key_for(namespace, id))
           end
         end
 
-        def self.set_write_location_for(namespace, id, location)
+        def set_write_location_for(namespace, id, location)
           Gitlab::Redis::SharedState.with do |redis|
             redis.set(redis_key_for(namespace, id), location, ex: EXPIRATION)
+            redis.set(old_redis_key_for(namespace, id), location, ex: EXPIRATION)
           end
         end
 
-        def self.last_write_location_for(namespace, id)
+        def last_write_location_for(namespace, id)
           Gitlab::Redis::SharedState.with do |redis|
-            redis.get(redis_key_for(namespace, id))
+            redis.get(redis_key_for(namespace, id)) ||
+              redis.get(old_redis_key_for(namespace, id))
           end
         end
 
-        def self.redis_key_for(namespace, id)
-          "database-load-balancing/write-location/#{namespace}/#{id}"
+        def redis_key_for(namespace, id)
+          name = @load_balancer.name
+
+          "database-load-balancing/write-location/#{name}/#{namespace}/#{id}"
         end
 
-        def self.load_balancer
-          LoadBalancing.proxy.load_balancer
+        def old_redis_key_for(namespace, id)
+          "database-load-balancing/write-location/#{namespace}/#{id}"
         end
       end
     end