Merge branch 'fix-load-balancer-sticking' into 'master'

Fix sticking of the database load balancer

See merge request !1502

app/models/ci/runner.rb

 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

app/models/ee/build.rb

@@ -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

app/models/ee/ci/runner.rb

+module EE
+  module Ci
+    module Runner
+      def tick_runner_queue
+        ::Gitlab::Database::LoadBalancing::Sticking.stick(:runner, token)
+
+        super
+      end
+    end
+  end
+end

app/services/ee/user_project_access_changed_service.rb

+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

app/services/user_project_access_changed_service.rb

 class UserProjectAccessChangedService
+  prepend EE::UserProjectAccessChangedService
+
   def initialize(user_ids)
     @user_ids = Array.wrap(user_ids)
   end

lib/api/helpers.rb

 module API
   module Helpers
+    prepend EE::API::Helpers
+
     include Gitlab::Utils
     include Helpers::Pagination
 

lib/api/helpers/runner.rb

 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

lib/ci/api/helpers.rb

 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

lib/ee/api/helpers.rb

+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

lib/ee/api/helpers/runner.rb

+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

lib/ee/ci/api/helpers.rb

+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

lib/gitlab/database/load_balancing.rb

@@ -50,7 +50,7 @@ module Gitlab
       end
 
       def self.program_name
-        File.basename($0)
+        @program_name ||= File.basename($0)
       end
 
       # Configures proxying of requests.

lib/gitlab/database/load_balancing/connection_proxy.rb

@@ -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

lib/gitlab/database/load_balancing/rack_middleware.rb

 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

lib/gitlab/database/load_balancing/session.rb

@@ -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

lib/gitlab/database/load_balancing/sticking.rb

+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

spec/lib/ee/api/helpers/runner_spec.rb

+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

spec/lib/ee/api/helpers_spec.rb

+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

spec/lib/ee/ci/api/helpers_spec.rb

+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

spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb

@@ -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

spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb

@@ -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

spec/lib/gitlab/database/load_balancing/session_spec.rb

@@ -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

spec/lib/gitlab/database/load_balancing/sticking_spec.rb

+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

spec/models/ee/ci/build_spec.rb

@@ -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

spec/models/ee/ci/runner_spec.rb

+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

spec/services/ee/user_project_access_changed_service_spec.rb

+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