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.

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

lib/api/runner.rb

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

lib/ci/api/builds.rb

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

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

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