1 files changed, 98 insertions, 0 deletions
diff --git a/lib/gitlab/database/load_balancing/rack_middleware.rb b/lib/gitlab/database/load_balancing/rack_middleware.rb
new file mode 100644
index 00000000000..4734ff99bd3
--- /dev/null
+++ b/lib/gitlab/database/load_balancing/rack_middleware.rb
@@ -0,0 +1,98 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module LoadBalancing
+      # Rack middleware to handle sticking when serving Rails requests. Grape
+      # API calls are handled separately as different API endpoints need to
+      # stick based on different objects.
+      class RackMiddleware
+        STICK_OBJECT = 'load_balancing.stick_object'
+
+        # 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.
+        def self.stick_or_unstick(env, namespace, id)
+          return unless LoadBalancing.enable?
+
+          Sticking.unstick_or_continue_sticking(namespace, id)
+
+          env[STICK_OBJECT] ||= Set.new
+          env[STICK_OBJECT] << [namespace, id]
+        end
+
+        def initialize(app)
+          @app = app
+        end
+
+        def call(env)
+          # Ensure that any state that may have run before the first request
+          # doesn't linger around.
+          clear
+
+          unstick_or_continue_sticking(env)
+
+          result = @app.call(env)
+
+          stick_if_necessary(env)
+
+          result
+        ensure
+          clear
+        end
+
+        # Determine if we need to stick based on currently available user data.
+        #
+        # Typically this code will only be reachable for Rails requests as
+        # Grape data is not yet available at this point.
+        def unstick_or_continue_sticking(env)
+          namespaces_and_ids = sticking_namespaces_and_ids(env)
+
+          namespaces_and_ids.each do |namespace, id|
+            Sticking.unstick_or_continue_sticking(namespace, id)
+          end
+        end
+
+        # Determine if we need to stick after handling a request.
+        def stick_if_necessary(env)
+          namespaces_and_ids = sticking_namespaces_and_ids(env)
+
+          namespaces_and_ids.each do |namespace, id|
+            Sticking.stick_if_necessary(namespace, id)
+          end
+        end
+
+        def clear
+          load_balancer.release_host
+          Session.clear_session
+        end
+
+        def load_balancer
+          LoadBalancing.proxy.load_balancer
+        end
+
+        # Determines the sticking namespace and identifier based on the Rack
+        # environment.
+        #
+        # For Rails requests this uses warden, but Grape and others have to
+        # manually set the right environment variable.
+        def sticking_namespaces_and_ids(env)
+          warden = env['warden']
+
+          if warden && warden.user
+            [[:user, warden.user.id]]
+          elsif env[STICK_OBJECT].present?
+            env[STICK_OBJECT].to_a
+          else
+            []
+          end
+        end
+      end
+    end
+  end
+end