Add latest changes from gitlab-org/gitlab@14-4-stable-eev14.4.0-rc42

11 files changed, 461 insertions, 361 deletions

diff --git a/spec/lib/gitlab/database/load_balancing/action_cable_callbacks_spec.rb b/spec/lib/gitlab/database/load_balancing/action_cable_callbacks_spec.rb
index ebbbafb855f..768855464c1 100644
--- a/spec/lib/gitlab/database/load_balancing/action_cable_callbacks_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/action_cable_callbacks_spec.rb
@@ -5,7 +5,7 @@ require 'spec_helper'
 RSpec.describe Gitlab::Database::LoadBalancing::ActionCableCallbacks, :request_store do
   describe '.wrapper' do
     it 'uses primary and then releases the connection and clears the session' do
-      expect(Gitlab::Database::LoadBalancing).to receive_message_chain(:proxy, :load_balancer, :release_host)
+      expect(Gitlab::Database::LoadBalancing).to receive(:release_hosts)
       expect(Gitlab::Database::LoadBalancing::Session).to receive(:clear_session)
 
       described_class.wrapper.call(
@@ -18,7 +18,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ActionCableCallbacks, :request_s
 
     context 'with an exception' do
       it 'releases the connection and clears the session' do
-        expect(Gitlab::Database::LoadBalancing).to receive_message_chain(:proxy, :load_balancer, :release_host)
+        expect(Gitlab::Database::LoadBalancing).to receive(:release_hosts)
         expect(Gitlab::Database::LoadBalancing::Session).to receive(:clear_session)
 
         expect do
diff --git a/spec/lib/gitlab/database/load_balancing/active_record_proxy_spec.rb b/spec/lib/gitlab/database/load_balancing/active_record_proxy_spec.rb
deleted file mode 100644
index 8886ce9756d..00000000000
--- a/spec/lib/gitlab/database/load_balancing/active_record_proxy_spec.rb
+++ /dev/null
@@ -1,20 +0,0 @@
-# frozen_string_literal: true
-
-require 'spec_helper'
-
-RSpec.describe Gitlab::Database::LoadBalancing::ActiveRecordProxy do
-  describe '#connection' do
-    it 'returns a connection proxy' do
-      dummy = Class.new do
-        include Gitlab::Database::LoadBalancing::ActiveRecordProxy
-      end
-
-      proxy = double(:proxy)
-
-      expect(Gitlab::Database::LoadBalancing).to receive(:proxy)
-        .and_return(proxy)
-
-      expect(dummy.new.connection).to eq(proxy)
-    end
-  end
-end
diff --git a/spec/lib/gitlab/database/load_balancing/configuration_spec.rb b/spec/lib/gitlab/database/load_balancing/configuration_spec.rb
index 6621e6276a5..3e5249a3dea 100644
--- a/spec/lib/gitlab/database/load_balancing/configuration_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/configuration_spec.rb
@@ -108,6 +108,14 @@ RSpec.describe Gitlab::Database::LoadBalancing::Configuration do
   end
 
   describe '#load_balancing_enabled?' do
+    it 'returns false when running inside a Rake task' do
+      config = described_class.new(ActiveRecord::Base, %w[foo bar])
+
+      allow(Gitlab::Runtime).to receive(:rake?).and_return(true)
+
+      expect(config.load_balancing_enabled?).to eq(false)
+    end
+
     it 'returns true when hosts are configured' do
       config = described_class.new(ActiveRecord::Base, %w[foo bar])
 
diff --git a/spec/lib/gitlab/database/load_balancing/host_spec.rb b/spec/lib/gitlab/database/load_balancing/host_spec.rb
index e2011692228..b040c7a76bd 100644
--- a/spec/lib/gitlab/database/load_balancing/host_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/host_spec.rb
@@ -172,6 +172,14 @@ RSpec.describe Gitlab::Database::LoadBalancing::Host do
 
         expect(host).not_to be_online
       end
+
+      it 'returns false when ActiveRecord::ConnectionNotEstablished is raised' do
+        allow(host)
+          .to receive(:check_replica_status?)
+          .and_raise(ActiveRecord::ConnectionNotEstablished)
+
+        expect(host).not_to be_online
+      end
     end
   end
 
diff --git a/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb b/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb
index 86fae14b961..f3ce5563e38 100644
--- a/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb
@@ -47,16 +47,27 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
   end
 
   describe '#initialize' do
-    it 'ignores the hosts when the primary_only option is enabled' do
+    it 'ignores the hosts when load balancing is disabled' do
       config = Gitlab::Database::LoadBalancing::Configuration
         .new(ActiveRecord::Base, [db_host])
-      lb = described_class.new(config, primary_only: true)
+
+      allow(config).to receive(:load_balancing_enabled?).and_return(false)
+
+      lb = described_class.new(config)
       hosts = lb.host_list.hosts
 
       expect(hosts.length).to eq(1)
       expect(hosts.first)
         .to be_instance_of(Gitlab::Database::LoadBalancing::PrimaryHost)
     end
+
+    it 'sets the name of the connection that is used' do
+      config =
+        Gitlab::Database::LoadBalancing::Configuration.new(ActiveRecord::Base)
+      lb = described_class.new(config)
+
+      expect(lb.name).to eq(:main)
+    end
   end
 
   describe '#read' do
@@ -140,10 +151,13 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
         .to yield_with_args(ActiveRecord::Base.retrieve_connection)
     end
 
-    it 'uses the primary when the primary_only option is enabled' do
+    it 'uses the primary when load balancing is disabled' do
       config = Gitlab::Database::LoadBalancing::Configuration
         .new(ActiveRecord::Base)
-      lb = described_class.new(config, primary_only: true)
+
+      allow(config).to receive(:load_balancing_enabled?).and_return(false)
+
+      lb = described_class.new(config)
 
       # When no hosts are configured, we don't want to produce any warnings, as
       # they aren't useful/too noisy.
@@ -274,34 +288,43 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
 
       expect { lb.retry_with_backoff { raise } }.to raise_error(RuntimeError)
     end
-  end
 
-  describe '#connection_error?' do
-    before do
-      stub_const('Gitlab::Database::LoadBalancing::LoadBalancer::CONNECTION_ERRORS',
-                 [NotImplementedError])
+    it 'skips retries when only the primary is used' do
+      allow(lb).to receive(:primary_only?).and_return(true)
+
+      expect(lb).not_to receive(:sleep)
+
+      expect { lb.retry_with_backoff { raise } }.to raise_error(RuntimeError)
     end
+  end
 
+  describe '#connection_error?' do
     it 'returns true for a connection error' do
-      error = NotImplementedError.new
+      error = ActiveRecord::ConnectionNotEstablished.new
 
       expect(lb.connection_error?(error)).to eq(true)
     end
 
+    it 'returns false for a missing database error' do
+      error = ActiveRecord::NoDatabaseError.new
+
+      expect(lb.connection_error?(error)).to eq(false)
+    end
+
     it 'returns true for a wrapped connection error' do
-      wrapped = wrapped_exception(ActiveRecord::StatementInvalid, NotImplementedError)
+      wrapped = wrapped_exception(ActiveRecord::StatementInvalid, ActiveRecord::ConnectionNotEstablished)
 
       expect(lb.connection_error?(wrapped)).to eq(true)
     end
 
     it 'returns true for a wrapped connection error from a view' do
-      wrapped = wrapped_exception(ActionView::Template::Error, NotImplementedError)
+      wrapped = wrapped_exception(ActionView::Template::Error, ActiveRecord::ConnectionNotEstablished)
 
       expect(lb.connection_error?(wrapped)).to eq(true)
     end
 
     it 'returns true for deeply wrapped/nested errors' do
-      top = twice_wrapped_exception(ActionView::Template::Error, ActiveRecord::StatementInvalid, NotImplementedError)
+      top = twice_wrapped_exception(ActionView::Template::Error, ActiveRecord::StatementInvalid, ActiveRecord::ConnectionNotEstablished)
 
       expect(lb.connection_error?(top)).to eq(true)
     end
diff --git a/spec/lib/gitlab/database/load_balancing/primary_host_spec.rb b/spec/lib/gitlab/database/load_balancing/primary_host_spec.rb
index a0e63a7ee4e..45d81808971 100644
--- a/spec/lib/gitlab/database/load_balancing/primary_host_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/primary_host_spec.rb
@@ -63,9 +63,8 @@ RSpec.describe Gitlab::Database::LoadBalancing::PrimaryHost do
   end
 
   describe '#primary_write_location' do
-    it 'returns the write location of the primary' do
-      expect(host.primary_write_location).to be_an_instance_of(String)
-      expect(host.primary_write_location).not_to be_empty
+    it 'raises NotImplementedError' do
+      expect { host.primary_write_location }.to raise_error(NotImplementedError)
     end
   end
 
@@ -76,51 +75,8 @@ RSpec.describe Gitlab::Database::LoadBalancing::PrimaryHost do
   end
 
   describe '#database_replica_location' do
-    let(:connection) { double(:connection) }
-
-    it 'returns the write ahead location of the replica', :aggregate_failures do
-      expect(host)
-        .to receive(:query_and_release)
-        .and_return({ 'location' => '0/D525E3A8' })
-
-      expect(host.database_replica_location).to be_an_instance_of(String)
-    end
-
-    it 'returns nil when the database query returned no rows' do
-      expect(host).to receive(:query_and_release).and_return({})
-
-      expect(host.database_replica_location).to be_nil
-    end
-
-    it 'returns nil when the database connection fails' do
-      allow(host).to receive(:connection).and_raise(PG::Error)
-
-      expect(host.database_replica_location).to be_nil
-    end
-  end
-
-  describe '#query_and_release' do
-    it 'executes a SQL query' do
-      results = host.query_and_release('SELECT 10 AS number')
-
-      expect(results).to be_an_instance_of(Hash)
-      expect(results['number'].to_i).to eq(10)
-    end
-
-    it 'releases the connection after running the query' do
-      expect(host)
-        .to receive(:release_connection)
-        .once
-
-      host.query_and_release('SELECT 10 AS number')
-    end
-
-    it 'returns an empty Hash in the event of an error' do
-      expect(host.connection)
-        .to receive(:select_all)
-        .and_raise(RuntimeError, 'kittens')
-
-      expect(host.query_and_release('SELECT 10 AS number')).to eq({})
+    it 'raises NotImplementedError' do
+      expect { host.database_replica_location }.to raise_error(NotImplementedError)
     end
   end
 end
diff --git a/spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb b/spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb
index ea0c7f781fd..af7e2a4b167 100644
--- a/spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb
@@ -6,12 +6,12 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
   let(:app) { double(:app) }
   let(:middleware) { described_class.new(app) }
   let(:warden_user) { double(:warden, user: double(:user, id: 42)) }
-  let(:single_sticking_object) { Set.new([[:user, 42]]) }
+  let(:single_sticking_object) { Set.new([[ActiveRecord::Base, :user, 42]]) }
   let(:multiple_sticking_objects) do
     Set.new([
-      [:user, 42],
-      [:runner, '123456789'],
-      [:runner, '1234']
+      [ActiveRecord::Base, :user, 42],
+      [ActiveRecord::Base, :runner, '123456789'],
+      [ActiveRecord::Base, :runner, '1234']
     ])
   end
 
@@ -19,47 +19,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
     Gitlab::Database::LoadBalancing::Session.clear_session
   end
 
-  describe '.stick_or_unstick' do
-    before do
-      allow(Gitlab::Database::LoadBalancing).to receive(:enable?)
-        .and_return(true)
-    end
-
-    it 'sticks or unsticks a single object and updates the Rack environment' do
-      expect(Gitlab::Database::LoadBalancing::Sticking)
-        .to receive(:unstick_or_continue_sticking)
-        .with(:user, 42)
-
-      env = {}
-
-      described_class.stick_or_unstick(env, :user, 42)
-
-      expect(env[described_class::STICK_OBJECT].to_a).to eq([[:user, 42]])
-    end
-
-    it 'sticks or unsticks multiple objects and updates the Rack environment' do
-      expect(Gitlab::Database::LoadBalancing::Sticking)
-        .to receive(:unstick_or_continue_sticking)
-        .with(:user, 42)
-        .ordered
-
-      expect(Gitlab::Database::LoadBalancing::Sticking)
-        .to receive(:unstick_or_continue_sticking)
-        .with(:runner, '123456789')
-        .ordered
-
-      env = {}
-
-      described_class.stick_or_unstick(env, :user, 42)
-      described_class.stick_or_unstick(env, :runner, '123456789')
-
-      expect(env[described_class::STICK_OBJECT].to_a).to eq([
-        [:user, 42],
-        [:runner, '123456789']
-      ])
-    end
-  end
-
   describe '#call' do
     it 'handles a request' do
       env = {}
@@ -82,7 +41,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
 
   describe '#unstick_or_continue_sticking' do
     it 'does not stick if no namespace and identifier could be found' do
-      expect(Gitlab::Database::LoadBalancing::Sticking)
+      expect(ApplicationRecord.sticking)
         .not_to receive(:unstick_or_continue_sticking)
 
       middleware.unstick_or_continue_sticking({})
@@ -91,9 +50,11 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
     it 'sticks to the primary if a warden user is found' do
       env = { 'warden' => warden_user }
 
-      expect(Gitlab::Database::LoadBalancing::Sticking)
-        .to receive(:unstick_or_continue_sticking)
-        .with(:user, 42)
+      Gitlab::Database::LoadBalancing.base_models.each do |model|
+        expect(model.sticking)
+          .to receive(:unstick_or_continue_sticking)
+          .with(:user, 42)
+      end
 
       middleware.unstick_or_continue_sticking(env)
     end
@@ -101,7 +62,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
     it 'sticks to the primary if a sticking namespace and identifier is found' do
       env = { described_class::STICK_OBJECT => single_sticking_object }
 
-      expect(Gitlab::Database::LoadBalancing::Sticking)
+      expect(ApplicationRecord.sticking)
         .to receive(:unstick_or_continue_sticking)
         .with(:user, 42)
 
@@ -111,17 +72,17 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
     it 'sticks to the primary if multiple sticking namespaces and identifiers were found' do
       env = { described_class::STICK_OBJECT => multiple_sticking_objects }
 
-      expect(Gitlab::Database::LoadBalancing::Sticking)
+      expect(ApplicationRecord.sticking)
         .to receive(:unstick_or_continue_sticking)
         .with(:user, 42)
         .ordered
 
-      expect(Gitlab::Database::LoadBalancing::Sticking)
+      expect(ApplicationRecord.sticking)
         .to receive(:unstick_or_continue_sticking)
         .with(:runner, '123456789')
         .ordered
 
-      expect(Gitlab::Database::LoadBalancing::Sticking)
+      expect(ApplicationRecord.sticking)
         .to receive(:unstick_or_continue_sticking)
         .with(:runner, '1234')
         .ordered
@@ -132,7 +93,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
 
   describe '#stick_if_necessary' do
     it 'does not stick to the primary if not necessary' do
-      expect(Gitlab::Database::LoadBalancing::Sticking)
+      expect(ApplicationRecord.sticking)
         .not_to receive(:stick_if_necessary)
 
       middleware.stick_if_necessary({})
@@ -141,9 +102,11 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
     it 'sticks to the primary if a warden user is found' do
       env = { 'warden' => warden_user }
 
-      expect(Gitlab::Database::LoadBalancing::Sticking)
-        .to receive(:stick_if_necessary)
-        .with(:user, 42)
+      Gitlab::Database::LoadBalancing.base_models.each do |model|
+        expect(model.sticking)
+          .to receive(:stick_if_necessary)
+          .with(:user, 42)
+      end
 
       middleware.stick_if_necessary(env)
     end
@@ -151,7 +114,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
     it 'sticks to the primary if a a single sticking object is found' do
       env = { described_class::STICK_OBJECT => single_sticking_object }
 
-      expect(Gitlab::Database::LoadBalancing::Sticking)
+      expect(ApplicationRecord.sticking)
         .to receive(:stick_if_necessary)
         .with(:user, 42)
 
@@ -161,17 +124,17 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
     it 'sticks to the primary if multiple sticking namespaces and identifiers were found' do
       env = { described_class::STICK_OBJECT => multiple_sticking_objects }
 
-      expect(Gitlab::Database::LoadBalancing::Sticking)
+      expect(ApplicationRecord.sticking)
         .to receive(:stick_if_necessary)
         .with(:user, 42)
         .ordered
 
-      expect(Gitlab::Database::LoadBalancing::Sticking)
+      expect(ApplicationRecord.sticking)
         .to receive(:stick_if_necessary)
         .with(:runner, '123456789')
         .ordered
 
-      expect(Gitlab::Database::LoadBalancing::Sticking)
+      expect(ApplicationRecord.sticking)
         .to receive(:stick_if_necessary)
         .with(:runner, '1234')
         .ordered
@@ -182,47 +145,34 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
 
   describe '#clear' do
     it 'clears the currently used host and session' do
-      lb = double(:lb)
       session = spy(:session)
 
-      allow(middleware).to receive(:load_balancer).and_return(lb)
-
-      expect(lb).to receive(:release_host)
-
       stub_const('Gitlab::Database::LoadBalancing::Session', session)
 
+      expect(Gitlab::Database::LoadBalancing).to receive(:release_hosts)
+
       middleware.clear
 
       expect(session).to have_received(:clear_session)
     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)
-
-      middleware.load_balancer
-    end
-  end
-
-  describe '#sticking_namespaces_and_ids' do
+  describe '#sticking_namespaces' do
     context 'using a Warden request' do
       it 'returns the warden user if present' do
         env = { 'warden' => warden_user }
+        ids = Gitlab::Database::LoadBalancing.base_models.map do |model|
+          [model, :user, 42]
+        end
 
-        expect(middleware.sticking_namespaces_and_ids(env)).to eq([[:user, 42]])
+        expect(middleware.sticking_namespaces(env)).to eq(ids)
       end
 
       it 'returns an empty Array if no user was present' do
         warden = double(:warden, user: nil)
         env = { 'warden' => warden }
 
-        expect(middleware.sticking_namespaces_and_ids(env)).to eq([])
+        expect(middleware.sticking_namespaces(env)).to eq([])
       end
     end
 
@@ -230,17 +180,17 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
       it 'returns the sticking object' do
         env = { described_class::STICK_OBJECT => multiple_sticking_objects }
 
-        expect(middleware.sticking_namespaces_and_ids(env)).to eq([
-          [:user, 42],
-          [:runner, '123456789'],
-          [:runner, '1234']
+        expect(middleware.sticking_namespaces(env)).to eq([
+          [ActiveRecord::Base, :user, 42],
+          [ActiveRecord::Base, :runner, '123456789'],
+          [ActiveRecord::Base, :runner, '1234']
         ])
       end
     end
 
     context 'using a regular request' do
       it 'returns an empty Array' do
-        expect(middleware.sticking_namespaces_and_ids({})).to eq([])
+        expect(middleware.sticking_namespaces({})).to eq([])
       end
     end
   end
diff --git a/spec/lib/gitlab/database/load_balancing/setup_spec.rb b/spec/lib/gitlab/database/load_balancing/setup_spec.rb
new file mode 100644
index 00000000000..01646bc76ef
--- /dev/null
+++ b/spec/lib/gitlab/database/load_balancing/setup_spec.rb
@@ -0,0 +1,119 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::LoadBalancing::Setup do
+  describe '#setup' do
+    it 'sets up the load balancer' do
+      setup = described_class.new(ActiveRecord::Base)
+
+      expect(setup).to receive(:disable_prepared_statements)
+      expect(setup).to receive(:setup_load_balancer)
+      expect(setup).to receive(:setup_service_discovery)
+
+      setup.setup
+    end
+  end
+
+  describe '#disable_prepared_statements' do
+    it 'disables prepared statements and reconnects to the database' do
+      config = double(
+        :config,
+        configuration_hash: { host: 'localhost' },
+        env_name: 'test',
+        name: 'main'
+      )
+      model = double(:model, connection_db_config: config)
+
+      expect(ActiveRecord::DatabaseConfigurations::HashConfig)
+        .to receive(:new)
+        .with('test', 'main', { host: 'localhost', prepared_statements: false })
+        .and_call_original
+
+      # HashConfig doesn't implement its own #==, so we can't directly compare
+      # the expected value with a pre-defined one.
+      expect(model)
+        .to receive(:establish_connection)
+        .with(an_instance_of(ActiveRecord::DatabaseConfigurations::HashConfig))
+
+      described_class.new(model).disable_prepared_statements
+    end
+  end
+
+  describe '#setup_load_balancer' do
+    it 'sets up the load balancer' do
+      model = Class.new(ActiveRecord::Base)
+      setup = described_class.new(model)
+      config = Gitlab::Database::LoadBalancing::Configuration.new(model)
+      lb = instance_spy(Gitlab::Database::LoadBalancing::LoadBalancer)
+
+      allow(lb).to receive(:configuration).and_return(config)
+
+      expect(Gitlab::Database::LoadBalancing::LoadBalancer)
+        .to receive(:new)
+        .with(setup.configuration)
+        .and_return(lb)
+
+      setup.setup_load_balancer
+
+      expect(model.connection.load_balancer).to eq(lb)
+      expect(model.sticking)
+        .to be_an_instance_of(Gitlab::Database::LoadBalancing::Sticking)
+    end
+  end
+
+  describe '#setup_service_discovery' do
+    context 'when service discovery is disabled' do
+      it 'does nothing' do
+        expect(Gitlab::Database::LoadBalancing::ServiceDiscovery)
+          .not_to receive(:new)
+
+        described_class.new(ActiveRecord::Base).setup_service_discovery
+      end
+    end
+
+    context 'when service discovery is enabled' do
+      it 'immediately performs service discovery' do
+        model = ActiveRecord::Base
+        setup = described_class.new(model)
+        sv = instance_spy(Gitlab::Database::LoadBalancing::ServiceDiscovery)
+        lb = model.connection.load_balancer
+
+        allow(setup.configuration)
+          .to receive(:service_discovery_enabled?)
+          .and_return(true)
+
+        allow(Gitlab::Database::LoadBalancing::ServiceDiscovery)
+          .to receive(:new)
+          .with(lb, setup.configuration.service_discovery)
+          .and_return(sv)
+
+        expect(sv).to receive(:perform_service_discovery)
+        expect(sv).not_to receive(:start)
+
+        setup.setup_service_discovery
+      end
+
+      it 'starts service discovery if needed' do
+        model = ActiveRecord::Base
+        setup = described_class.new(model, start_service_discovery: true)
+        sv = instance_spy(Gitlab::Database::LoadBalancing::ServiceDiscovery)
+        lb = model.connection.load_balancer
+
+        allow(setup.configuration)
+          .to receive(:service_discovery_enabled?)
+          .and_return(true)
+
+        allow(Gitlab::Database::LoadBalancing::ServiceDiscovery)
+          .to receive(:new)
+          .with(lb, setup.configuration.service_discovery)
+          .and_return(sv)
+
+        expect(sv).to receive(:perform_service_discovery)
+        expect(sv).to receive(:start)
+
+        setup.setup_service_discovery
+      end
+    end
+  end
+end
diff --git a/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb b/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb
index f683ade978a..08dd6a0a788 100644
--- a/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb
@@ -5,14 +5,12 @@ require 'spec_helper'
 RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
   let(:middleware) { described_class.new }
 
-  let(:load_balancer) { double.as_null_object }
   let(:worker_class) { 'TestDataConsistencyWorker' }
   let(:job) { { "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e" } }
 
   before do
     skip_feature_flags_yaml_validation
     skip_default_enabled_yaml_check
-    allow(::Gitlab::Database::LoadBalancing).to receive_message_chain(:proxy, :load_balancer).and_return(load_balancer)
   end
 
   after do
@@ -23,7 +21,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
     middleware.call(worker_class, job, nil, nil) {}
   end
 
-  describe '#call' do
+  describe '#call', :database_replica do
     shared_context 'data consistency worker class' do |data_consistency, feature_flag|
       let(:expected_consistency) { data_consistency }
       let(:worker_class) do
@@ -85,9 +83,15 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
         end
 
         it 'passes database_replica_location' do
-          expected_location = { Gitlab::Database::MAIN_DATABASE_NAME.to_sym => location }
+          expected_location = {}
 
-          expect(load_balancer).to receive_message_chain(:host, "database_replica_location").and_return(location)
+          Gitlab::Database::LoadBalancing.each_load_balancer do |lb|
+            expect(lb.host)
+              .to receive(:database_replica_location)
+              .and_return(location)
+
+            expected_location[lb.name] = location
+          end
 
           run_middleware
 
@@ -103,9 +107,15 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
         end
 
         it 'passes primary write location', :aggregate_failures do
-          expected_location = { Gitlab::Database::MAIN_DATABASE_NAME.to_sym => location }
+          expected_location = {}
 
-          expect(load_balancer).to receive(:primary_write_location).and_return(location)
+          Gitlab::Database::LoadBalancing.each_load_balancer do |lb|
+            expect(lb)
+              .to receive(:primary_write_location)
+              .and_return(location)
+
+            expected_location[lb.name] = location
+          end
 
           run_middleware
 
@@ -137,8 +147,10 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
       let(:job) { { "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", 'wal_locations' => wal_locations } }
 
       before do
-        allow(load_balancer).to receive(:primary_write_location).and_return(new_location)
-        allow(load_balancer).to receive(:database_replica_location).and_return(new_location)
+        Gitlab::Database::LoadBalancing.each_load_balancer do |lb|
+          allow(lb).to receive(:primary_write_location).and_return(new_location)
+          allow(lb).to receive(:database_replica_location).and_return(new_location)
+        end
       end
 
       shared_examples_for 'does not set database location again' do |use_primary|
diff --git a/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb b/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb
index 9f23eb0094f..06efdcd8f99 100644
--- a/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb
@@ -2,20 +2,17 @@
 
 require 'spec_helper'
 
-RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware do
+RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware, :clean_gitlab_redis_queues do
   let(:middleware) { described_class.new }
-
-  let(:load_balancer) { double.as_null_object }
-
   let(:worker) { worker_class.new }
   let(:job) { { "retry" => 3, "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", 'database_replica_location' => '0/D525E3A8' } }
 
   before do
     skip_feature_flags_yaml_validation
     skip_default_enabled_yaml_check
-    allow(::Gitlab::Database::LoadBalancing).to receive_message_chain(:proxy, :load_balancer).and_return(load_balancer)
 
     replication_lag!(false)
+    Gitlab::Database::LoadBalancing::Session.clear_session
   end
 
   after do
@@ -67,7 +64,10 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware do
       let(:wal_locations) { { Gitlab::Database::MAIN_DATABASE_NAME.to_sym => location } }
 
       it 'does not stick to the primary', :aggregate_failures do
-        expect(load_balancer).to receive(:select_up_to_date_host).with(location).and_return(true)
+        expect(ActiveRecord::Base.connection.load_balancer)
+          .to receive(:select_up_to_date_host)
+          .with(location)
+          .and_return(true)
 
         run_middleware do
           expect(Gitlab::Database::LoadBalancing::Session.current.use_primary?).not_to be_truthy
@@ -92,7 +92,12 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware do
         let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e', 'wal_locations' => wal_locations } }
 
         before do
-          allow(load_balancer).to receive(:select_up_to_date_host).with(location).and_return(true)
+          Gitlab::Database::LoadBalancing.each_load_balancer do |lb|
+            allow(lb)
+              .to receive(:select_up_to_date_host)
+              .with(location)
+              .and_return(true)
+          end
         end
 
         it_behaves_like 'replica is up to date', 'replica'
@@ -102,7 +107,10 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware do
         let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e', 'dedup_wal_locations' => wal_locations } }
 
         before do
-          allow(load_balancer).to receive(:select_up_to_date_host).with(wal_locations[:main]).and_return(true)
+          allow(ActiveRecord::Base.connection.load_balancer)
+            .to receive(:select_up_to_date_host)
+            .with(wal_locations[:main])
+            .and_return(true)
         end
 
         it_behaves_like 'replica is up to date', 'replica'
@@ -112,7 +120,10 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware do
         let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e', 'database_write_location' => '0/D525E3A8' } }
 
         before do
-          allow(load_balancer).to receive(:select_up_to_date_host).with('0/D525E3A8').and_return(true)
+          allow(ActiveRecord::Base.connection.load_balancer)
+            .to receive(:select_up_to_date_host)
+            .with('0/D525E3A8')
+            .and_return(true)
         end
 
         it_behaves_like 'replica is up to date', 'replica'
@@ -158,18 +169,15 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware do
               process_job(job)
             end.to raise_error(Sidekiq::JobRetry::Skip)
 
-            expect(job['error_class']).to eq('Gitlab::Database::LoadBalancing::SidekiqServerMiddleware::JobReplicaNotUpToDate')
+            job_for_retry = Sidekiq::RetrySet.new.first
+            expect(job_for_retry['error_class']).to eq('Gitlab::Database::LoadBalancing::SidekiqServerMiddleware::JobReplicaNotUpToDate')
           end
 
           include_examples 'load balancing strategy', 'retry'
         end
 
         context 'when job is retried' do
-          before do
-            expect do
-              process_job(job)
-            end.to raise_error(Sidekiq::JobRetry::Skip)
-          end
+          let(:job) { { "retry" => 3, "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", 'database_replica_location' => '0/D525E3A8', 'retry_count' => 0 } }
 
           context 'and replica still lagging behind' do
             include_examples 'stick to the primary', 'primary'
@@ -191,7 +199,9 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware do
 
       context 'when replica is not up to date' do
         before do
-          allow(load_balancer).to receive(:select_up_to_date_host).and_return(false)
+          Gitlab::Database::LoadBalancing.each_load_balancer do |lb|
+            allow(lb).to receive(:select_up_to_date_host).and_return(false)
+          end
         end
 
         include_examples 'stick to the primary', 'primary'
@@ -199,8 +209,47 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware do
     end
   end
 
+  describe '#databases_in_sync?' do
+    it 'treats load balancers without WAL entries as in sync' do
+      expect(middleware.send(:databases_in_sync?, {}))
+        .to eq(true)
+    end
+
+    it 'returns true when all load balancers are in sync' do
+      locations = {}
+
+      Gitlab::Database::LoadBalancing.each_load_balancer do |lb|
+        locations[lb.name] = 'foo'
+
+        expect(lb)
+          .to receive(:select_up_to_date_host)
+          .with('foo')
+          .and_return(true)
+      end
+
+      expect(middleware.send(:databases_in_sync?, locations))
+        .to eq(true)
+    end
+
+    it 'returns false when the load balancers are not in sync' do
+      locations = {}
+
+      Gitlab::Database::LoadBalancing.each_load_balancer do |lb|
+        locations[lb.name] = 'foo'
+
+        allow(lb)
+          .to receive(:select_up_to_date_host)
+          .with('foo')
+          .and_return(false)
+      end
+
+      expect(middleware.send(:databases_in_sync?, locations))
+        .to eq(false)
+    end
+  end
+
   def process_job(job)
-    Sidekiq::JobRetry.new.local(worker_class, job, 'default') do
+    Sidekiq::JobRetry.new.local(worker_class, job.to_json, 'default') do
       worker_class.process_job(job)
     end
   end
@@ -212,6 +261,8 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware do
   end
 
   def replication_lag!(exists)
-    allow(load_balancer).to receive(:select_up_to_date_host).and_return(!exists)
+    Gitlab::Database::LoadBalancing.each_load_balancer do |lb|
+      allow(lb).to receive(:select_up_to_date_host).and_return(!exists)
+    end
   end
 end
diff --git a/spec/lib/gitlab/database/load_balancing/sticking_spec.rb b/spec/lib/gitlab/database/load_balancing/sticking_spec.rb
index cf52e59db3a..8ceda52ee85 100644
--- a/spec/lib/gitlab/database/load_balancing/sticking_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/sticking_spec.rb
@@ -3,55 +3,82 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
+  let(:sticking) do
+    described_class.new(ActiveRecord::Base.connection.load_balancer)
+  end
+
   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)
+  describe '#stick_or_unstick_request' do
+    it 'sticks or unsticks a single object and updates the Rack environment' do
+      expect(sticking)
+        .to receive(:unstick_or_continue_sticking)
+        .with(:user, 42)
 
-        described_class.stick_if_necessary(:user, 42)
-      end
+      env = {}
+
+      sticking.stick_or_unstick_request(env, :user, 42)
+
+      expect(env[Gitlab::Database::LoadBalancing::RackMiddleware::STICK_OBJECT].to_a)
+        .to eq([[ActiveRecord::Base, :user, 42]])
     end
 
-    context 'when sticking is enabled' do
-      before do
-        allow(Gitlab::Database::LoadBalancing).to receive(:enable?)
-          .and_return(true)
-      end
+    it 'sticks or unsticks multiple objects and updates the Rack environment' do
+      expect(sticking)
+        .to receive(:unstick_or_continue_sticking)
+        .with(:user, 42)
+        .ordered
 
-      it 'does not stick if no write was performed' do
-        allow(Gitlab::Database::LoadBalancing::Session.current)
-          .to receive(:performed_write?)
-          .and_return(false)
+      expect(sticking)
+        .to receive(:unstick_or_continue_sticking)
+        .with(:runner, '123456789')
+        .ordered
 
-        expect(described_class).not_to receive(:stick)
+      env = {}
 
-        described_class.stick_if_necessary(:user, 42)
-      end
+      sticking.stick_or_unstick_request(env, :user, 42)
+      sticking.stick_or_unstick_request(env, :runner, '123456789')
 
-      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(env[Gitlab::Database::LoadBalancing::RackMiddleware::STICK_OBJECT].to_a).to eq([
+        [ActiveRecord::Base, :user, 42],
+        [ActiveRecord::Base, :runner, '123456789']
+      ])
+    end
+  end
 
-        expect(described_class).to receive(:stick).with(:user, 42)
+  describe '#stick_if_necessary' do
+    it 'does not stick if no write was performed' do
+      allow(Gitlab::Database::LoadBalancing::Session.current)
+        .to receive(:performed_write?)
+        .and_return(false)
 
-        described_class.stick_if_necessary(:user, 42)
-      end
+      expect(sticking).not_to receive(:stick)
+
+      sticking.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(sticking)
+        .to receive(:stick)
+        .with(:user, 42)
+
+      sticking.stick_if_necessary(:user, 42)
     end
   end
 
-  describe '.all_caught_up?' do
-    let(:lb) { double(:lb) }
+  describe '#all_caught_up?' do
+    let(:lb) { ActiveRecord::Base.connection.load_balancer }
     let(:last_write_location) { 'foo' }
 
     before do
-      allow(described_class).to receive(:load_balancer).and_return(lb)
-
-      allow(described_class).to receive(:last_write_location_for)
+      allow(sticking)
+        .to receive(:last_write_location_for)
         .with(:user, 42)
         .and_return(last_write_location)
     end
@@ -60,13 +87,9 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
       let(:last_write_location) { nil }
 
       it 'returns true' do
-        allow(described_class).to receive(:last_write_location_for)
-          .with(:user, 42)
-          .and_return(nil)
-
         expect(lb).not_to receive(:select_up_to_date_host)
 
-        expect(described_class.all_caught_up?(:user, 42)).to eq(true)
+        expect(sticking.all_caught_up?(:user, 42)).to eq(true)
       end
     end
 
@@ -76,9 +99,11 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
       end
 
       it 'returns true, and unsticks' do
-        expect(described_class).to receive(:unstick).with(:user, 42)
+        expect(sticking)
+          .to receive(:unstick)
+          .with(:user, 42)
 
-        expect(described_class.all_caught_up?(:user, 42)).to eq(true)
+        expect(sticking.all_caught_up?(:user, 42)).to eq(true)
       end
 
       it 'notifies with the proper event payload' do
@@ -87,7 +112,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
           .with('caught_up_replica_pick.load_balancing', { result: true })
           .and_call_original
 
-        described_class.all_caught_up?(:user, 42)
+        sticking.all_caught_up?(:user, 42)
       end
     end
 
@@ -97,7 +122,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
       end
 
       it 'returns false' do
-        expect(described_class.all_caught_up?(:user, 42)).to eq(false)
+        expect(sticking.all_caught_up?(:user, 42)).to eq(false)
       end
 
       it 'notifies with the proper event payload' do
@@ -106,42 +131,43 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
           .with('caught_up_replica_pick.load_balancing', { result: false })
           .and_call_original
 
-        described_class.all_caught_up?(:user, 42)
+        sticking.all_caught_up?(: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
+  describe '#unstick_or_continue_sticking' do
+    let(:lb) { ActiveRecord::Base.connection.load_balancer }
 
     it 'simply returns if no write location could be found' do
-      allow(described_class).to receive(:last_write_location_for)
+      allow(sticking)
+        .to receive(:last_write_location_for)
         .with(:user, 42)
         .and_return(nil)
 
       expect(lb).not_to receive(:select_up_to_date_host)
 
-      described_class.unstick_or_continue_sticking(:user, 42)
+      sticking.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)
+      allow(sticking)
+        .to receive(:last_write_location_for)
         .with(:user, 42)
         .and_return('foo')
 
       allow(lb).to receive(:select_up_to_date_host).with('foo').and_return(true)
 
-      expect(described_class).to receive(:unstick).with(:user, 42)
+      expect(sticking)
+        .to receive(:unstick)
+        .with(:user, 42)
 
-      described_class.unstick_or_continue_sticking(:user, 42)
+      sticking.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)
+      allow(sticking)
+        .to receive(:last_write_location_for)
         .with(:user, 42)
         .and_return('foo')
 
@@ -150,184 +176,151 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
       expect(Gitlab::Database::LoadBalancing::Session.current)
         .to receive(:use_primary!)
 
-      described_class.unstick_or_continue_sticking(:user, 42)
+      sticking.unstick_or_continue_sticking(:user, 42)
     end
   end
 
   RSpec.shared_examples 'sticking' do
-    context 'when sticking is disabled' do
-      it 'does not perform any sticking', :aggregate_failures do
-        expect(described_class).not_to receive(:set_write_location_for)
-        expect(Gitlab::Database::LoadBalancing::Session.current).not_to receive(:use_primary!)
-
-        described_class.bulk_stick(:user, ids)
-      end
+    before do
+      allow(ActiveRecord::Base.connection.load_balancer)
+        .to receive(:primary_write_location)
+        .and_return('foo')
     end
 
-    context 'when sticking is enabled' do
-      before do
-        allow(Gitlab::Database::LoadBalancing).to receive(:configured?).and_return(true)
-
-        lb = double(:lb, primary_write_location: 'foo')
+    it 'sticks an entity to the primary', :aggregate_failures do
+      allow(ActiveRecord::Base.connection.load_balancer)
+        .to receive(:primary_only?)
+        .and_return(false)
 
-        allow(described_class).to receive(:load_balancer).and_return(lb)
+      ids.each do |id|
+        expect(sticking)
+          .to receive(:set_write_location_for)
+          .with(:user, id, 'foo')
       end
 
-      it 'sticks an entity to the primary', :aggregate_failures do
-        ids.each do |id|
-          expect(described_class).to receive(:set_write_location_for)
-                                       .with(:user, id, 'foo')
-        end
+      expect(Gitlab::Database::LoadBalancing::Session.current)
+        .to receive(:use_primary!)
 
-        expect(Gitlab::Database::LoadBalancing::Session.current)
-          .to receive(:use_primary!)
+      subject
+    end
 
-        subject
-      end
+    it 'does not update the write location when no replicas are used' do
+      expect(sticking).not_to receive(:set_write_location_for)
+
+      subject
     end
   end
 
-  describe '.stick' do
+  describe '#stick' do
     it_behaves_like 'sticking' do
       let(:ids) { [42] }
-      subject { described_class.stick(:user, ids.first) }
+      subject { sticking.stick(:user, ids.first) }
     end
   end
 
-  describe '.bulk_stick' do
+  describe '#bulk_stick' do
     it_behaves_like 'sticking' do
       let(:ids) { [42, 43] }
-      subject { described_class.bulk_stick(:user, ids) }
+      subject { sticking.bulk_stick(:user, ids) }
     end
   end
 
-  describe '.mark_primary_write_location' do
-    context 'when enabled' do
-      before do
-        allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
-        allow(Gitlab::Database::LoadBalancing).to receive(:configured?).and_return(true)
-      end
-
-      it 'updates the write location with the load balancer' do
-        lb = double(:lb, primary_write_location: 'foo')
+  describe '#mark_primary_write_location' do
+    it 'updates the write location with the load balancer' do
+      allow(ActiveRecord::Base.connection.load_balancer)
+        .to receive(:primary_write_location)
+        .and_return('foo')
 
-        allow(described_class).to receive(:load_balancer).and_return(lb)
+      allow(ActiveRecord::Base.connection.load_balancer)
+        .to receive(:primary_only?)
+        .and_return(false)
 
-        expect(described_class).to receive(:set_write_location_for)
-          .with(:user, 42, 'foo')
+      expect(sticking)
+        .to receive(:set_write_location_for)
+        .with(:user, 42, 'foo')
 
-        described_class.mark_primary_write_location(:user, 42)
-      end
+      sticking.mark_primary_write_location(:user, 42)
     end
 
-    context 'when load balancing is configured but not enabled' do
-      before do
-        allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false)
-        allow(Gitlab::Database::LoadBalancing).to receive(:configured?).and_return(true)
-      end
-
-      it 'updates the write location with the main ActiveRecord connection' do
-        allow(described_class).to receive(:load_balancer).and_return(nil)
-        expect(ActiveRecord::Base).to receive(:connection).and_call_original
-        expect(described_class).to receive(:set_write_location_for)
-          .with(:user, 42, anything)
+    it 'does nothing when no replicas are used' do
+      expect(sticking).not_to receive(:set_write_location_for)
 
-        described_class.mark_primary_write_location(:user, 42)
-      end
-
-      context 'when write location is nil' do
-        before do
-          allow(Gitlab::Database.main).to receive(:get_write_location).and_return(nil)
-        end
+      sticking.mark_primary_write_location(:user, 42)
+    end
+  end
 
-        it 'does not update the write location' do
-          expect(described_class).not_to receive(:set_write_location_for)
+  describe '#unstick' do
+    it 'removes the sticking data from Redis' do
+      sticking.set_write_location_for(:user, 4, 'foo')
+      sticking.unstick(:user, 4)
 
-          described_class.mark_primary_write_location(:user, 42)
-        end
-      end
+      expect(sticking.last_write_location_for(:user, 4)).to be_nil
     end
 
-    context 'when load balancing is disabled' do
-      before do
-        allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false)
-        allow(Gitlab::Database::LoadBalancing).to receive(:configured?).and_return(false)
+    it 'removes the old key' do
+      Gitlab::Redis::SharedState.with do |redis|
+        redis.set(sticking.send(:old_redis_key_for, :user, 4), 'foo', ex: 30)
       end
 
-      it 'updates the write location with the main ActiveRecord connection' do
-        expect(described_class).not_to receive(:set_write_location_for)
-
-        described_class.mark_primary_write_location(:user, 42)
-      end
+      sticking.unstick(:user, 4)
+      expect(sticking.last_write_location_for(:user, 4)).to be_nil
     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)
+  describe '#last_write_location_for' do
+    it 'returns the last WAL write location for a user' do
+      sticking.set_write_location_for(:user, 4, 'foo')
 
-      expect(described_class.last_write_location_for(:user, 4)).to be_nil
+      expect(sticking.last_write_location_for(:user, 4)).to eq('foo')
     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')
+    it 'falls back to reading the old key' do
+      Gitlab::Redis::SharedState.with do |redis|
+        redis.set(sticking.send(:old_redis_key_for, :user, 4), 'foo', ex: 30)
+      end
 
-      expect(described_class.last_write_location_for(:user, 4)).to eq('foo')
+      expect(sticking.last_write_location_for(:user, 4)).to eq('foo')
     end
   end
 
-  describe '.redis_key_for' do
+  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')
+      expect(sticking.redis_key_for(:user, 42))
+        .to eq('database-load-balancing/write-location/main/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
-
-  describe '.select_caught_up_replicas' do
-    let(:lb) { double(:lb) }
-
-    before do
-      allow(described_class).to receive(:load_balancer).and_return(lb)
-    end
+  describe '#select_caught_up_replicas' do
+    let(:lb) { ActiveRecord::Base.connection.load_balancer }
 
     context 'with no write location' do
       before do
-        allow(described_class).to receive(:last_write_location_for)
-          .with(:project, 42).and_return(nil)
+        allow(sticking)
+          .to receive(:last_write_location_for)
+          .with(:project, 42)
+          .and_return(nil)
       end
 
       it 'returns false and does not try to find caught up hosts' do
         expect(lb).not_to receive(:select_up_to_date_host)
-        expect(described_class.select_caught_up_replicas(:project, 42)).to be false
+        expect(sticking.select_caught_up_replicas(:project, 42)).to be false
       end
     end
 
     context 'with write location' do
       before do
-        allow(described_class).to receive(:last_write_location_for)
-          .with(:project, 42).and_return('foo')
+        allow(sticking)
+          .to receive(:last_write_location_for)
+          .with(:project, 42)
+          .and_return('foo')
       end
 
       it 'returns true, selects hosts, and unsticks if any secondary has caught up' do
         expect(lb).to receive(:select_up_to_date_host).and_return(true)
-        expect(described_class).to receive(:unstick).with(:project, 42)
-        expect(described_class.select_caught_up_replicas(:project, 42)).to be true
+        expect(sticking)
+          .to receive(:unstick)
+          .with(:project, 42)
+        expect(sticking.select_caught_up_replicas(:project, 42)).to be true
       end
     end
   end