Commit 424c2771 authored by Yorick Peterse's avatar Yorick Peterse

Fix sticking of the database load balancer

There were 3 problems with sticking:

1. When sticking to a request (based on a previous request) we would
   always overwrite the WAL pointer, leading to a user being stuck to
   the primary for too long.

2. Sticking was not working for the Grape API. The API in particular is
   tricky because we have to stick either using a user, CI runner, or a
   CI build.

3. Refreshing of permissions did not lead to the refreshed users being
   stuck to the primary. This could result in users not being able to
   access new resources for a brief moment of time, or them still being
   able to see old resources.

To solve this some of the Grape related logic is handled by injecting EE
specific modules in the right places. This ensures we can stick as early
as possible, using the right data. The same technique is applied for
various parts of the CI codebase.
parent c4b21866
module Ci
class Runner < ActiveRecord::Base
extend Ci::Model
prepend EE::Ci::Runner
RUNNER_QUEUE_EXPIRY_TIME = 60.minutes
LAST_CONTACT_TIME = 1.hour.ago
......
......@@ -6,8 +6,19 @@ module EE
module Build
extend ActiveSupport::Concern
included do
after_save :stick_build_if_status_changed
end
def shared_runners_minutes_limit_enabled?
runner && runner.shared? && project.shared_runners_minutes_limit_enabled?
end
def stick_build_if_status_changed
return unless status_changed?
return unless running?
::Gitlab::Database::LoadBalancing::Sticking.stick(:build, id)
end
end
end
module EE
module Ci
module Runner
def tick_runner_queue
::Gitlab::Database::LoadBalancing::Sticking.stick(:runner, token)
super
end
end
end
end
module EE
module UserProjectAccessChangedService
def execute
result = super
@user_ids.each do |id|
::Gitlab::Database::LoadBalancing::Sticking.stick(:user, id)
end
result
end
end
end
class UserProjectAccessChangedService
prepend EE::UserProjectAccessChangedService
def initialize(user_ids)
@user_ids = Array.wrap(user_ids)
end
......
module API
module Helpers
prepend EE::API::Helpers
include Gitlab::Utils
include Helpers::Pagination
......
module API
module Helpers
module Runner
prepend EE::API::Helpers::Runner
JOB_TOKEN_HEADER = 'HTTP_JOB_TOKEN'.freeze
JOB_TOKEN_PARAM = :token
UPDATE_RUNNER_EVERY = 10 * 60
......@@ -50,10 +52,14 @@ module API
forbidden!('Job has been erased!') if job.erased?
end
def authenticate_job!(job)
def authenticate_job!
job = Ci::Build.find_by_id(params[:id])
validate_job!(job) do
forbidden! unless job_token_valid?(job)
end
job
end
def job_token_valid?(job)
......
......@@ -113,8 +113,7 @@ module API
optional :state, type: String, desc: %q(Job's status: success, failed)
end
put '/:id' do
job = Ci::Build.find_by_id(params[:id])
authenticate_job!(job)
job = authenticate_job!
job.update_attributes(trace: params[:trace]) if params[:trace]
......@@ -140,8 +139,7 @@ module API
optional :token, type: String, desc: %q(Job's authentication token)
end
patch '/:id/trace' do
job = Ci::Build.find_by_id(params[:id])
authenticate_job!(job)
job = authenticate_job!
error!('400 Missing header Content-Range', 400) unless request.headers.has_key?('Content-Range')
content_range = request.headers['Content-Range']
......@@ -175,8 +173,7 @@ module API
require_gitlab_workhorse!
Gitlab::Workhorse.verify_api_request!(headers)
job = Ci::Build.find_by_id(params[:id])
authenticate_job!(job)
job = authenticate_job!
forbidden!('Job is not running') unless job.running?
if params[:filesize]
......@@ -212,8 +209,7 @@ module API
not_allowed! unless Gitlab.config.artifacts.enabled
require_gitlab_workhorse!
job = Ci::Build.find_by_id(params[:id])
authenticate_job!(job)
job = authenticate_job!
forbidden!('Job is not running!') unless job.running?
artifacts_upload_path = ArtifactUploader.artifacts_upload_path
......@@ -245,8 +241,7 @@ module API
optional :token, type: String, desc: %q(Job's authentication token)
end
get '/:id/artifacts' do
job = Ci::Build.find_by_id(params[:id])
authenticate_job!(job)
job = authenticate_job!
artifacts_file = job.artifacts_file
unless artifacts_file.file_storage?
......
......@@ -86,8 +86,7 @@ module Ci
# Example Request:
# PATCH /builds/:id/trace.txt
patch ":id/trace.txt" do
build = Ci::Build.find_by_id(params[:id])
authenticate_build!(build)
build = authenticate_build!
error!('400 Missing header Content-Range', 400) unless request.headers.has_key?('Content-Range')
content_range = request.headers['Content-Range']
......@@ -117,8 +116,7 @@ module Ci
require_gitlab_workhorse!
Gitlab::Workhorse.verify_api_request!(headers)
not_allowed! unless Gitlab.config.artifacts.enabled
build = Ci::Build.find_by_id(params[:id])
authenticate_build!(build)
build = authenticate_build!
forbidden!('build is not running') unless build.running?
if params[:filesize]
......@@ -154,8 +152,7 @@ module Ci
post ":id/artifacts" do
require_gitlab_workhorse!
not_allowed! unless Gitlab.config.artifacts.enabled
build = Ci::Build.find_by_id(params[:id])
authenticate_build!(build)
build = authenticate_build!
forbidden!('Build is not running!') unless build.running?
artifacts_upload_path = ArtifactUploader.artifacts_upload_path
......@@ -189,8 +186,7 @@ module Ci
# Example Request:
# GET /builds/:id/artifacts
get ":id/artifacts" do
build = Ci::Build.find_by_id(params[:id])
authenticate_build!(build)
build = authenticate_build!
artifacts_file = build.artifacts_file
unless artifacts_file.file_storage?
......@@ -214,8 +210,7 @@ module Ci
# Example Request:
# DELETE /builds/:id/artifacts
delete ":id/artifacts" do
build = Ci::Build.find_by_id(params[:id])
authenticate_build!(build)
build = authenticate_build!
status(200)
build.erase_artifacts!
......
module Ci
module API
module Helpers
prepend EE::Ci::API::Helpers
BUILD_TOKEN_HEADER = "HTTP_BUILD_TOKEN".freeze
BUILD_TOKEN_PARAM = :token
UPDATE_RUNNER_EVERY = 10 * 60
......@@ -13,10 +15,14 @@ module Ci
forbidden! unless current_runner
end
def authenticate_build!(build)
def authenticate_build!
build = Ci::Build.find_by_id(params[:id])
validate_build!(build) do
forbidden! unless build_token_valid?(build)
end
build
end
def validate_build!(build)
......
module EE
module API
module Helpers
def current_user
user = super
::Gitlab::Database::LoadBalancing::RackMiddleware.
stick_or_unstick(env, :user, user.id) if user
user
end
end
end
end
module EE
module API
module Helpers
module Runner
def authenticate_job!
id = params[:id]
::Gitlab::Database::LoadBalancing::RackMiddleware.
stick_or_unstick(env, :build, id) if id
super
end
def current_runner
token = params[:token]
::Gitlab::Database::LoadBalancing::RackMiddleware.
stick_or_unstick(env, :runner, token) if token
super
end
end
end
end
end
module EE
module Ci
module API
module Helpers
def authenticate_build!
id = params[:id]
::Gitlab::Database::LoadBalancing::RackMiddleware.
stick_or_unstick(env, :build, id) if id
super
end
def current_runner
token = params[:token]
::Gitlab::Database::LoadBalancing::RackMiddleware.
stick_or_unstick(env, :runner, token) if token
super
end
end
end
end
end
......@@ -50,7 +50,7 @@ module Gitlab
end
def self.program_name
File.basename($0)
@program_name ||= File.basename($0)
end
# Configures proxying of requests.
......
......@@ -70,7 +70,7 @@ module Gitlab
# Sticking has to be enabled before calling the method. Not doing so
# could lead to methods called in a block still being performed on a
# secondary instead of on a primary (when necessary).
Session.current.use_primary! if sticky
Session.current.write! if sticky
connection.send(name, *args, &block)
end
......
module Gitlab
module Database
module LoadBalancing
# Rack middleware for managing load balancing.
# 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
SESSION_KEY = :gitlab_load_balancer
# The number of seconds after which a session should stop reading from
# the primary.
EXPIRATION = 30
STICK_OBJECT = 'load_balancing.stick_object'.freeze
# 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] = [namespace, id]
end
def initialize(app)
@app = app
......@@ -18,34 +32,36 @@ module Gitlab
# doesn't linger around.
clear
user = user_for_request(env)
check_primary_requirement(user) if user
unstick_or_continue_sticking(env)
result = @app.call(env)
assign_primary_for_user(user) if Session.current.use_primary? && user
stick_if_necessary(env)
result
ensure
clear
end
# Checks if we need to use the primary for the current user.
def check_primary_requirement(user)
location = last_write_location_for(user)
# 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)
namespace, id = sticking_namespace_and_id(env)
return unless location
if load_balancer.all_caught_up?(location)
delete_write_location_for(user)
else
Session.current.use_primary!
if namespace && id
Sticking.unstick_or_continue_sticking(namespace, id)
end
end
def assign_primary_for_user(user)
set_write_location_for(user, load_balancer.primary_write_location)
# Determine if we need to stick after handling a request.
def stick_if_necessary(env)
namespace, id = sticking_namespace_and_id(env)
if namespace && id
Sticking.stick_if_necessary(namespace, id)
end
end
def clear
......@@ -57,44 +73,22 @@ module Gitlab
LoadBalancing.proxy.load_balancer
end
# Returns the User object for the currently authenticated user, if any.
def user_for_request(env)
api = env['api.endpoint']
# 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_namespace_and_id(env)
warden = env['warden']
if api && api.respond_to?(:current_user)
# The current request is an API request. In this case we can use our
# `current_user` helper method.
api.current_user
elsif warden && warden.user
# Used by the Rails app, and sometimes by the API.
warden.user
if warden && warden.user
[:user, warden.user.id]
elsif env[STICK_OBJECT]
env[STICK_OBJECT]
else
nil
end
end
def last_write_location_for(user)
Gitlab::Redis.with do |redis|
redis.get(redis_key_for(user))
[]
end
end
def delete_write_location_for(user)
Gitlab::Redis.with do |redis|
redis.del(redis_key_for(user))
end
end
def set_write_location_for(user, location)
Gitlab::Redis.with do |redis|
redis.set(redis_key_for(user), location, ex: EXPIRATION)
end
end
def redis_key_for(user)
"database-load-balancing/write-location/#{user.id}"
end
end
end
end
......
......@@ -19,6 +19,7 @@ module Gitlab
def initialize
@use_primary = false
@performed_write = false
end
def use_primary?
......@@ -28,6 +29,15 @@ module Gitlab
def use_primary!
@use_primary = true
end
def write!
@performed_write = true
use_primary!
end
def performed_write?
@performed_write
end
end
end
end
......
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
# 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)
location = last_write_location_for(namespace, id)
return unless location
if load_balancer.all_caught_up?(location)
unstick(namespace, id)
else
Session.current.use_primary!
end
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?
location = load_balancer.primary_write_location
set_write_location_for(namespace, id, location)
Session.current.use_primary!
end
# Stops sticking to the primary.
def self.unstick(namespace, id)
Gitlab::Redis.with do |redis|
redis.del(redis_key_for(namespace, id))
end
end
def self.set_write_location_for(namespace, id, location)
Gitlab::Redis.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.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
require 'spec_helper'
describe EE::API::Helpers::Runner do
let(:helper) { Class.new { include API::Helpers::Runner }.new }
before do
allow(helper).to receive(:env).and_return({})
end
describe '#authenticate_job' do
let(:build) { create(:ci_build) }
before do
allow(helper).to receive(:validate_job!)
end
it 'handles sticking of a build when a build ID is specified' do
allow(helper).to receive(:params).and_return(id: build.id)
expect(Gitlab::Database::LoadBalancing::RackMiddleware).
to receive(:stick_or_unstick).
with({}, :build, build.id)
helper.authenticate_job!
end
it 'does not handle sticking if no build ID was specified' do
allow(helper).to receive(:params).and_return({})
expect(Gitlab::Database::LoadBalancing::RackMiddleware).
not_to receive(:stick_or_unstick)
helper.authenticate_job!
end
it 'returns the build if one could be found' do
allow(helper).to receive(:params).and_return(id: build.id)
expect(helper.authenticate_job!).to eq(build)
end
end
describe '#current_runner' do
let(:runner) { create(:ci_runner, token: 'foo') }
it 'handles sticking of a runner if a token is specified' do
allow(helper).to receive(:params).and_return(token: runner.token)
expect(Gitlab::Database::LoadBalancing::RackMiddleware).
to receive(:stick_or_unstick).
with({}, :runner, runner.token)
helper.current_runner
end
it 'does not handle sticking if no token was specified' do
allow(helper).to receive(:params).and_return({})
expect(Gitlab::Database::LoadBalancing::RackMiddleware).
not_to receive(:stick_or_unstick)
helper.current_runner
end
it 'returns the runner if one could be found' do
allow(helper).to receive(:params).and_return(token: runner.token)
expect(helper.current_runner).to eq(runner)
end
end
end
require 'spec_helper'
describe EE::API::Helpers do
let(:helper) { Class.new { include API::Helpers }.new }
before do
allow(helper).to receive(:env).and_return({})
allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
end
describe '#current_user' do
let(:user) { build(:user, id: 42) }
before do
allow(helper).to receive(:sudo!)
end
it 'handles sticking when a user could be found' do
allow(helper).to receive(:initial_current_user).and_return(user)
expect(Gitlab::Database::LoadBalancing::RackMiddleware).
to receive(:stick_or_unstick).with({}, :user, 42)
helper.current_user
end
it 'does not handle sticking if no user could be found' do
allow(helper).to receive(:initial_current_user).and_return(nil)
expect(Gitlab::Database::LoadBalancing::RackMiddleware).
not_to receive(:stick_or_unstick)
helper.current_user
end
it 'returns the user if one could be found' do
allow(helper).to receive(:initial_current_user).and_return(user)
expect(helper.current_user).to eq(user)
end
end
end
require 'spec_helper'
describe EE::Ci::API::Helpers do
let(:helper) { Class.new { include Ci::API::Helpers }.new }
before do
allow(helper).to receive(:env).and_return({})
end
describe '#authenticate_build' do
let(:build) { create(:ci_build) }
before do
allow(helper).to receive(:validate_build!)
end
it 'handles sticking of a build when a build ID is specified' do
allow(helper).to receive(:params).and_return(id: build.id)
expect(Gitlab::Database::LoadBalancing::RackMiddleware).
to receive(:stick_or_unstick).
with({}, :build, build.id)
helper.authenticate_build!
end
it 'does not handle sticking if no build ID was specified' do
allow(helper).to receive(:params).and_return({})
expect(Gitlab::Database::LoadBalancing::RackMiddleware).
not_to receive(:stick_or_unstick)
helper.authenticate_build!
end
it 'returns the build if one could be found' do
allow(helper).to receive(:params).and_return(id: build.id)
expect(helper.authenticate_build!).to eq(build)
end
end
describe '#current_runner' do
let(:runner) { create(:ci_runner, token: 'foo') }
it 'handles sticking of a runner if a token is specified' do
allow(helper).to receive(:params).and_return(token: runner.token)
expect(Gitlab::Database::LoadBalancing::RackMiddleware).
to receive(:stick_or_unstick).
with({}, :runner, runner.token)
helper.current_runner
end
it 'does not handle sticking if no token was specified' do
allow(helper).to receive(:params).and_return({})
expect(Gitlab::Database::LoadBalancing::RackMiddleware).
not_to receive(:stick_or_unstick)
helper.current_runner
end
it 'returns the runner if one could be found' do
allow(helper).to receive(:params).and_return(token: runner.token)
expect(helper.current_runner).to eq(runner)
end
end
end
......@@ -124,10 +124,10 @@ describe Gitlab::Database::LoadBalancing::ConnectionProxy do
and_return(session)
end
it 'uses the primary' do
it 'it uses but does not stick to the primary when sticking is disabled' do
expect(proxy.load_balancer).to receive(:read_write).and_yield(connection)
expect(connection).to receive(:foo).with('foo')
expect(session).not_to receive(:use_primary!)
expect(session).not_to receive(:write!)
proxy.write_using_load_balancer(:foo, 'foo')
end
......@@ -135,7 +135,7 @@ describe Gitlab::Database::LoadBalancing::ConnectionProxy do
it 'sticks to the primary when sticking is enabled' do
expect(proxy.load_balancer).to receive(:read_write).and_yield(connection)
expect(connection).to receive(:foo).with('foo')
expect(session).to receive(:use_primary!)
expect(session).to receive(:write!)
proxy.write_using_load_balancer(:foo, 'foo', sticky: true)
end
......
......@@ -8,134 +8,100 @@ describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
Gitlab::Database::LoadBalancing::Session.clear_session
end
describe '#call' do
let(:lb) { double(:lb) }
let(:user) { double(:user, id: 42) }
before do
expect(app).to receive(:call).with(an_instance_of(Hash))
allow(middleware).to receive(:load_balancer).and_return(lb)
expect(middleware).to receive(:clear).twice
end
describe '.stick_or_unstick' do
it 'sticks or unsticks and updates the Rack environment' do
allow(Gitlab::Database::LoadBalancing).to receive(:enable?).
and_return(true)
context 'when the primary was used' do
it 'assigns the user to the primary' do
allow(middleware).to receive(:user_for_request).and_return(user)
expect(Gitlab::Database::LoadBalancing::Sticking).
to receive(:unstick_or_continue_sticking).
with(:user, 42)
allow(middleware).to receive(:last_write_location_for).
with(user).
and_return('123')
env = {}
allow(lb).to receive(:all_caught_up?).with('123').and_return(false)
described_class.stick_or_unstick(env, :user, 42)
expect(middleware).to receive(:assign_primary_for_user).with(user)
middleware.call({})
end
end
context 'when a primary was not used' do
it 'does not assign the user to the primary' do
allow(middleware).to receive(:user_for_request).and_return(user)
allow(middleware).to receive(:last_write_location_for).
with(user).
and_return('123')
allow(lb).to receive(:all_caught_up?).with('123').and_return(true)
expect(middleware).not_to receive(:assign_primary_for_user)
middleware.call({})
end
expect(env[described_class::STICK_OBJECT]).to eq([:user, 42])
end
end
describe '#check_primary_requirement' do
let(:lb) { double(:lb) }
let(:user) { double(:user, id: 42) }
before do
allow(middleware).to receive(:load_balancer).and_return(lb)
end
describe '#call' do
it 'handles a request' do
env = {}
it 'marks the primary as the host to use when necessary' do
expect(middleware).to receive(:last_write_location_for).
with(user).
and_return('foo')
expect(middleware).to receive(:clear).twice
expect(lb).to receive(:all_caught_up?).with('foo').and_return(false)
expect(middleware).to receive(:unstick_or_continue_sticking).with(env)
expect(middleware).to receive(:stick_if_necessary).with(env)
expect(Gitlab::Database::LoadBalancing::Session.current).
to receive(:use_primary!)
expect(app).to receive(:call).with(env).and_return(10)
middleware.check_primary_requirement(user)
expect(middleware.call(env)).to eq(10)
end
end
it 'does not use the primary when there is no cached write location' do
expect(middleware).to receive(:last_write_location_for).
with(user).
and_return(nil)
expect(lb).not_to receive(:all_caught_up?)
expect(Gitlab::Database::LoadBalancing::Session.current).
not_to receive(:use_primary!)
describe '#unstick_or_continue_sticking' do
it 'does not stick if no namespace and identifier could be found' do
expect(Gitlab::Database::LoadBalancing::Sticking).
not_to receive(:unstick_or_continue_sticking)
middleware.check_primary_requirement(user)
middleware.unstick_or_continue_sticking({})
end
it 'does not use the primary when all hosts have caught up' do
expect(middleware).to receive(:last_write_location_for).
with(user).
and_return('foo')
it 'sticks to the primary if a sticking namespace and identifier were found' do
env = { described_class::STICK_OBJECT => [:user, 42] }
expect(lb).to receive(:all_caught_up?).with('foo').and_return(true)
expect(Gitlab::Database::LoadBalancing::Sticking).
to receive(:unstick_or_continue_sticking).
with(:user, 42)
expect(middleware).to receive(:delete_write_location_for).with(user)
middleware.check_primary_requirement(user)
middleware.unstick_or_continue_sticking(env)
end
end
describe '#assign_primary_for_user' do
it 'stores primary instance details for the current user' do
user = double(:user, id: 42)
describe '#stick_if_necessary' do
it 'does not stick to the primary if not necessary' do
expect(Gitlab::Database::LoadBalancing::Sticking).
not_to receive(:stick_if_necessary)
lb = double(:load_balancer, primary_write_location: '123')
middleware.stick_if_necessary({})
end
allow(middleware).to receive(:load_balancer).and_return(lb)
it 'sticks to the primary if necessary' do
env = { described_class::STICK_OBJECT => [:user, 42] }
expect(middleware).to receive(:set_write_location_for).with(user, '123')
expect(Gitlab::Database::LoadBalancing::Sticking).
to receive(:stick_if_necessary).
with(:user, 42)
middleware.assign_primary_for_user(user)
middleware.stick_if_necessary(env)
end
end
describe '#clear' do
it 'clears the currently used host and session' do
proxy = double(:proxy)
lb = double(:lb)
session = double(:session)
allow(middleware).to receive(:load_balancer).and_return(lb)
allow(Gitlab::Database::LoadBalancing).to receive(:proxy).and_return(proxy)
allow(proxy).to receive(:load_balancer).and_return(lb)
expect(lb).to receive(:release_host)
middleware.clear
stub_const('Gitlab::Database::LoadBalancing::RackMiddleware::Session',
session)
thread_key = Gitlab::Database::LoadBalancing::Session::CACHE_KEY
expect(session).to receive(:clear_session)
expect(RequestStore[thread_key]).to be_nil
middleware.clear
end
end
describe '#load_balancer' do
it 'returns the load balancer' do
describe '.load_balancer' do
it 'returns a the load balancer' do
proxy = double(:proxy)
allow(Gitlab::Database::LoadBalancing).to receive(:proxy).and_return(proxy)
expect(Gitlab::Database::LoadBalancing).to receive(:proxy).
and_return(proxy)
expect(proxy).to receive(:load_balancer)
......@@ -143,63 +109,36 @@ describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
end
end
describe '#user_for_request' do
let(:user) { double(:user, id: 42) }
it 'returns the current user for a Grape request' do
env = { 'api.endpoint' => double(:api, current_user: user) }
expect(middleware.user_for_request(env)).to eq(user)
end
it 'returns the current user for a Rails request' do
env = { 'warden' => double(:warden, user: user) }
expect(middleware.user_for_request(env)).to eq(user)
end
it 'returns nil if no user could be found' do
expect(middleware.user_for_request({})).to be_nil
end
end
describe '#sticking_namespace_and_id' do
context 'using a Warden request' do
it 'returns the warden user if present' do
user = double(:user, id: 42)
warden = double(:warden, user: user)
env = { 'warden' => warden }
describe '#last_write_location_for' do
it 'returns the last WAL write location for a user' do
user = double(:user, id: 42)
expect(middleware.sticking_namespace_and_id(env)).to eq([:user, 42])
end
middleware.set_write_location_for(user, '123')
it 'returns an empty Array if no user was present' do
warden = double(:warden, user: nil)
env = { 'warden' => warden }
expect(middleware.last_write_location_for(user)).to eq('123')
expect(middleware.sticking_namespace_and_id(env)).to eq([])
end
end
end
describe '#delete_write_location' do
it 'removes the WAL write location from Redis' do
user = double(:user, id: 42)
middleware.set_write_location_for(user, '123')
middleware.delete_write_location_for(user)
expect(middleware.last_write_location_for(user)).to be_nil
end
end
context 'using a request with a manually set sticking object' do
it 'returns the sticking object' do
env = { described_class::STICK_OBJECT => [:user, 42] }
describe '#set_write_location' do
it 'stores the WAL write location in Redis' do
user = double(:user, id: 42)
middleware.set_write_location_for(user, '123')
expect(middleware.last_write_location_for(user)).to eq('123')
expect(middleware.sticking_namespace_and_id(env)).to eq([:user, 42])
end
end
end
describe '#redis_key_for' do
it 'returns a String' do
user = double(:user, id: 42)
expect(middleware.redis_key_for(user)).
to eq('database-load-balancing/write-location/42')
context 'using a regular request' do
it 'returns an empty Array' do
expect(middleware.sticking_namespace_and_id({})).to eq([])
end
end
end
end
......@@ -32,5 +32,23 @@ describe Gitlab::Database::LoadBalancing::Session do
it 'returns false when a secondary should be used' do
expect(described_class.new.use_primary?).to eq(false)
end
it 'returns true when a write was performed' do
instance = described_class.new
instance.write!
expect(instance.use_primary?).to eq(true)
end
end
describe '#performed_write?' do
it 'returns true if a write was performed' do
instance = described_class.new
instance.write!
expect(instance.performed_write?).to eq(true)
end
end
end
require 'spec_helper'
describe Gitlab::Database::LoadBalancing::Sticking, :redis do
after do
Gitlab::Database::LoadBalancing::Session.clear_session
end
describe '.stick_if_necessary' do
context 'when sticking is disabled' do
it 'does not perform any sticking' do
expect(described_class).not_to receive(:stick)
described_class.stick_if_necessary(:user, 42)
end
end
context 'when sticking is enabled' do
before do
allow(Gitlab::Database::LoadBalancing).to receive(:enable?).
and_return(true)
end
it 'does not stick if no write was performed' do
allow(Gitlab::Database::LoadBalancing::Session.current).
to receive(:performed_write?).
and_return(false)
expect(described_class).not_to receive(:stick)
described_class.stick_if_necessary(:user, 42)
end
it 'sticks to the primary if a write was performed' do
allow(Gitlab::Database::LoadBalancing::Session.current).
to receive(:performed_write?).
and_return(true)
expect(described_class).to receive(:stick).with(:user, 42)
described_class.stick_if_necessary(:user, 42)
end
end
end
describe '.unstick_or_continue_sticking' do
let(:lb) { double(:lb) }
before do
allow(described_class).to receive(:load_balancer).and_return(lb)
end
it 'simply returns if no write location could be found' do
allow(described_class).to receive(:last_write_location_for).
with(:user, 42).
and_return(nil)
expect(lb).not_to receive(:all_caught_up?)
described_class.unstick_or_continue_sticking(:user, 42)
end
it 'unsticks if all secondaries have caught up' do
allow(described_class).to receive(:last_write_location_for).
with(:user, 42).
and_return('foo')
allow(lb).to receive(:all_caught_up?).with('foo').and_return(true)
expect(described_class).to receive(:unstick).with(:user, 42)
described_class.unstick_or_continue_sticking(:user, 42)
end
it 'continues using the primary if the secondaries have not yet caught up' do
allow(described_class).to receive(:last_write_location_for).
with(:user, 42).
and_return('foo')
allow(lb).to receive(:all_caught_up?).with('foo').and_return(false)
expect(Gitlab::Database::LoadBalancing::Session.current).
to receive(:use_primary!)
described_class.unstick_or_continue_sticking(:user, 42)
end
end
describe '.stick' do
context 'when sticking is disabled' do
it 'does not perform any sticking' do
expect(described_class).not_to receive(:set_write_location_for)
described_class.stick(:user, 42)
end
end
context 'when sticking is enabled' do
it 'sticks an entity to the primary' do
allow(Gitlab::Database::LoadBalancing).to receive(:enable?).
and_return(true)
lb = double(:lb, primary_write_location: 'foo')
allow(described_class).to receive(:load_balancer).and_return(lb)
expect(described_class).to receive(:set_write_location_for).
with(:user, 42, 'foo')
expect(Gitlab::Database::LoadBalancing::Session.current).
to receive(:use_primary!)
described_class.stick(:user, 42)
end
end
end
describe '.unstick' do
it 'removes the sticking data from Redis' do
described_class.set_write_location_for(:user, 4, 'foo')
described_class.unstick(:user, 4)
expect(described_class.last_write_location_for(:user, 4)).to be_nil
end
end
describe '.last_write_location_for' do
it 'returns the last WAL write location for a user' do
described_class.set_write_location_for(:user, 4, 'foo')
expect(described_class.last_write_location_for(:user, 4)).to eq('foo')
end
end
describe '.redis_key_for' do
it 'returns a String' do
expect(described_class.redis_key_for(:user, 42)).
to eq('database-load-balancing/write-location/user/42')
end
end
describe '.load_balancer' do
it 'returns a the load balancer' do
proxy = double(:proxy)
expect(Gitlab::Database::LoadBalancing).to receive(:proxy).
and_return(proxy)
expect(proxy).to receive(:load_balancer)
described_class.load_balancer
end
end
end
......@@ -53,4 +53,18 @@ describe Ci::Build, models: true do
end
end
end
describe '#stick_build_if_status_changed' do
it 'sticks the build if the status changed' do
build = create(:ci_build, :pending)
allow(Gitlab::Database::LoadBalancing).to receive(:enable?).
and_return(true)
expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:stick).
with(:build, build.id)
build.update(status: :running)
end
end
end
require 'spec_helper'
describe EE::Ci::Runner, models: true do
describe '#tick_runner_queue' do
it 'sticks the runner to the primary and calls the original method' do
runner = create(:ci_runner)
allow(Gitlab::Database::LoadBalancing).to receive(:enable?).
and_return(true)
expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:stick).
with(:runner, runner.token)
expect(Gitlab::Workhorse).to receive(:set_key_and_notify)
runner.tick_runner_queue
end
end
end
require 'spec_helper'
describe EE::UserProjectAccessChangedService do
let(:service) { UserProjectAccessChangedService.new([1, 2]) }
describe '#execute' do
it 'sticks all the updated users and returns the original result' do
allow(Gitlab::Database::LoadBalancing).to receive(:enable?).
and_return(true)
expect(AuthorizedProjectsWorker).to receive(:bulk_perform_and_wait).
with([[1], [2]]).
and_return(10)
[1, 2].each do |id|
expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:stick).
with(:user, id).
ordered
end
expect(service.execute).to eq(10)
end
end
end
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment