99 lines
2.8 KiB
Ruby
99 lines
2.8 KiB
Ruby
# 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
|