Add latest changes from gitlab-org/gitlab@14-0-stable-eev14.0.0-rc42

24 files changed, 4388 insertions, 18 deletions

diff --git a/spec/lib/gitlab/database/background_migration/batched_job_spec.rb b/spec/lib/gitlab/database/background_migration/batched_job_spec.rb
index 78e0b7627e9..2de784d3e16 100644
--- a/spec/lib/gitlab/database/background_migration/batched_job_spec.rb
+++ b/spec/lib/gitlab/database/background_migration/batched_job_spec.rb
@@ -49,16 +49,6 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedJob, type: :model d
     let(:batched_job) { build(:batched_background_migration_job) }
     let(:batched_migration) { batched_job.batched_migration }
 
-    describe '#migration_aborted?' do
-      before do
-        batched_migration.status = :aborted
-      end
-
-      it 'returns the migration aborted?' do
-        expect(batched_job.migration_aborted?).to eq(batched_migration.aborted?)
-      end
-    end
-
     describe '#migration_job_class' do
       it 'returns the migration job_class' do
         expect(batched_job.migration_job_class).to eq(batched_migration.job_class)
diff --git a/spec/lib/gitlab/database/background_migration/batched_migration_spec.rb b/spec/lib/gitlab/database/background_migration/batched_migration_spec.rb
index 43e34325419..d881390cd52 100644
--- a/spec/lib/gitlab/database/background_migration/batched_migration_spec.rb
+++ b/spec/lib/gitlab/database/background_migration/batched_migration_spec.rb
@@ -19,6 +19,12 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
     end
   end
 
+  describe 'validations' do
+    subject { build(:batched_background_migration) }
+
+    it { is_expected.to validate_uniqueness_of(:job_arguments).scoped_to(:job_class_name, :table_name, :column_name) }
+  end
+
   describe '.queue_order' do
     let!(:migration1) { create(:batched_background_migration) }
     let!(:migration2) { create(:batched_background_migration) }
@@ -36,6 +42,38 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
 
     it 'returns the first active migration according to queue order' do
       expect(described_class.active_migration).to eq(migration2)
+      create(:batched_background_migration_job, batched_migration: migration1, batch_size: 1000, status: :succeeded)
+    end
+  end
+
+  describe '.queued' do
+    let!(:migration1) { create(:batched_background_migration, :finished) }
+    let!(:migration2) { create(:batched_background_migration, :paused) }
+    let!(:migration3) { create(:batched_background_migration, :active) }
+
+    it 'returns active and paused migrations' do
+      expect(described_class.queued).to contain_exactly(migration2, migration3)
+    end
+  end
+
+  describe '.successful_rows_counts' do
+    let!(:migration1) { create(:batched_background_migration) }
+    let!(:migration2) { create(:batched_background_migration) }
+    let!(:migration_without_jobs) { create(:batched_background_migration) }
+
+    before do
+      create(:batched_background_migration_job, batched_migration: migration1, batch_size: 1000, status: :succeeded)
+      create(:batched_background_migration_job, batched_migration: migration1, batch_size: 200, status: :failed)
+      create(:batched_background_migration_job, batched_migration: migration2, batch_size: 500, status: :succeeded)
+      create(:batched_background_migration_job, batched_migration: migration2, batch_size: 200, status: :running)
+    end
+
+    it 'returns totals from successful jobs' do
+      results = described_class.successful_rows_counts([migration1, migration2, migration_without_jobs])
+
+      expect(results[migration1.id]).to eq(1000)
+      expect(results[migration2.id]).to eq(500)
+      expect(results[migration_without_jobs.id]).to eq(nil)
     end
   end
 
@@ -324,4 +362,29 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
       subject
     end
   end
+
+  describe '.for_configuration' do
+    let!(:migration) do
+      create(
+        :batched_background_migration,
+        job_class_name: 'MyJobClass',
+        table_name: :projects,
+        column_name: :id,
+        job_arguments: [[:id], [:id_convert_to_bigint]]
+      )
+    end
+
+    before do
+      create(:batched_background_migration, job_class_name: 'OtherClass')
+      create(:batched_background_migration, table_name: 'other_table')
+      create(:batched_background_migration, column_name: 'other_column')
+      create(:batched_background_migration, job_arguments: %w[other arguments])
+    end
+
+    it 'finds the migration matching the given configuration parameters' do
+      actual = described_class.for_configuration('MyJobClass', :projects, :id, [[:id], [:id_convert_to_bigint]])
+
+      expect(actual).to contain_exactly(migration)
+    end
+  end
 end
diff --git a/spec/lib/gitlab/database/consistency_spec.rb b/spec/lib/gitlab/database/consistency_spec.rb
new file mode 100644
index 00000000000..35fa65512ae
--- /dev/null
+++ b/spec/lib/gitlab/database/consistency_spec.rb
@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Consistency do
+  let(:session) do
+    Gitlab::Database::LoadBalancing::Session.current
+  end
+
+  describe '.with_read_consistency' do
+    it 'sticks to primary database' do
+      expect(session).not_to be_using_primary
+
+      block = -> (&control) do
+        described_class.with_read_consistency do
+          expect(session).to be_using_primary
+
+          control.call
+        end
+      end
+
+      expect { |probe| block.call(&probe) }.to yield_control
+    end
+  end
+end
diff --git a/spec/lib/gitlab/database/count/reltuples_count_strategy_spec.rb b/spec/lib/gitlab/database/count/reltuples_count_strategy_spec.rb
index 324ed498abc..cdcc862c376 100644
--- a/spec/lib/gitlab/database/count/reltuples_count_strategy_spec.rb
+++ b/spec/lib/gitlab/database/count/reltuples_count_strategy_spec.rb
@@ -27,7 +27,7 @@ RSpec.describe Gitlab::Database::Count::ReltuplesCountStrategy do
     end
 
     context 'when models using single-type inheritance are used' do
-      let(:models) { [Group, CiService, Namespace] }
+      let(:models) { [Group, Integrations::BaseCi, Namespace] }
 
       before do
         models.each do |model|
diff --git a/spec/lib/gitlab/database/dynamic_model_helpers_spec.rb b/spec/lib/gitlab/database/dynamic_model_helpers_spec.rb
index 23ad621d0ee..0844616ee1c 100644
--- a/spec/lib/gitlab/database/dynamic_model_helpers_spec.rb
+++ b/spec/lib/gitlab/database/dynamic_model_helpers_spec.rb
@@ -3,12 +3,12 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::DynamicModelHelpers do
+  let(:including_class) { Class.new.include(described_class) }
+  let(:table_name) { 'projects' }
+
   describe '#define_batchable_model' do
     subject { including_class.new.define_batchable_model(table_name) }
 
-    let(:including_class) { Class.new.include(described_class) }
-    let(:table_name) { 'projects' }
-
     it 'is an ActiveRecord model' do
       expect(subject.ancestors).to include(ActiveRecord::Base)
     end
@@ -25,4 +25,86 @@ RSpec.describe Gitlab::Database::DynamicModelHelpers do
       expect(subject.inheritance_column).to eq('_type_disabled')
     end
   end
+
+  describe '#each_batch' do
+    subject { including_class.new }
+
+    before do
+      create_list(:project, 2)
+    end
+
+    context 'when no transaction is open' do
+      before do
+        allow(subject).to receive(:transaction_open?).and_return(false)
+      end
+
+      it 'iterates table in batches' do
+        each_batch_size = ->(&block) do
+          subject.each_batch(table_name, of: 1) do |batch|
+            block.call(batch.size)
+          end
+        end
+
+        expect { |b| each_batch_size.call(&b) }
+          .to yield_successive_args(1, 1)
+      end
+    end
+
+    context 'when transaction is open' do
+      before do
+        allow(subject).to receive(:transaction_open?).and_return(true)
+      end
+
+      it 'raises an error' do
+        expect { subject.each_batch(table_name, of: 1) { |batch| batch.size } }
+          .to raise_error(RuntimeError, /each_batch should not run inside a transaction/)
+      end
+    end
+  end
+
+  describe '#each_batch_range' do
+    subject { including_class.new }
+
+    let(:first_project) { create(:project) }
+    let(:second_project) { create(:project) }
+
+    context 'when no transaction is open' do
+      before do
+        allow(subject).to receive(:transaction_open?).and_return(false)
+      end
+
+      it 'iterates table in batch ranges' do
+        expect { |b| subject.each_batch_range(table_name, of: 1, &b) }
+          .to yield_successive_args(
+            [first_project.id, first_project.id],
+            [second_project.id, second_project.id]
+          )
+      end
+
+      it 'yields only one batch if bigger than the table size' do
+        expect { |b| subject.each_batch_range(table_name, of: 2, &b) }
+          .to yield_successive_args([first_project.id, second_project.id])
+      end
+
+      it 'makes it possible to apply a scope' do
+        each_batch_limited = ->(&b) do
+          subject.each_batch_range(table_name, scope: ->(table) { table.limit(1) }, of: 1, &b)
+        end
+
+        expect { |b| each_batch_limited.call(&b) }
+          .to yield_successive_args([first_project.id, first_project.id])
+      end
+    end
+
+    context 'when transaction is open' do
+      before do
+        allow(subject).to receive(:transaction_open?).and_return(true)
+      end
+
+      it 'raises an error' do
+        expect { subject.each_batch_range(table_name, of: 1) { 1 } }
+          .to raise_error(RuntimeError, /each_batch should not run inside a transaction/)
+      end
+    end
+  end
 end
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
new file mode 100644
index 00000000000..8886ce9756d
--- /dev/null
+++ b/spec/lib/gitlab/database/load_balancing/active_record_proxy_spec.rb
@@ -0,0 +1,20 @@
+# 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/connection_proxy_spec.rb b/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb
new file mode 100644
index 00000000000..015dd2ba8d2
--- /dev/null
+++ b/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb
@@ -0,0 +1,316 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
+  let(:proxy) { described_class.new }
+
+  describe '#select' do
+    it 'performs a read' do
+      expect(proxy).to receive(:read_using_load_balancer).with(:select, ['foo'])
+
+      proxy.select('foo')
+    end
+  end
+
+  describe '#select_all' do
+    let(:override_proxy) { ActiveRecord::Base.connection.class }
+
+    # We can't use :Gitlab::Utils::Override because this method is dynamically prepended
+    it 'method signatures match' do
+      expect(proxy.method(:select_all).parameters).to eq(override_proxy.instance_method(:select_all).parameters)
+    end
+
+    describe 'using a SELECT query' do
+      it 'runs the query on a secondary' do
+        arel = double(:arel)
+
+        expect(proxy).to receive(:read_using_load_balancer)
+          .with(:select_all, [arel, 'foo', []])
+
+        proxy.select_all(arel, 'foo')
+      end
+    end
+
+    describe 'using a SELECT FOR UPDATE query' do
+      it 'runs the query on the primary and sticks to it' do
+        arel = double(:arel, locked: true)
+
+        expect(proxy).to receive(:write_using_load_balancer)
+          .with(:select_all, [arel, 'foo', []], sticky: true)
+
+        proxy.select_all(arel, 'foo')
+      end
+    end
+  end
+
+  Gitlab::Database::LoadBalancing::ConnectionProxy::NON_STICKY_READS.each do |name|
+    describe "#{name}" do
+      it 'runs the query on the replica' do
+        expect(proxy).to receive(:read_using_load_balancer)
+          .with(name, ['foo'])
+
+        proxy.send(name, 'foo')
+      end
+    end
+  end
+
+  Gitlab::Database::LoadBalancing::ConnectionProxy::STICKY_WRITES.each do |name|
+    describe "#{name}" do
+      it 'runs the query on the primary and sticks to it' do
+        expect(proxy).to receive(:write_using_load_balancer)
+          .with(name, ['foo'], sticky: true)
+
+        proxy.send(name, 'foo')
+      end
+    end
+  end
+
+  describe '.insert_all!' do
+    before do
+      ActiveRecord::Schema.define do
+        create_table :connection_proxy_bulk_insert, force: true do |t|
+          t.string :name, null: true
+        end
+      end
+    end
+
+    after do
+      ActiveRecord::Schema.define do
+        drop_table :connection_proxy_bulk_insert, force: true
+      end
+    end
+
+    let(:model_class) do
+      Class.new(ApplicationRecord) do
+        self.table_name = "connection_proxy_bulk_insert"
+      end
+    end
+
+    it 'inserts data in bulk' do
+      expect(model_class).to receive(:connection)
+        .at_least(:once)
+        .and_return(proxy)
+
+      expect(proxy).to receive(:write_using_load_balancer)
+        .at_least(:once)
+        .and_call_original
+
+      expect do
+        model_class.insert_all! [
+          { name: "item1" },
+          { name: "item2" }
+        ]
+      end.to change { model_class.count }.by(2)
+    end
+  end
+
+  # We have an extra test for #transaction here to make sure that nested queries
+  # are also sent to a primary.
+  describe '#transaction' do
+    let(:session) { double(:session) }
+
+    before do
+      allow(Gitlab::Database::LoadBalancing::Session).to receive(:current)
+        .and_return(session)
+    end
+
+    context 'session fallbacks ambiguous queries to replicas' do
+      let(:replica) { double(:connection) }
+
+      before do
+        allow(session).to receive(:fallback_to_replicas_for_ambiguous_queries?).and_return(true)
+        allow(session).to receive(:use_primary?).and_return(false)
+        allow(replica).to receive(:transaction).and_yield
+        allow(replica).to receive(:select)
+      end
+
+      context 'with a read query' do
+        it 'runs the transaction and any nested queries on the replica' do
+          expect(proxy.load_balancer).to receive(:read)
+            .twice.and_yield(replica)
+          expect(proxy.load_balancer).not_to receive(:read_write)
+          expect(session).not_to receive(:write!)
+
+          proxy.transaction { proxy.select('true') }
+        end
+      end
+
+      context 'with a write query' do
+        it 'raises an exception' do
+          allow(proxy.load_balancer).to receive(:read).and_yield(replica)
+          allow(proxy.load_balancer).to receive(:read_write).and_yield(replica)
+
+          expect do
+            proxy.transaction { proxy.insert('something') }
+          end.to raise_error(Gitlab::Database::LoadBalancing::ConnectionProxy::WriteInsideReadOnlyTransactionError)
+        end
+      end
+    end
+
+    context 'session does not fallback to replicas for ambiguous queries' do
+      let(:primary) { double(:connection) }
+
+      before do
+        allow(session).to receive(:fallback_to_replicas_for_ambiguous_queries?).and_return(false)
+        allow(session).to receive(:use_replicas_for_read_queries?).and_return(false)
+        allow(session).to receive(:use_primary?).and_return(true)
+        allow(primary).to receive(:transaction).and_yield
+        allow(primary).to receive(:select)
+        allow(primary).to receive(:insert)
+      end
+
+      context 'with a read query' do
+        it 'runs the transaction and any nested queries on the primary and stick to it' do
+          expect(proxy.load_balancer).to receive(:read_write)
+            .twice.and_yield(primary)
+          expect(proxy.load_balancer).not_to receive(:read)
+          expect(session).to receive(:write!)
+
+          proxy.transaction { proxy.select('true') }
+        end
+      end
+
+      context 'with a write query' do
+        it 'runs the transaction and any nested queries on the primary and stick to it' do
+          expect(proxy.load_balancer).to receive(:read_write)
+            .twice.and_yield(primary)
+          expect(proxy.load_balancer).not_to receive(:read)
+          expect(session).to receive(:write!).twice
+
+          proxy.transaction { proxy.insert('something') }
+        end
+      end
+    end
+  end
+
+  describe '#method_missing' do
+    it 'runs the query on the primary without sticking to it' do
+      expect(proxy).to receive(:write_using_load_balancer)
+        .with(:foo, ['foo'])
+
+      proxy.foo('foo')
+    end
+
+    it 'properly forwards trailing hash arguments' do
+      allow(proxy.load_balancer).to receive(:read_write)
+
+      expect(proxy).to receive(:write_using_load_balancer).and_call_original
+
+      expect { proxy.case_sensitive_comparison(:table, :attribute, :column, { value: :value, format: :format }) }
+        .not_to raise_error
+    end
+
+    context 'current session prefers to fallback ambiguous queries to replicas' do
+      let(:session) { double(:session) }
+
+      before do
+        allow(Gitlab::Database::LoadBalancing::Session).to receive(:current)
+          .and_return(session)
+        allow(session).to receive(:fallback_to_replicas_for_ambiguous_queries?).and_return(true)
+        allow(session).to receive(:use_primary?).and_return(false)
+      end
+
+      it 'runs the query on the replica' do
+        expect(proxy).to receive(:read_using_load_balancer).with(:foo, ['foo'])
+
+        proxy.foo('foo')
+      end
+
+      it 'properly forwards trailing hash arguments' do
+        allow(proxy.load_balancer).to receive(:read)
+
+        expect(proxy).to receive(:read_using_load_balancer).and_call_original
+
+        expect { proxy.case_sensitive_comparison(:table, :attribute, :column, { value: :value, format: :format }) }
+          .not_to raise_error
+      end
+    end
+  end
+
+  describe '#read_using_load_balancer' do
+    let(:session) { double(:session) }
+    let(:connection) { double(:connection) }
+
+    before do
+      allow(Gitlab::Database::LoadBalancing::Session).to receive(:current)
+        .and_return(session)
+    end
+
+    context 'with a regular session' do
+      it 'uses a secondary' do
+        allow(session).to receive(:use_primary?).and_return(false)
+        allow(session).to receive(:use_replicas_for_read_queries?).and_return(false)
+
+        expect(connection).to receive(:foo).with('foo')
+        expect(proxy.load_balancer).to receive(:read).and_yield(connection)
+
+        proxy.read_using_load_balancer(:foo, ['foo'])
+      end
+    end
+
+    context 'with a regular session and forcing all reads to replicas' do
+      it 'uses a secondary' do
+        allow(session).to receive(:use_primary?).and_return(false)
+        allow(session).to receive(:use_replicas_for_read_queries?).and_return(true)
+
+        expect(connection).to receive(:foo).with('foo')
+        expect(proxy.load_balancer).to receive(:read).and_yield(connection)
+
+        proxy.read_using_load_balancer(:foo, ['foo'])
+      end
+    end
+
+    context 'with a session using the primary but forcing all reads to replicas' do
+      it 'uses a secondary' do
+        allow(session).to receive(:use_primary?).and_return(true)
+        allow(session).to receive(:use_replicas_for_read_queries?).and_return(true)
+
+        expect(connection).to receive(:foo).with('foo')
+        expect(proxy.load_balancer).to receive(:read).and_yield(connection)
+
+        proxy.read_using_load_balancer(:foo, ['foo'])
+      end
+    end
+
+    describe 'with a session using the primary' do
+      it 'uses the primary' do
+        allow(session).to receive(:use_primary?).and_return(true)
+        allow(session).to receive(:use_replicas_for_read_queries?).and_return(false)
+
+        expect(connection).to receive(:foo).with('foo')
+
+        expect(proxy.load_balancer).to receive(:read_write)
+          .and_yield(connection)
+
+        proxy.read_using_load_balancer(:foo, ['foo'])
+      end
+    end
+  end
+
+  describe '#write_using_load_balancer' do
+    let(:session) { double(:session) }
+    let(:connection) { double(:connection) }
+
+    before do
+      allow(Gitlab::Database::LoadBalancing::Session).to receive(:current)
+        .and_return(session)
+    end
+
+    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(:write!)
+
+      proxy.write_using_load_balancer(:foo, ['foo'])
+    end
+
+    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(:write!)
+
+      proxy.write_using_load_balancer(:foo, ['foo'], sticky: true)
+    end
+  end
+end
diff --git a/spec/lib/gitlab/database/load_balancing/host_list_spec.rb b/spec/lib/gitlab/database/load_balancing/host_list_spec.rb
new file mode 100644
index 00000000000..873b599f84d
--- /dev/null
+++ b/spec/lib/gitlab/database/load_balancing/host_list_spec.rb
@@ -0,0 +1,188 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::LoadBalancing::HostList do
+  def expect_metrics(hosts)
+    expect(Gitlab::Metrics.registry.get(:db_load_balancing_hosts).get({})).to eq(hosts)
+  end
+
+  before do
+    allow(Gitlab::Database)
+      .to receive(:create_connection_pool)
+      .and_return(ActiveRecord::Base.connection_pool)
+  end
+
+  let(:load_balancer) { double(:load_balancer) }
+  let(:host_count) { 2 }
+
+  let(:host_list) do
+    hosts = Array.new(host_count) do
+      Gitlab::Database::LoadBalancing::Host.new('localhost', load_balancer, port: 5432)
+    end
+
+    described_class.new(hosts)
+  end
+
+  describe '#initialize' do
+    it 'sets metrics for current number of hosts and current index' do
+      host_list
+
+      expect_metrics(2)
+    end
+  end
+
+  describe '#length' do
+    it 'returns the number of hosts in the list' do
+      expect(host_list.length).to eq(2)
+    end
+  end
+
+  describe '#host_names_and_ports' do
+    context 'with ports' do
+      it 'returns the host names of all hosts' do
+        hosts = [
+          ['localhost', 5432],
+          ['localhost', 5432]
+        ]
+
+        expect(host_list.host_names_and_ports).to eq(hosts)
+      end
+    end
+
+    context 'without ports' do
+      let(:host_list) do
+        hosts = Array.new(2) do
+          Gitlab::Database::LoadBalancing::Host.new('localhost', load_balancer)
+        end
+
+        described_class.new(hosts)
+      end
+
+      it 'returns the host names of all hosts' do
+        hosts = [
+          ['localhost', nil],
+          ['localhost', nil]
+        ]
+
+        expect(host_list.host_names_and_ports).to eq(hosts)
+      end
+    end
+  end
+
+  describe '#manage_pool?' do
+    before do
+      allow(Gitlab::Database).to receive(:create_connection_pool) { double(:connection) }
+    end
+
+    context 'when the testing pool belongs to one host of the host list' do
+      it 'returns true' do
+        pool = host_list.hosts.first.pool
+
+        expect(host_list.manage_pool?(pool)).to be(true)
+      end
+    end
+
+    context 'when the testing pool belongs to a former host of the host list' do
+      it 'returns false' do
+        pool = host_list.hosts.first.pool
+        host_list.hosts = [
+          Gitlab::Database::LoadBalancing::Host.new('foo', load_balancer)
+        ]
+
+        expect(host_list.manage_pool?(pool)).to be(false)
+      end
+    end
+
+    context 'when the testing pool belongs to a new host of the host list' do
+      it 'returns true' do
+        host = Gitlab::Database::LoadBalancing::Host.new('foo', load_balancer)
+        host_list.hosts = [host]
+
+        expect(host_list.manage_pool?(host.pool)).to be(true)
+      end
+    end
+
+    context 'when the testing pool does not have any relation with the host list' do
+      it 'returns false' do
+        host = Gitlab::Database::LoadBalancing::Host.new('foo', load_balancer)
+
+        expect(host_list.manage_pool?(host.pool)).to be(false)
+      end
+    end
+  end
+
+  describe '#hosts' do
+    it 'returns a copy of the host' do
+      first = host_list.hosts
+
+      expect(host_list.hosts).to eq(first)
+      expect(host_list.hosts.object_id).not_to eq(first.object_id)
+    end
+  end
+
+  describe '#hosts=' do
+    it 'updates the list of hosts to use' do
+      host_list.hosts = [
+        Gitlab::Database::LoadBalancing::Host.new('foo', load_balancer)
+      ]
+
+      expect(host_list.length).to eq(1)
+      expect(host_list.hosts[0].host).to eq('foo')
+      expect_metrics(1)
+    end
+  end
+
+  describe '#next' do
+    it 'returns a host' do
+      expect(host_list.next)
+        .to be_an_instance_of(Gitlab::Database::LoadBalancing::Host)
+    end
+
+    it 'cycles through all available hosts' do
+      expect(host_list.next).to eq(host_list.hosts[0])
+      expect_metrics(2)
+
+      expect(host_list.next).to eq(host_list.hosts[1])
+      expect_metrics(2)
+
+      expect(host_list.next).to eq(host_list.hosts[0])
+      expect_metrics(2)
+    end
+
+    it 'skips hosts that are offline' do
+      allow(host_list.hosts[0]).to receive(:online?).and_return(false)
+
+      expect(host_list.next).to eq(host_list.hosts[1])
+      expect_metrics(2)
+    end
+
+    it 'returns nil if no hosts are online' do
+      host_list.hosts.each do |host|
+        allow(host).to receive(:online?).and_return(false)
+      end
+
+      expect(host_list.next).to be_nil
+      expect_metrics(2)
+    end
+
+    it 'returns nil if no hosts are available' do
+      expect(described_class.new.next).to be_nil
+    end
+  end
+
+  describe '#shuffle' do
+    let(:host_count) { 3 }
+
+    it 'randomizes the list' do
+      2.times do
+        all_hosts = host_list.hosts
+
+        host_list.shuffle
+
+        expect(host_list.length).to eq(host_count)
+        expect(host_list.hosts).to contain_exactly(*all_hosts)
+      end
+    end
+  end
+end
diff --git a/spec/lib/gitlab/database/load_balancing/host_spec.rb b/spec/lib/gitlab/database/load_balancing/host_spec.rb
new file mode 100644
index 00000000000..4dfddef68c8
--- /dev/null
+++ b/spec/lib/gitlab/database/load_balancing/host_spec.rb
@@ -0,0 +1,445 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::LoadBalancing::Host do
+  let(:load_balancer) do
+    Gitlab::Database::LoadBalancing::LoadBalancer.new(%w[localhost])
+  end
+
+  let(:host) { load_balancer.host_list.hosts.first }
+
+  before do
+    allow(Gitlab::Database).to receive(:create_connection_pool)
+      .and_return(ActiveRecord::Base.connection_pool)
+  end
+
+  def raise_and_wrap(wrapper, original)
+    raise original
+  rescue original.class
+    raise wrapper, 'boom'
+  end
+
+  def wrapped_exception(wrapper, original)
+    raise_and_wrap(wrapper, original.new)
+  rescue wrapper => error
+    error
+  end
+
+  describe '#connection' do
+    it 'returns a connection from the pool' do
+      expect(host.pool).to receive(:connection)
+
+      host.connection
+    end
+  end
+
+  describe '#disconnect!' do
+    it 'disconnects the pool' do
+      connection = double(:connection, in_use?: false)
+      pool = double(:pool, connections: [connection])
+
+      allow(host)
+        .to receive(:pool)
+        .and_return(pool)
+
+      expect(host)
+        .not_to receive(:sleep)
+
+      expect(host.pool)
+        .to receive(:disconnect!)
+
+      host.disconnect!
+    end
+
+    it 'disconnects the pool when waiting for connections takes too long' do
+      connection = double(:connection, in_use?: true)
+      pool = double(:pool, connections: [connection])
+
+      allow(host)
+        .to receive(:pool)
+        .and_return(pool)
+
+      expect(host.pool)
+        .to receive(:disconnect!)
+
+      host.disconnect!(1)
+    end
+  end
+
+  describe '#release_connection' do
+    it 'releases the current connection from the pool' do
+      expect(host.pool).to receive(:release_connection)
+
+      host.release_connection
+    end
+  end
+
+  describe '#offline!' do
+    it 'marks the host as offline' do
+      expect(host.pool).to receive(:disconnect!)
+
+      expect(Gitlab::Database::LoadBalancing::Logger).to receive(:warn)
+        .with(hash_including(event: :host_offline))
+        .and_call_original
+
+      host.offline!
+    end
+  end
+
+  describe '#online?' do
+    context 'when the replica status is recent enough' do
+      before do
+        expect(host).to receive(:check_replica_status?).and_return(false)
+      end
+
+      it 'returns the latest status' do
+        expect(host).not_to receive(:refresh_status)
+        expect(Gitlab::Database::LoadBalancing::Logger).not_to receive(:info)
+        expect(Gitlab::Database::LoadBalancing::Logger).not_to receive(:warn)
+
+        expect(host).to be_online
+      end
+
+      it 'returns an offline status' do
+        host.offline!
+
+        expect(host).not_to receive(:refresh_status)
+        expect(Gitlab::Database::LoadBalancing::Logger).not_to receive(:info)
+        expect(Gitlab::Database::LoadBalancing::Logger).not_to receive(:warn)
+
+        expect(host).not_to be_online
+      end
+    end
+
+    context 'when the replica status is outdated' do
+      before do
+        expect(host)
+          .to receive(:check_replica_status?)
+          .and_return(true)
+      end
+
+      it 'refreshes the status' do
+        expect(Gitlab::Database::LoadBalancing::Logger).to receive(:info)
+          .with(hash_including(event: :host_online))
+          .and_call_original
+
+        expect(host).to be_online
+      end
+
+      context 'and replica is not up to date' do
+        before do
+          expect(host).to receive(:replica_is_up_to_date?).and_return(false)
+        end
+
+        it 'marks the host offline' do
+          expect(Gitlab::Database::LoadBalancing::Logger).to receive(:warn)
+            .with(hash_including(event: :host_offline))
+            .and_call_original
+
+          expect(host).not_to be_online
+        end
+      end
+    end
+
+    context 'when the replica is not online' do
+      it 'returns false when ActionView::Template::Error is raised' do
+        wrapped_error = wrapped_exception(ActionView::Template::Error, StandardError)
+
+        allow(host)
+          .to receive(:check_replica_status?)
+          .and_raise(wrapped_error)
+
+        expect(host).not_to be_online
+      end
+
+      it 'returns false when ActiveRecord::StatementInvalid is raised' do
+        allow(host)
+          .to receive(:check_replica_status?)
+          .and_raise(ActiveRecord::StatementInvalid.new('foo'))
+
+        expect(host).not_to be_online
+      end
+
+      it 'returns false when PG::Error is raised' do
+        allow(host)
+          .to receive(:check_replica_status?)
+          .and_raise(PG::Error)
+
+        expect(host).not_to be_online
+      end
+    end
+  end
+
+  describe '#refresh_status' do
+    it 'refreshes the status' do
+      host.offline!
+
+      expect(host)
+        .to receive(:replica_is_up_to_date?)
+        .and_call_original
+
+      host.refresh_status
+
+      expect(host).to be_online
+    end
+  end
+
+  describe '#check_replica_status?' do
+    it 'returns true when we need to check the replica status' do
+      allow(host)
+        .to receive(:last_checked_at)
+        .and_return(1.year.ago)
+
+      expect(host.check_replica_status?).to eq(true)
+    end
+
+    it 'returns false when we do not need to check the replica status' do
+      freeze_time do
+        allow(host)
+          .to receive(:last_checked_at)
+          .and_return(Time.zone.now)
+
+        expect(host.check_replica_status?).to eq(false)
+      end
+    end
+  end
+
+  describe '#replica_is_up_to_date?' do
+    context 'when the lag time is below the threshold' do
+      it 'returns true' do
+        expect(host)
+          .to receive(:replication_lag_below_threshold?)
+          .and_return(true)
+
+        expect(host.replica_is_up_to_date?).to eq(true)
+      end
+    end
+
+    context 'when the lag time exceeds the threshold' do
+      before do
+        allow(host)
+          .to receive(:replication_lag_below_threshold?)
+          .and_return(false)
+      end
+
+      it 'returns true if the data is recent enough' do
+        expect(host)
+          .to receive(:data_is_recent_enough?)
+          .and_return(true)
+
+        expect(host.replica_is_up_to_date?).to eq(true)
+      end
+
+      it 'returns false when the data is not recent enough' do
+        expect(host)
+          .to receive(:data_is_recent_enough?)
+          .and_return(false)
+
+        expect(host.replica_is_up_to_date?).to eq(false)
+      end
+    end
+  end
+
+  describe '#replication_lag_below_threshold' do
+    it 'returns true when the lag time is below the threshold' do
+      expect(host)
+        .to receive(:replication_lag_time)
+        .and_return(1)
+
+      expect(host.replication_lag_below_threshold?).to eq(true)
+    end
+
+    it 'returns false when the lag time exceeds the threshold' do
+      expect(host)
+        .to receive(:replication_lag_time)
+        .and_return(9000)
+
+      expect(host.replication_lag_below_threshold?).to eq(false)
+    end
+
+    it 'returns false when no lag time could be calculated' do
+      expect(host)
+        .to receive(:replication_lag_time)
+        .and_return(nil)
+
+      expect(host.replication_lag_below_threshold?).to eq(false)
+    end
+  end
+
+  describe '#data_is_recent_enough?' do
+    it 'returns true when the data is recent enough' do
+      expect(host.data_is_recent_enough?).to eq(true)
+    end
+
+    it 'returns false when the data is not recent enough' do
+      diff = Gitlab::Database::LoadBalancing.max_replication_difference * 2
+
+      expect(host)
+        .to receive(:query_and_release)
+        .and_return({ 'diff' => diff })
+
+      expect(host.data_is_recent_enough?).to eq(false)
+    end
+
+    it 'returns false when no lag size could be calculated' do
+      expect(host)
+        .to receive(:replication_lag_size)
+        .and_return(nil)
+
+      expect(host.data_is_recent_enough?).to eq(false)
+    end
+  end
+
+  describe '#replication_lag_time' do
+    it 'returns the lag time as a Float' do
+      expect(host.replication_lag_time).to be_an_instance_of(Float)
+    end
+
+    it 'returns nil when the database query returned no rows' do
+      expect(host)
+        .to receive(:query_and_release)
+        .and_return({})
+
+      expect(host.replication_lag_time).to be_nil
+    end
+  end
+
+  describe '#replication_lag_size' do
+    it 'returns the lag size as an Integer' do
+      expect(host.replication_lag_size).to be_an_instance_of(Integer)
+    end
+
+    it 'returns nil when the database query returned no rows' do
+      expect(host)
+        .to receive(:query_and_release)
+        .and_return({})
+
+      expect(host.replication_lag_size).to be_nil
+    end
+
+    it 'returns nil when the database connection fails' do
+      wrapped_error = wrapped_exception(ActionView::Template::Error, StandardError)
+
+      allow(host)
+        .to receive(:connection)
+        .and_raise(wrapped_error)
+
+      expect(host.replication_lag_size).to be_nil
+    end
+  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
+    end
+  end
+
+  describe '#caught_up?' do
+    let(:connection) { double(:connection) }
+
+    before do
+      allow(connection).to receive(:quote).and_return('foo')
+    end
+
+    it 'returns true when a host has caught up' do
+      allow(host).to receive(:connection).and_return(connection)
+      expect(connection).to receive(:select_all).and_return([{ 'result' => 't' }])
+
+      expect(host.caught_up?('foo')).to eq(true)
+    end
+
+    it 'returns true when a host has caught up' do
+      allow(host).to receive(:connection).and_return(connection)
+      expect(connection).to receive(:select_all).and_return([{ 'result' => true }])
+
+      expect(host.caught_up?('foo')).to eq(true)
+    end
+
+    it 'returns false when a host has not caught up' do
+      allow(host).to receive(:connection).and_return(connection)
+      expect(connection).to receive(:select_all).and_return([{ 'result' => 'f' }])
+
+      expect(host.caught_up?('foo')).to eq(false)
+    end
+
+    it 'returns false when a host has not caught up' do
+      allow(host).to receive(:connection).and_return(connection)
+      expect(connection).to receive(:select_all).and_return([{ 'result' => false }])
+
+      expect(host.caught_up?('foo')).to eq(false)
+    end
+
+    it 'returns false when the connection fails' do
+      wrapped_error = wrapped_exception(ActionView::Template::Error, StandardError)
+
+      allow(host)
+        .to receive(:connection)
+        .and_raise(wrapped_error)
+
+      expect(host.caught_up?('foo')).to eq(false)
+    end
+  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
+      wrapped_error = wrapped_exception(ActionView::Template::Error, StandardError)
+
+      allow(host)
+        .to receive(:connection)
+              .and_raise(wrapped_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({})
+    end
+  end
+
+  describe '#host' do
+    it 'returns the hostname' do
+      expect(host.host).to eq('localhost')
+    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
new file mode 100644
index 00000000000..4705bb23885
--- /dev/null
+++ b/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb
@@ -0,0 +1,522 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
+  let(:pool) { Gitlab::Database.create_connection_pool(2) }
+  let(:conflict_error) { Class.new(RuntimeError) }
+
+  let(:lb) { described_class.new(%w(localhost localhost)) }
+
+  before do
+    allow(Gitlab::Database).to receive(:create_connection_pool)
+      .and_return(pool)
+    stub_const(
+      'Gitlab::Database::LoadBalancing::LoadBalancer::PG::TRSerializationFailure',
+      conflict_error
+    )
+  end
+
+  def raise_and_wrap(wrapper, original)
+    raise original
+  rescue original.class
+    raise wrapper, 'boop'
+  end
+
+  def wrapped_exception(wrapper, original)
+    raise_and_wrap(wrapper, original.new)
+  rescue wrapper => error
+    error
+  end
+
+  def twice_wrapped_exception(top, middle, original)
+    begin
+      raise_and_wrap(middle, original.new)
+    rescue middle => middle_error
+      raise_and_wrap(top, middle_error)
+    end
+  rescue top => top_error
+    top_error
+  end
+
+  describe '#read' do
+    it 'yields a connection for a read' do
+      connection = double(:connection)
+      host = double(:host)
+
+      allow(lb).to receive(:host).and_return(host)
+      allow(host).to receive(:query_cache_enabled).and_return(true)
+
+      expect(host).to receive(:connection).and_return(connection)
+
+      expect { |b| lb.read(&b) }.to yield_with_args(connection)
+    end
+
+    it 'ensures that query cache is enabled' do
+      connection = double(:connection)
+      host = double(:host)
+
+      allow(lb).to receive(:host).and_return(host)
+      allow(host).to receive(:query_cache_enabled).and_return(false)
+      allow(host).to receive(:connection).and_return(connection)
+
+      expect(host).to receive(:enable_query_cache!).once
+
+      lb.read { 10 }
+    end
+
+    it 'marks hosts that are offline' do
+      allow(lb).to receive(:connection_error?).and_return(true)
+
+      expect(lb.host_list.hosts[0]).to receive(:offline!)
+      expect(lb).to receive(:release_host)
+
+      raised = false
+
+      returned = lb.read do
+        unless raised
+          raised = true
+          raise
+        end
+
+        10
+      end
+
+      expect(returned).to eq(10)
+    end
+
+    it 'retries a query in the event of a serialization failure' do
+      raised = false
+
+      expect(lb).to receive(:release_host)
+
+      returned = lb.read do
+        unless raised
+          raised = true
+          raise conflict_error
+        end
+
+        10
+      end
+
+      expect(returned).to eq(10)
+    end
+
+    it 'retries every host at most 3 times when a query conflict is raised' do
+      expect(lb).to receive(:release_host).exactly(6).times
+      expect(lb).to receive(:read_write)
+
+      lb.read { raise conflict_error }
+    end
+
+    it 'uses the primary if no secondaries are available' do
+      allow(lb).to receive(:connection_error?).and_return(true)
+
+      expect(lb.host_list.hosts).to all(receive(:online?).and_return(false))
+
+      expect(lb).to receive(:read_write).and_call_original
+
+      expect { |b| lb.read(&b) }
+        .to yield_with_args(ActiveRecord::Base.retrieve_connection)
+    end
+  end
+
+  describe '#read_write' do
+    it 'yields a connection for a write' do
+      expect { |b| lb.read_write(&b) }
+        .to yield_with_args(ActiveRecord::Base.retrieve_connection)
+    end
+
+    it 'uses a retry with exponential backoffs' do
+      expect(lb).to receive(:retry_with_backoff).and_yield
+
+      lb.read_write { 10 }
+    end
+  end
+
+  describe '#db_role_for_connection' do
+    context 'when the load balancer creates the connection with #read' do
+      it 'returns :replica' do
+        role = nil
+        lb.read do |connection|
+          role = lb.db_role_for_connection(connection)
+        end
+
+        expect(role).to be(:replica)
+      end
+    end
+
+    context 'when the load balancer uses nested #read' do
+      it 'returns :replica' do
+        roles = []
+        lb.read do |connection_1|
+          lb.read do |connection_2|
+            roles << lb.db_role_for_connection(connection_2)
+          end
+          roles << lb.db_role_for_connection(connection_1)
+        end
+
+        expect(roles).to eq([:replica, :replica])
+      end
+    end
+
+    context 'when the load balancer creates the connection with #read_write' do
+      it 'returns :primary' do
+        role = nil
+        lb.read_write do |connection|
+          role = lb.db_role_for_connection(connection)
+        end
+
+        expect(role).to be(:primary)
+      end
+    end
+
+    context 'when the load balancer uses nested #read_write' do
+      it 'returns :primary' do
+        roles = []
+        lb.read_write do |connection_1|
+          lb.read_write do |connection_2|
+            roles << lb.db_role_for_connection(connection_2)
+          end
+          roles << lb.db_role_for_connection(connection_1)
+        end
+
+        expect(roles).to eq([:primary, :primary])
+      end
+    end
+
+    context 'when the load balancer falls back the connection creation to primary' do
+      it 'returns :primary' do
+        allow(lb).to receive(:serialization_failure?).and_return(true)
+
+        role = nil
+        raised = 7 # 2 hosts = 6 retries
+
+        lb.read do |connection|
+          if raised > 0
+            raised -= 1
+            raise
+          end
+
+          role = lb.db_role_for_connection(connection)
+        end
+
+        expect(role).to be(:primary)
+      end
+    end
+
+    context 'when the load balancer uses replica after recovery from a failure' do
+      it 'returns :replica' do
+        allow(lb).to receive(:connection_error?).and_return(true)
+
+        role = nil
+        raised = false
+
+        lb.read do |connection|
+          unless raised
+            raised = true
+            raise
+          end
+
+          role = lb.db_role_for_connection(connection)
+        end
+
+        expect(role).to be(:replica)
+      end
+    end
+
+    context 'when the connection comes from a pool managed by the host list' do
+      it 'returns :replica' do
+        connection = double(:connection)
+        allow(connection).to receive(:pool).and_return(lb.host_list.hosts.first.pool)
+
+        expect(lb.db_role_for_connection(connection)).to be(:replica)
+      end
+    end
+
+    context 'when the connection comes from the primary pool' do
+      it 'returns :primary' do
+        connection = double(:connection)
+        allow(connection).to receive(:pool).and_return(ActiveRecord::Base.connection_pool)
+
+        expect(lb.db_role_for_connection(connection)).to be(:primary)
+      end
+    end
+
+    context 'when the connection does not come from any known pool' do
+      it 'returns nil' do
+        connection = double(:connection)
+        pool = double(:connection_pool)
+        allow(connection).to receive(:pool).and_return(pool)
+
+        expect(lb.db_role_for_connection(connection)).to be(nil)
+      end
+    end
+  end
+
+  describe '#host' do
+    it 'returns the secondary host to use' do
+      expect(lb.host).to be_an_instance_of(Gitlab::Database::LoadBalancing::Host)
+    end
+
+    it 'stores the host in a thread-local variable' do
+      RequestStore.delete(described_class::CACHE_KEY)
+      RequestStore.delete(described_class::VALID_HOSTS_CACHE_KEY)
+
+      expect(lb.host_list).to receive(:next).once.and_call_original
+
+      lb.host
+      lb.host
+    end
+  end
+
+  describe '#release_host' do
+    it 'releases the host and its connection' do
+      host = lb.host
+
+      expect(host).to receive(:disable_query_cache!)
+
+      lb.release_host
+
+      expect(RequestStore[described_class::CACHE_KEY]).to be_nil
+      expect(RequestStore[described_class::VALID_HOSTS_CACHE_KEY]).to be_nil
+    end
+  end
+
+  describe '#release_primary_connection' do
+    it 'releases the connection to the primary' do
+      expect(ActiveRecord::Base.connection_pool).to receive(:release_connection)
+
+      lb.release_primary_connection
+    end
+  end
+
+  describe '#primary_write_location' do
+    it 'returns a String in the right format' do
+      expect(lb.primary_write_location).to match(%r{[A-F0-9]{1,8}/[A-F0-9]{1,8}})
+    end
+
+    it 'raises an error if the write location could not be retrieved' do
+      connection = double(:connection)
+
+      allow(lb).to receive(:read_write).and_yield(connection)
+      allow(connection).to receive(:select_all).and_return([])
+
+      expect { lb.primary_write_location }.to raise_error(RuntimeError)
+    end
+  end
+
+  describe '#all_caught_up?' do
+    it 'returns true if all hosts caught up to the write location' do
+      expect(lb.host_list.hosts).to all(receive(:caught_up?).with('foo').and_return(true))
+
+      expect(lb.all_caught_up?('foo')).to eq(true)
+    end
+
+    it 'returns false if a host has not yet caught up' do
+      expect(lb.host_list.hosts[0]).to receive(:caught_up?)
+        .with('foo')
+        .and_return(true)
+
+      expect(lb.host_list.hosts[1]).to receive(:caught_up?)
+        .with('foo')
+        .and_return(false)
+
+      expect(lb.all_caught_up?('foo')).to eq(false)
+    end
+  end
+
+  describe '#retry_with_backoff' do
+    it 'returns the value returned by the block' do
+      value = lb.retry_with_backoff { 10 }
+
+      expect(value).to eq(10)
+    end
+
+    it 're-raises errors not related to database connections' do
+      expect(lb).not_to receive(:sleep) # to make sure we're not retrying
+
+      expect { lb.retry_with_backoff { raise 'boop' } }
+        .to raise_error(RuntimeError)
+    end
+
+    it 'retries the block when a connection error is raised' do
+      allow(lb).to receive(:connection_error?).and_return(true)
+      expect(lb).to receive(:sleep).with(2)
+      expect(lb).to receive(:release_primary_connection)
+
+      raised = false
+      returned = lb.retry_with_backoff do
+        unless raised
+          raised = true
+          raise
+        end
+
+        10
+      end
+
+      expect(returned).to eq(10)
+    end
+
+    it 're-raises the connection error if the retries did not succeed' do
+      allow(lb).to receive(:connection_error?).and_return(true)
+      expect(lb).to receive(:sleep).with(2).ordered
+      expect(lb).to receive(:sleep).with(4).ordered
+      expect(lb).to receive(:sleep).with(16).ordered
+
+      expect(lb).to receive(:release_primary_connection).exactly(3).times
+
+      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])
+    end
+
+    it 'returns true for a connection error' do
+      error = NotImplementedError.new
+
+      expect(lb.connection_error?(error)).to eq(true)
+    end
+
+    it 'returns true for a wrapped connection error' do
+      wrapped = wrapped_exception(ActiveRecord::StatementInvalid, NotImplementedError)
+
+      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)
+
+      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)
+
+      expect(lb.connection_error?(top)).to eq(true)
+    end
+
+    it 'returns true for an invalid encoding error' do
+      error = RuntimeError.new('invalid encoding name: unicode')
+
+      expect(lb.connection_error?(error)).to eq(true)
+    end
+
+    it 'returns false for errors not related to database connections' do
+      error = RuntimeError.new
+
+      expect(lb.connection_error?(error)).to eq(false)
+    end
+  end
+
+  describe '#serialization_failure?' do
+    let(:conflict_error) { Class.new(RuntimeError) }
+
+    before do
+      stub_const(
+        'Gitlab::Database::LoadBalancing::LoadBalancer::PG::TRSerializationFailure',
+        conflict_error
+      )
+    end
+
+    it 'returns for a serialization error' do
+      expect(lb.serialization_failure?(conflict_error.new)).to eq(true)
+    end
+
+    it 'returns true for a wrapped error' do
+      wrapped = wrapped_exception(ActionView::Template::Error, conflict_error)
+
+      expect(lb.serialization_failure?(wrapped)).to eq(true)
+    end
+  end
+
+  describe '#select_caught_up_hosts' do
+    let(:location) { 'AB/12345'}
+    let(:hosts) { lb.host_list.hosts }
+    let(:valid_host_list) { RequestStore[described_class::VALID_HOSTS_CACHE_KEY] }
+    let(:valid_hosts) { valid_host_list.hosts }
+
+    subject { lb.select_caught_up_hosts(location) }
+
+    context 'when all replicas are caught up' do
+      before do
+        expect(hosts).to all(receive(:caught_up?).with(location).and_return(true))
+      end
+
+      it 'returns true and sets all hosts to valid' do
+        expect(subject).to be true
+        expect(valid_host_list).to be_a(Gitlab::Database::LoadBalancing::HostList)
+        expect(valid_hosts).to contain_exactly(*hosts)
+      end
+    end
+
+    context 'when none of the replicas are caught up' do
+      before do
+        expect(hosts).to all(receive(:caught_up?).with(location).and_return(false))
+      end
+
+      it 'returns false and does not set the valid hosts' do
+        expect(subject).to be false
+        expect(valid_host_list).to be_nil
+      end
+    end
+
+    context 'when one of the replicas is caught up' do
+      before do
+        expect(hosts[0]).to receive(:caught_up?).with(location).and_return(false)
+        expect(hosts[1]).to receive(:caught_up?).with(location).and_return(true)
+      end
+
+      it 'returns true and sets one host to valid' do
+        expect(subject).to be true
+        expect(valid_host_list).to be_a(Gitlab::Database::LoadBalancing::HostList)
+        expect(valid_hosts).to contain_exactly(hosts[1])
+      end
+
+      it 'host always returns the caught-up replica' do
+        subject
+
+        3.times do
+          expect(lb.host).to eq(hosts[1])
+          RequestStore.delete(described_class::CACHE_KEY)
+        end
+      end
+    end
+  end
+
+  describe '#select_caught_up_hosts' do
+    let(:location) { 'AB/12345'}
+    let(:hosts) { lb.host_list.hosts }
+    let(:set_host) { RequestStore[described_class::CACHE_KEY] }
+
+    subject { lb.select_up_to_date_host(location) }
+
+    context 'when none of the replicas are caught up' do
+      before do
+        expect(hosts).to all(receive(:caught_up?).with(location).and_return(false))
+      end
+
+      it 'returns false and does not update the host thread-local variable' do
+        expect(subject).to be false
+        expect(set_host).to be_nil
+      end
+    end
+
+    context 'when any of the replicas is caught up' do
+      before do
+        # `allow` for non-caught up host, because we may not even check it, if will find the caught up one earlier
+        allow(hosts[0]).to receive(:caught_up?).with(location).and_return(false)
+        expect(hosts[1]).to receive(:caught_up?).with(location).and_return(true)
+      end
+
+      it 'returns true and sets host thread-local variable' do
+        expect(subject).to be true
+        expect(set_host).to eq(hosts[1])
+      end
+    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
new file mode 100644
index 00000000000..01367716518
--- /dev/null
+++ b/spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb
@@ -0,0 +1,243 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+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(:multiple_sticking_objects) do
+    Set.new([
+      [:user, 42],
+      [:runner, '123456789'],
+      [:runner, '1234']
+    ])
+  end
+
+  after 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 = {}
+
+      expect(middleware).to receive(:clear).twice
+
+      expect(middleware).to receive(:unstick_or_continue_sticking).with(env)
+      expect(middleware).to receive(:stick_if_necessary).with(env)
+
+      expect(app).to receive(:call).with(env).and_return(10)
+
+      expect(middleware.call(env)).to eq(10)
+    end
+  end
+
+  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.unstick_or_continue_sticking({})
+    end
+
+    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)
+
+      middleware.unstick_or_continue_sticking(env)
+    end
+
+    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)
+        .to receive(:unstick_or_continue_sticking)
+        .with(:user, 42)
+
+      middleware.unstick_or_continue_sticking(env)
+    end
+
+    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)
+        .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
+
+      expect(Gitlab::Database::LoadBalancing::Sticking)
+        .to receive(:unstick_or_continue_sticking)
+        .with(:runner, '1234')
+        .ordered
+
+      middleware.unstick_or_continue_sticking(env)
+    end
+  end
+
+  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)
+
+      middleware.stick_if_necessary({})
+    end
+
+    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)
+
+      middleware.stick_if_necessary(env)
+    end
+
+    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)
+        .to receive(:stick_if_necessary)
+        .with(:user, 42)
+
+      middleware.stick_if_necessary(env)
+    end
+
+    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)
+        .to receive(:stick_if_necessary)
+        .with(:user, 42)
+        .ordered
+
+      expect(Gitlab::Database::LoadBalancing::Sticking)
+        .to receive(:stick_if_necessary)
+        .with(:runner, '123456789')
+        .ordered
+
+      expect(Gitlab::Database::LoadBalancing::Sticking)
+        .to receive(:stick_if_necessary)
+        .with(:runner, '1234')
+        .ordered
+
+      middleware.stick_if_necessary(env)
+    end
+  end
+
+  describe '#clear' do
+    it 'clears the currently used host and session' do
+      lb = double(:lb)
+      session = double(:session)
+
+      allow(middleware).to receive(:load_balancer).and_return(lb)
+
+      expect(lb).to receive(:release_host)
+
+      stub_const('Gitlab::Database::LoadBalancing::RackMiddleware::Session',
+                 session)
+
+      expect(session).to receive(:clear_session)
+
+      middleware.clear
+    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
+    context 'using a Warden request' do
+      it 'returns the warden user if present' do
+        env = { 'warden' => warden_user }
+
+        expect(middleware.sticking_namespaces_and_ids(env)).to eq([[:user, 42]])
+      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([])
+      end
+    end
+
+    context 'using a request with a manually set sticking object' 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']
+        ])
+      end
+    end
+
+    context 'using a regular request' do
+      it 'returns an empty Array' do
+        expect(middleware.sticking_namespaces_and_ids({})).to eq([])
+      end
+    end
+  end
+end
diff --git a/spec/lib/gitlab/database/load_balancing/resolver_spec.rb b/spec/lib/gitlab/database/load_balancing/resolver_spec.rb
new file mode 100644
index 00000000000..0051cf50255
--- /dev/null
+++ b/spec/lib/gitlab/database/load_balancing/resolver_spec.rb
@@ -0,0 +1,80 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::LoadBalancing::Resolver do
+  describe '#resolve' do
+    let(:ip_addr) { IPAddr.new('127.0.0.2') }
+
+    context 'when nameserver is an IP' do
+      it 'returns an IPAddr object' do
+        service = described_class.new('127.0.0.2')
+
+        expect(service.resolve).to eq(ip_addr)
+      end
+    end
+
+    context 'when nameserver is not an IP' do
+      subject { described_class.new('localhost').resolve }
+
+      it 'looks the nameserver up in the hosts file' do
+        allow_next_instance_of(Resolv::Hosts) do |instance|
+          allow(instance).to receive(:getaddress).with('localhost').and_return('127.0.0.2')
+        end
+
+        expect(subject).to eq(ip_addr)
+      end
+
+      context 'when nameserver is not in the hosts file' do
+        it 'looks the nameserver up in DNS' do
+          resource = double(:resource, address: ip_addr)
+          packet = double(:packet, answer: [resource])
+
+          allow_next_instance_of(Resolv::Hosts) do |instance|
+            allow(instance).to receive(:getaddress).with('localhost').and_raise(Resolv::ResolvError)
+          end
+
+          allow(Net::DNS::Resolver).to receive(:start)
+            .with('localhost', Net::DNS::A)
+            .and_return(packet)
+
+          expect(subject).to eq(ip_addr)
+        end
+
+        context 'when nameserver is not in DNS' do
+          it 'raises an exception' do
+            allow_next_instance_of(Resolv::Hosts) do |instance|
+              allow(instance).to receive(:getaddress).with('localhost').and_raise(Resolv::ResolvError)
+            end
+
+            allow(Net::DNS::Resolver).to receive(:start)
+              .with('localhost', Net::DNS::A)
+              .and_return(double(:packet, answer: []))
+
+            expect { subject }.to raise_exception(
+              described_class::UnresolvableNameserverError,
+              'could not resolve localhost'
+            )
+          end
+        end
+
+        context 'when DNS does not respond' do
+          it 'raises an exception' do
+            allow_next_instance_of(Resolv::Hosts) do |instance|
+              allow(instance).to receive(:getaddress).with('localhost').and_raise(Resolv::ResolvError)
+            end
+
+            allow(Net::DNS::Resolver).to receive(:start)
+              .with('localhost', Net::DNS::A)
+              .and_raise(Net::DNS::Resolver::NoResponseError)
+
+            expect { subject }.to raise_exception(
+              described_class::UnresolvableNameserverError,
+              'no response from DNS server(s)'
+            )
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb b/spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb
new file mode 100644
index 00000000000..7fc7b5e8d11
--- /dev/null
+++ b/spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb
@@ -0,0 +1,252 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
+  let(:service) do
+    described_class.new(nameserver: 'localhost', port: 8600, record: 'foo')
+  end
+
+  before do
+    resource = double(:resource, address: IPAddr.new('127.0.0.1'))
+    packet = double(:packet, answer: [resource])
+
+    allow(Net::DNS::Resolver).to receive(:start)
+      .with('localhost', Net::DNS::A)
+      .and_return(packet)
+  end
+
+  describe '#initialize' do
+    describe ':record_type' do
+      subject { described_class.new(nameserver: 'localhost', port: 8600, record: 'foo', record_type: record_type) }
+
+      context 'with a supported type' do
+        let(:record_type) { 'SRV' }
+
+        it { expect(subject.record_type).to eq Net::DNS::SRV }
+      end
+
+      context 'with an unsupported type' do
+        let(:record_type) { 'AAAA' }
+
+        it 'raises an argument error' do
+          expect { subject }.to raise_error(ArgumentError, 'Unsupported record type: AAAA')
+        end
+      end
+    end
+  end
+
+  describe '#start' do
+    before do
+      allow(service)
+        .to receive(:loop)
+        .and_yield
+    end
+
+    it 'starts service discovery in a new thread' do
+      expect(service)
+        .to receive(:refresh_if_necessary)
+        .and_return(5)
+
+      expect(service)
+        .to receive(:rand)
+        .and_return(2)
+
+      expect(service)
+        .to receive(:sleep)
+        .with(7)
+
+      service.start.join
+    end
+
+    it 'reports exceptions to Sentry' do
+      error = StandardError.new
+
+      expect(service)
+        .to receive(:refresh_if_necessary)
+        .and_raise(error)
+
+      expect(Gitlab::ErrorTracking)
+        .to receive(:track_exception)
+        .with(error)
+
+      expect(service)
+        .to receive(:rand)
+        .and_return(2)
+
+      expect(service)
+        .to receive(:sleep)
+        .with(62)
+
+      service.start.join
+    end
+  end
+
+  describe '#refresh_if_necessary' do
+    let(:address_foo) { described_class::Address.new('foo') }
+    let(:address_bar) { described_class::Address.new('bar') }
+
+    context 'when a refresh is necessary' do
+      before do
+        allow(service)
+          .to receive(:addresses_from_load_balancer)
+          .and_return(%w[localhost])
+
+        allow(service)
+          .to receive(:addresses_from_dns)
+          .and_return([10, [address_foo, address_bar]])
+      end
+
+      it 'refreshes the load balancer hosts' do
+        expect(service)
+          .to receive(:replace_hosts)
+          .with([address_foo, address_bar])
+
+        expect(service.refresh_if_necessary).to eq(10)
+      end
+    end
+
+    context 'when a refresh is not necessary' do
+      before do
+        allow(service)
+          .to receive(:addresses_from_load_balancer)
+          .and_return(%w[localhost])
+
+        allow(service)
+          .to receive(:addresses_from_dns)
+          .and_return([10, %w[localhost]])
+      end
+
+      it 'does not refresh the load balancer hosts' do
+        expect(service)
+          .not_to receive(:replace_hosts)
+
+        expect(service.refresh_if_necessary).to eq(10)
+      end
+    end
+  end
+
+  describe '#replace_hosts' do
+    let(:address_foo) { described_class::Address.new('foo') }
+    let(:address_bar) { described_class::Address.new('bar') }
+
+    let(:load_balancer) do
+      Gitlab::Database::LoadBalancing::LoadBalancer.new([address_foo])
+    end
+
+    before do
+      allow(service)
+        .to receive(:load_balancer)
+        .and_return(load_balancer)
+    end
+
+    it 'replaces the hosts of the load balancer' do
+      service.replace_hosts([address_bar])
+
+      expect(load_balancer.host_list.host_names_and_ports).to eq([['bar', nil]])
+    end
+
+    it 'disconnects the old connections' do
+      host = load_balancer.host_list.hosts.first
+
+      allow(service)
+        .to receive(:disconnect_timeout)
+        .and_return(2)
+
+      expect(host)
+        .to receive(:disconnect!)
+        .with(2)
+
+      service.replace_hosts([address_bar])
+    end
+  end
+
+  describe '#addresses_from_dns' do
+    let(:service) { described_class.new(nameserver: 'localhost', port: 8600, record: 'foo', record_type: record_type) }
+    let(:packet) { double(:packet, answer: [res1, res2]) }
+
+    before do
+      allow(service.resolver)
+        .to receive(:search)
+        .with('foo', described_class::RECORD_TYPES[record_type])
+        .and_return(packet)
+    end
+
+    context 'with an A record' do
+      let(:record_type) { 'A' }
+
+      let(:res1) { double(:resource, address: IPAddr.new('255.255.255.0'), ttl: 90) }
+      let(:res2) { double(:resource, address: IPAddr.new('127.0.0.1'), ttl: 90) }
+
+      it 'returns a TTL and ordered list of IP addresses' do
+        addresses = [
+          described_class::Address.new('127.0.0.1'),
+          described_class::Address.new('255.255.255.0')
+        ]
+
+        expect(service.addresses_from_dns).to eq([90, addresses])
+      end
+    end
+
+    context 'with an SRV record' do
+      let(:record_type) { 'SRV' }
+
+      let(:res1) { double(:resource, host: 'foo1.service.consul.', port: 5432, weight: 1, priority: 1, ttl: 90) }
+      let(:res2) { double(:resource, host: 'foo2.service.consul.', port: 5433, weight: 1, priority: 1, ttl: 90) }
+      let(:res3) { double(:resource, host: 'foo3.service.consul.', port: 5434, weight: 1, priority: 1, ttl: 90) }
+      let(:packet) { double(:packet, answer: [res1, res2, res3], additional: []) }
+
+      before do
+        expect_next_instance_of(Gitlab::Database::LoadBalancing::SrvResolver) do |resolver|
+          allow(resolver).to receive(:address_for).with('foo1.service.consul.').and_return(IPAddr.new('255.255.255.0'))
+          allow(resolver).to receive(:address_for).with('foo2.service.consul.').and_return(IPAddr.new('127.0.0.1'))
+          allow(resolver).to receive(:address_for).with('foo3.service.consul.').and_return(nil)
+        end
+      end
+
+      it 'returns a TTL and ordered list of hosts' do
+        addresses = [
+          described_class::Address.new('127.0.0.1', 5433),
+          described_class::Address.new('255.255.255.0', 5432)
+        ]
+
+        expect(service.addresses_from_dns).to eq([90, addresses])
+      end
+    end
+  end
+
+  describe '#new_wait_time_for' do
+    it 'returns the DNS TTL if greater than the default interval' do
+      res = double(:resource, ttl: 90)
+
+      expect(service.new_wait_time_for([res])).to eq(90)
+    end
+
+    it 'returns the default interval if greater than the DNS TTL' do
+      res = double(:resource, ttl: 10)
+
+      expect(service.new_wait_time_for([res])).to eq(60)
+    end
+
+    it 'returns the default interval if no resources are given' do
+      expect(service.new_wait_time_for([])).to eq(60)
+    end
+  end
+
+  describe '#addresses_from_load_balancer' do
+    it 'returns the ordered host names of the load balancer' do
+      load_balancer = Gitlab::Database::LoadBalancing::LoadBalancer.new(%w[b a])
+
+      allow(service)
+        .to receive(:load_balancer)
+        .and_return(load_balancer)
+
+      addresses = [
+        described_class::Address.new('a'),
+        described_class::Address.new('b')
+      ]
+
+      expect(service.addresses_from_load_balancer).to eq(addresses)
+    end
+  end
+end
diff --git a/spec/lib/gitlab/database/load_balancing/session_spec.rb b/spec/lib/gitlab/database/load_balancing/session_spec.rb
new file mode 100644
index 00000000000..74512f76fd4
--- /dev/null
+++ b/spec/lib/gitlab/database/load_balancing/session_spec.rb
@@ -0,0 +1,353 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::LoadBalancing::Session do
+  after do
+    described_class.clear_session
+  end
+
+  describe '.current' do
+    it 'returns the current session' do
+      expect(described_class.current).to be_an_instance_of(described_class)
+    end
+  end
+
+  describe '.clear_session' do
+    it 'clears the current session' do
+      described_class.current
+      described_class.clear_session
+
+      expect(RequestStore[described_class::CACHE_KEY]).to be_nil
+    end
+  end
+
+  describe '.without_sticky_writes' do
+    it 'ignores sticky write events sent by a connection proxy' do
+      described_class.without_sticky_writes do
+        described_class.current.write!
+      end
+
+      session = described_class.current
+
+      expect(session).not_to be_using_primary
+    end
+
+    it 'still is aware of write that happened' do
+      described_class.without_sticky_writes do
+        described_class.current.write!
+      end
+
+      session = described_class.current
+
+      expect(session.performed_write?).to be true
+    end
+  end
+
+  describe '#use_primary?' do
+    it 'returns true when the primary should be used' do
+      instance = described_class.new
+
+      instance.use_primary!
+
+      expect(instance.use_primary?).to eq(true)
+    end
+
+    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 '#use_primary' do
+    let(:instance) { described_class.new }
+
+    context 'when primary was used before' do
+      before do
+        instance.write!
+      end
+
+      it 'restores state after use' do
+        expect { |blk| instance.use_primary(&blk) }.to yield_with_no_args
+
+        expect(instance.use_primary?).to eq(true)
+      end
+    end
+
+    context 'when primary was not used' do
+      it 'restores state after use' do
+        expect { |blk| instance.use_primary(&blk) }.to yield_with_no_args
+
+        expect(instance.use_primary?).to eq(false)
+      end
+    end
+
+    it 'uses primary during block' do
+      expect do |blk|
+        instance.use_primary do
+          expect(instance.use_primary?).to eq(true)
+
+          # call yield probe
+          blk.to_proc.call
+        end
+      end.to yield_control
+    end
+
+    it 'continues using primary when write was performed' do
+      instance.use_primary do
+        instance.write!
+      end
+
+      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
+
+  describe '#ignore_writes' do
+    it 'ignores write events' do
+      instance = described_class.new
+
+      instance.ignore_writes { instance.write! }
+
+      expect(instance).not_to be_using_primary
+      expect(instance.performed_write?).to eq true
+    end
+
+    it 'does not prevent using primary if an exception is raised' do
+      instance = described_class.new
+
+      instance.ignore_writes { raise ArgumentError } rescue ArgumentError
+      instance.write!
+
+      expect(instance).to be_using_primary
+    end
+  end
+
+  describe '#use_replicas_for_read_queries' do
+    let(:instance) { described_class.new }
+
+    it 'sets the flag inside the block' do
+      expect do |blk|
+        instance.use_replicas_for_read_queries do
+          expect(instance.use_replicas_for_read_queries?).to eq(true)
+
+          # call yield probe
+          blk.to_proc.call
+        end
+      end.to yield_control
+
+      expect(instance.use_replicas_for_read_queries?).to eq(false)
+    end
+
+    it 'restores state after use' do
+      expect do |blk|
+        instance.use_replicas_for_read_queries do
+          instance.use_replicas_for_read_queries do
+            expect(instance.use_replicas_for_read_queries?).to eq(true)
+
+            # call yield probe
+            blk.to_proc.call
+          end
+
+          expect(instance.use_replicas_for_read_queries?).to eq(true)
+        end
+      end.to yield_control
+
+      expect(instance.use_replicas_for_read_queries?).to eq(false)
+    end
+
+    context 'when primary was used before' do
+      before do
+        instance.use_primary!
+      end
+
+      it 'sets the flag inside the block' do
+        expect do |blk|
+          instance.use_replicas_for_read_queries do
+            expect(instance.use_replicas_for_read_queries?).to eq(true)
+
+            # call yield probe
+            blk.to_proc.call
+          end
+        end.to yield_control
+
+        expect(instance.use_replicas_for_read_queries?).to eq(false)
+      end
+    end
+
+    context 'when a write query is performed before' do
+      before do
+        instance.write!
+      end
+
+      it 'sets the flag inside the block' do
+        expect do |blk|
+          instance.use_replicas_for_read_queries do
+            expect(instance.use_replicas_for_read_queries?).to eq(true)
+
+            # call yield probe
+            blk.to_proc.call
+          end
+        end.to yield_control
+
+        expect(instance.use_replicas_for_read_queries?).to eq(false)
+      end
+    end
+  end
+
+  describe '#fallback_to_replicas_for_ambiguous_queries' do
+    let(:instance) { described_class.new }
+
+    it 'sets the flag inside the block' do
+      expect do |blk|
+        instance.fallback_to_replicas_for_ambiguous_queries do
+          expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(true)
+
+          # call yield probe
+          blk.to_proc.call
+        end
+      end.to yield_control
+
+      expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false)
+    end
+
+    it 'restores state after use' do
+      expect do |blk|
+        instance.fallback_to_replicas_for_ambiguous_queries do
+          instance.fallback_to_replicas_for_ambiguous_queries do
+            expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(true)
+
+            # call yield probe
+            blk.to_proc.call
+          end
+
+          expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(true)
+        end
+      end.to yield_control
+
+      expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false)
+    end
+
+    context 'when primary was used before' do
+      before do
+        instance.use_primary!
+      end
+
+      it 'uses primary during block' do
+        expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false)
+
+        expect do |blk|
+          instance.fallback_to_replicas_for_ambiguous_queries do
+            expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false)
+
+            # call yield probe
+            blk.to_proc.call
+          end
+        end.to yield_control
+
+        expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false)
+      end
+    end
+
+    context 'when a write was performed before' do
+      before do
+        instance.write!
+      end
+
+      it 'uses primary during block' do
+        expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false)
+
+        expect do |blk|
+          instance.fallback_to_replicas_for_ambiguous_queries do
+            expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false)
+
+            # call yield probe
+            blk.to_proc.call
+          end
+        end.to yield_control
+
+        expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false)
+      end
+    end
+
+    context 'when primary was used inside the block' do
+      it 'uses primary aterward' do
+        expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false)
+
+        instance.fallback_to_replicas_for_ambiguous_queries do
+          expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(true)
+
+          instance.use_primary!
+
+          expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false)
+        end
+
+        expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false)
+      end
+
+      it 'restores state after use' do
+        instance.fallback_to_replicas_for_ambiguous_queries do
+          instance.fallback_to_replicas_for_ambiguous_queries do
+            expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(true)
+
+            instance.use_primary!
+
+            expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false)
+          end
+
+          expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false)
+        end
+
+        expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false)
+      end
+    end
+
+    context 'when a write was performed inside the block' do
+      it 'uses primary aterward' do
+        expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false)
+
+        instance.fallback_to_replicas_for_ambiguous_queries do
+          expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(true)
+
+          instance.write!
+
+          expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false)
+        end
+
+        expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false)
+      end
+
+      it 'restores state after use' do
+        instance.fallback_to_replicas_for_ambiguous_queries do
+          instance.fallback_to_replicas_for_ambiguous_queries do
+            expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(true)
+
+            instance.write!
+
+            expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false)
+          end
+
+          expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false)
+        end
+
+        expect(instance.fallback_to_replicas_for_ambiguous_queries?).to eq(false)
+      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
new file mode 100644
index 00000000000..90051172fca
--- /dev/null
+++ b/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb
@@ -0,0 +1,161 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
+  let(:middleware) { described_class.new }
+
+  after do
+    Gitlab::Database::LoadBalancing::Session.clear_session
+  end
+
+  describe '#call' do
+    shared_context 'data consistency worker class' do |data_consistency, feature_flag|
+      let(:worker_class) do
+        Class.new do
+          def self.name
+            'TestDataConsistencyWorker'
+          end
+
+          include ApplicationWorker
+
+          data_consistency data_consistency, feature_flag: feature_flag
+
+          def perform(*args)
+          end
+        end
+      end
+
+      before do
+        stub_const('TestDataConsistencyWorker', worker_class)
+      end
+    end
+
+    shared_examples_for 'does not pass database locations' do
+      it 'does not pass database locations', :aggregate_failures do
+        middleware.call(worker_class, job, double(:queue), redis_pool) { 10 }
+
+        expect(job['database_replica_location']).to be_nil
+        expect(job['database_write_location']).to be_nil
+      end
+    end
+
+    shared_examples_for 'mark data consistency location' do |data_consistency|
+      include_context 'data consistency worker class', data_consistency, :load_balancing_for_test_data_consistency_worker
+
+      let(:location) { '0/D525E3A8' }
+
+      context 'when feature flag load_balancing_for_sidekiq is disabled' do
+        before do
+          stub_feature_flags(load_balancing_for_test_data_consistency_worker: false)
+        end
+
+        include_examples 'does not pass database locations'
+      end
+
+      context 'when write was not performed' do
+        before do
+          allow(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_primary?).and_return(false)
+        end
+
+        it 'passes database_replica_location' do
+          expect(middleware).to receive_message_chain(:load_balancer, :host, "database_replica_location").and_return(location)
+
+          middleware.call(worker_class, job, double(:queue), redis_pool) { 10 }
+
+          expect(job['database_replica_location']).to eq(location)
+        end
+      end
+
+      context 'when write was performed' do
+        before do
+          allow(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_primary?).and_return(true)
+        end
+
+        it 'passes primary write location', :aggregate_failures do
+          expect(middleware).to receive_message_chain(:load_balancer, :primary_write_location).and_return(location)
+
+          middleware.call(worker_class, job, double(:queue), redis_pool) { 10 }
+
+          expect(job['database_write_location']).to eq(location)
+        end
+      end
+    end
+
+    shared_examples_for 'database location was already provided' do |provided_database_location, other_location|
+      shared_examples_for 'does not set database location again' do |use_primary|
+        before do
+          allow(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_primary?).and_return(use_primary)
+        end
+
+        it 'does not set database locations again' do
+          middleware.call(worker_class, job, double(:queue), redis_pool) { 10 }
+
+          expect(job[provided_database_location]).to eq(old_location)
+          expect(job[other_location]).to be_nil
+        end
+      end
+
+      let(:old_location) { '0/D525E3A8' }
+      let(:new_location) { 'AB/12345' }
+      let(:job) { { "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", provided_database_location => old_location } }
+
+      before do
+        allow(middleware).to receive_message_chain(:load_balancer, :primary_write_location).and_return(new_location)
+        allow(middleware).to receive_message_chain(:load_balancer, :database_replica_location).and_return(new_location)
+      end
+
+      context "when write was performed" do
+        include_examples 'does not set database location again', true
+      end
+
+      context "when write was not performed" do
+        include_examples 'does not set database location again', false
+      end
+    end
+
+    let(:queue) { 'default' }
+    let(:redis_pool) { Sidekiq.redis_pool }
+    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
+    end
+
+    context 'when worker cannot be constantized' do
+      let(:worker_class) { 'ActionMailer::MailDeliveryJob' }
+
+      include_examples 'does not pass database locations'
+    end
+
+    context 'when worker class does not include ApplicationWorker' do
+      let(:worker_class) { ActiveJob::QueueAdapters::SidekiqAdapter::JobWrapper }
+
+      include_examples 'does not pass database locations'
+    end
+
+    context 'database write location was already provided' do
+      include_examples 'database location was already provided', 'database_write_location', 'database_replica_location'
+    end
+
+    context 'database replica location was already provided' do
+      include_examples 'database location was already provided', 'database_replica_location', 'database_write_location'
+    end
+
+    context 'when worker data consistency is :always' do
+      include_context 'data consistency worker class', :always, :load_balancing_for_test_data_consistency_worker
+
+      include_examples 'does not pass database locations'
+    end
+
+    context 'when worker data consistency is :delayed' do
+      include_examples  'mark data consistency location', :delayed
+    end
+
+    context 'when worker data consistency is :sticky' do
+      include_examples  'mark data consistency location', :sticky
+    end
+  end
+end
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
new file mode 100644
index 00000000000..b7cd0caa922
--- /dev/null
+++ b/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb
@@ -0,0 +1,201 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware do
+  let(:middleware) { described_class.new }
+
+  after do
+    Gitlab::Database::LoadBalancing::Session.clear_session
+  end
+
+  describe '#call' do
+    shared_context 'data consistency worker class' do |data_consistency, feature_flag|
+      let(:worker_class) do
+        Class.new do
+          def self.name
+            'TestDataConsistencyWorker'
+          end
+
+          include ApplicationWorker
+
+          data_consistency data_consistency, feature_flag: feature_flag
+
+          def perform(*args)
+          end
+        end
+      end
+
+      before do
+        stub_const('TestDataConsistencyWorker', worker_class)
+      end
+    end
+
+    shared_examples_for 'stick to the primary' do
+      it 'sticks to the primary' do
+        middleware.call(worker, job, double(:queue)) do
+          expect(Gitlab::Database::LoadBalancing::Session.current.use_primary?).to be_truthy
+        end
+      end
+    end
+
+    shared_examples_for 'replica is up to date' do |location, data_consistency|
+      it 'does not stick to the primary', :aggregate_failures do
+        expect(middleware).to receive(:replica_caught_up?).with(location).and_return(true)
+
+        middleware.call(worker, job, double(:queue)) do
+          expect(Gitlab::Database::LoadBalancing::Session.current.use_primary?).not_to be_truthy
+        end
+
+        expect(job[:database_chosen]).to eq('replica')
+      end
+
+      it "updates job hash with data_consistency :#{data_consistency}" do
+        middleware.call(worker, job, double(:queue)) do
+          expect(job).to include(data_consistency: data_consistency.to_s)
+        end
+      end
+    end
+
+    shared_examples_for 'sticks based on data consistency' do |data_consistency|
+      include_context 'data consistency worker class', data_consistency, :load_balancing_for_test_data_consistency_worker
+
+      context 'when load_balancing_for_test_data_consistency_worker is disabled' do
+        before do
+          stub_feature_flags(load_balancing_for_test_data_consistency_worker: false)
+        end
+
+        include_examples 'stick to the primary'
+      end
+
+      context 'when database replica location is set' do
+        let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e', 'database_replica_location' => '0/D525E3A8' } }
+
+        before do
+          allow(middleware).to receive(:replica_caught_up?).and_return(true)
+        end
+
+        it_behaves_like 'replica is up to date', '0/D525E3A8', data_consistency
+      end
+
+      context 'when database primary location is set' do
+        let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e', 'database_write_location' => '0/D525E3A8' } }
+
+        before do
+          allow(middleware).to receive(:replica_caught_up?).and_return(true)
+        end
+
+        it_behaves_like 'replica is up to date', '0/D525E3A8', data_consistency
+      end
+
+      context 'when database location is not set' do
+        let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e' } }
+
+        it_behaves_like 'stick to the primary', nil
+      end
+    end
+
+    let(:queue) { 'default' }
+    let(:redis_pool) { Sidekiq.redis_pool }
+    let(:worker) { worker_class.new }
+    let(:job) { { "retry" => 3, "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", 'database_replica_location' => '0/D525E3A8' } }
+    let(:block) { 10 }
+
+    before do
+      skip_feature_flags_yaml_validation
+      skip_default_enabled_yaml_check
+      allow(middleware).to receive(:clear)
+      allow(Gitlab::Database::LoadBalancing::Session.current).to receive(:performed_write?).and_return(true)
+    end
+
+    context 'when worker class does not include ApplicationWorker' do
+      let(:worker) { ActiveJob::QueueAdapters::SidekiqAdapter::JobWrapper.new }
+
+      include_examples 'stick to the primary'
+    end
+
+    context 'when worker data consistency is :always' do
+      include_context 'data consistency worker class', :always, :load_balancing_for_test_data_consistency_worker
+
+      include_examples 'stick to the primary'
+    end
+
+    context 'when worker data consistency is :delayed' do
+      include_examples 'sticks based on data consistency', :delayed
+
+      context 'when replica is not up to date' do
+        before do
+          allow(::Gitlab::Database::LoadBalancing).to receive_message_chain(:proxy, :load_balancer, :release_host)
+          allow(::Gitlab::Database::LoadBalancing).to receive_message_chain(:proxy, :load_balancer, :select_up_to_date_host).and_return(false)
+        end
+
+        around do |example|
+          with_sidekiq_server_middleware do |chain|
+            chain.add described_class
+            Sidekiq::Testing.disable! { example.run }
+          end
+        end
+
+        context 'when job is executed first' do
+          it 'raise an error and retries', :aggregate_failures do
+            expect do
+              process_job(job)
+            end.to raise_error(Sidekiq::JobRetry::Skip)
+
+            expect(job['error_class']).to eq('Gitlab::Database::LoadBalancing::SidekiqServerMiddleware::JobReplicaNotUpToDate')
+            expect(job[:database_chosen]).to eq('retry')
+          end
+        end
+
+        context 'when job is retried' do
+          it 'stick to the primary', :aggregate_failures do
+            expect do
+              process_job(job)
+            end.to raise_error(Sidekiq::JobRetry::Skip)
+
+            process_job(job)
+            expect(job[:database_chosen]).to eq('primary')
+          end
+        end
+
+        context 'replica selection mechanism feature flag rollout' do
+          before do
+            stub_feature_flags(sidekiq_load_balancing_rotate_up_to_date_replica: false)
+          end
+
+          it 'uses different implmentation' do
+            expect(::Gitlab::Database::LoadBalancing).to receive_message_chain(:proxy, :load_balancer, :host, :caught_up?).and_return(false)
+
+            expect do
+              process_job(job)
+            end.to raise_error(Sidekiq::JobRetry::Skip)
+          end
+        end
+      end
+    end
+
+    context 'when worker data consistency is :sticky' do
+      include_examples 'sticks based on data consistency', :sticky
+
+      context 'when replica is not up to date' do
+        before do
+          allow(middleware).to receive(:replica_caught_up?).and_return(false)
+        end
+
+        include_examples 'stick to the primary'
+
+        it 'updates job hash with primary database chosen', :aggregate_failures do
+          expect { |b| middleware.call(worker, job, double(:queue), &b) }.to yield_control
+
+          expect(job[:database_chosen]).to eq('primary')
+        end
+      end
+    end
+  end
+
+  def process_job(job)
+    Sidekiq::JobRetry.new.local(worker_class, job, queue) do
+      worker_class.process_job(job)
+    end
+  end
+end
diff --git a/spec/lib/gitlab/database/load_balancing/srv_resolver_spec.rb b/spec/lib/gitlab/database/load_balancing/srv_resolver_spec.rb
new file mode 100644
index 00000000000..6ac0608d485
--- /dev/null
+++ b/spec/lib/gitlab/database/load_balancing/srv_resolver_spec.rb
@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::LoadBalancing::SrvResolver do
+  let(:resolver) { Net::DNS::Resolver.new(nameservers: '127.0.0.1', port: 8600, use_tcp: true) }
+  let(:additional) { dns_response_packet_from_fixture('srv_with_a_rr_in_additional_section').additional }
+
+  describe '#address_for' do
+    let(:host) { 'patroni-02-db-gstg.node.east-us-2.consul.' }
+
+    subject { described_class.new(resolver, additional).address_for(host) }
+
+    context 'when additional section contains an A record' do
+      it 'returns an IP4 address' do
+        expect(subject).to eq(IPAddr.new('10.224.29.102'))
+      end
+    end
+
+    context 'when additional section contains an AAAA record' do
+      let(:host) { 'a.gtld-servers.net.' }
+      let(:additional) { dns_response_packet_from_fixture('a_with_aaaa_rr_in_additional_section').additional }
+
+      it 'returns an IP6 address' do
+        expect(subject).to eq(IPAddr.new('2001:503:a83e::2:30'))
+      end
+    end
+
+    context 'when additional section does not contain A nor AAAA records' do
+      let(:additional) { [] }
+
+      context 'when host resolves to an A record' do
+        before do
+          allow(resolver).to receive(:search).with(host, Net::DNS::ANY).and_return(dns_response_packet_from_fixture('a_rr'))
+        end
+
+        it 'returns an IP4 address' do
+          expect(subject).to eq(IPAddr.new('10.224.29.102'))
+        end
+      end
+
+      context 'when host does resolves to an AAAA record' do
+        before do
+          allow(resolver).to receive(:search).with(host, Net::DNS::ANY).and_return(dns_response_packet_from_fixture('aaaa_rr'))
+        end
+
+        it 'returns an IP6 address' do
+          expect(subject).to eq(IPAddr.new('2a00:1450:400e:80a::200e'))
+        end
+      end
+    end
+  end
+
+  def dns_response_packet_from_fixture(fixture_name)
+    fixture         = File.read(Rails.root + "spec/fixtures/dns/#{fixture_name}.json")
+    encoded_payload = Gitlab::Json.parse(fixture)['payload']
+    payload         = Base64.decode64(encoded_payload)
+
+    Net::DNS::Packet.parse(payload)
+  end
+end
diff --git a/spec/lib/gitlab/database/load_balancing/sticking_spec.rb b/spec/lib/gitlab/database/load_balancing/sticking_spec.rb
new file mode 100644
index 00000000000..bf4e3756e0e
--- /dev/null
+++ b/spec/lib/gitlab/database/load_balancing/sticking_spec.rb
@@ -0,0 +1,307 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.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 '.all_caught_up?' do
+    let(:lb) { double(:lb) }
+
+    before do
+      allow(described_class).to receive(:load_balancer).and_return(lb)
+    end
+
+    it 'returns true 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?)
+
+      expect(described_class.all_caught_up?(:user, 42)).to eq(true)
+    end
+
+    it 'returns true, and 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)
+
+      expect(described_class.all_caught_up?(:user, 42)).to eq(true)
+    end
+
+    it 'return false 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(described_class.all_caught_up?(:user, 42)).to eq(false)
+    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
+
+  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
+    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')
+
+        allow(described_class).to receive(:load_balancer).and_return(lb)
+      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!)
+
+        subject
+      end
+    end
+  end
+
+  describe '.stick' do
+    it_behaves_like 'sticking' do
+      let(:ids) { [42] }
+      subject { described_class.stick(:user, ids.first) }
+    end
+  end
+
+  describe '.bulk_stick' do
+    it_behaves_like 'sticking' do
+      let(:ids) { [42, 43] }
+      subject { described_class.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')
+
+        allow(described_class).to receive(:load_balancer).and_return(lb)
+
+        expect(described_class).to receive(:set_write_location_for)
+          .with(:user, 42, 'foo')
+
+        described_class.mark_primary_write_location(:user, 42)
+      end
+    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)
+
+        described_class.mark_primary_write_location(:user, 42)
+      end
+
+      context 'when write location is nil' do
+        before do
+          allow(Gitlab::Database).to receive(:get_write_location).and_return(nil)
+        end
+
+        it 'does not update the write location' do
+          expect(described_class).not_to receive(:set_write_location_for)
+
+          described_class.mark_primary_write_location(:user, 42)
+        end
+      end
+    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)
+      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
+    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
+
+  describe '.select_caught_up_replicas' do
+    let(:lb) { double(:lb) }
+
+    before do
+      allow(described_class).to receive(:load_balancer).and_return(lb)
+    end
+
+    context 'with no write location' do
+      before do
+        allow(described_class).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(described_class).not_to receive(:select_caught_up_hosts)
+        expect(described_class.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')
+      end
+
+      it 'returns true, selects hosts, and unsticks if any secondary has caught up' do
+        expect(lb).to receive(:select_caught_up_hosts).and_return(true)
+        expect(described_class).to receive(:unstick).with(:project, 42)
+        expect(described_class.select_caught_up_replicas(:project, 42)).to be true
+      end
+    end
+  end
+end
diff --git a/spec/lib/gitlab/database/load_balancing_spec.rb b/spec/lib/gitlab/database/load_balancing_spec.rb
new file mode 100644
index 00000000000..e7de7f2b43b
--- /dev/null
+++ b/spec/lib/gitlab/database/load_balancing_spec.rb
@@ -0,0 +1,834 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::LoadBalancing do
+  include_context 'clear DB Load Balancing configuration'
+
+  before do
+    stub_env('ENABLE_LOAD_BALANCING_FOR_FOSS', 'true')
+  end
+
+  describe '.proxy' do
+    context 'when configured' do
+      before do
+        allow(ActiveRecord::Base.singleton_class).to receive(:prepend)
+        subject.configure_proxy
+      end
+
+      it 'returns the connection proxy' do
+        expect(subject.proxy).to be_an_instance_of(subject::ConnectionProxy)
+      end
+    end
+
+    context 'when not configured' do
+      it 'returns nil' do
+        expect(subject.proxy).to be_nil
+      end
+
+      it 'tracks an error to sentry' do
+        expect(Gitlab::ErrorTracking).to receive(:track_exception).with(
+          an_instance_of(subject::ProxyNotConfiguredError)
+        )
+
+        subject.proxy
+      end
+    end
+  end
+
+  describe '.configuration' do
+    it 'returns a Hash' do
+      lb_config = { 'hosts' => %w(foo) }
+
+      original_db_config = Gitlab::Database.config
+      modified_db_config = original_db_config.merge(load_balancing: lb_config)
+      expect(Gitlab::Database).to receive(:config).and_return(modified_db_config)
+
+      expect(described_class.configuration).to eq(lb_config)
+    end
+  end
+
+  describe '.max_replication_difference' do
+    context 'without an explicitly configured value' do
+      it 'returns the default value' do
+        allow(described_class)
+          .to receive(:configuration)
+          .and_return({})
+
+        expect(described_class.max_replication_difference).to eq(8.megabytes)
+      end
+    end
+
+    context 'with an explicitly configured value' do
+      it 'returns the configured value' do
+        allow(described_class)
+          .to receive(:configuration)
+          .and_return({ 'max_replication_difference' => 4 })
+
+        expect(described_class.max_replication_difference).to eq(4)
+      end
+    end
+  end
+
+  describe '.max_replication_lag_time' do
+    context 'without an explicitly configured value' do
+      it 'returns the default value' do
+        allow(described_class)
+          .to receive(:configuration)
+          .and_return({})
+
+        expect(described_class.max_replication_lag_time).to eq(60)
+      end
+    end
+
+    context 'with an explicitly configured value' do
+      it 'returns the configured value' do
+        allow(described_class)
+          .to receive(:configuration)
+          .and_return({ 'max_replication_lag_time' => 4 })
+
+        expect(described_class.max_replication_lag_time).to eq(4)
+      end
+    end
+  end
+
+  describe '.replica_check_interval' do
+    context 'without an explicitly configured value' do
+      it 'returns the default value' do
+        allow(described_class)
+          .to receive(:configuration)
+          .and_return({})
+
+        expect(described_class.replica_check_interval).to eq(60)
+      end
+    end
+
+    context 'with an explicitly configured value' do
+      it 'returns the configured value' do
+        allow(described_class)
+          .to receive(:configuration)
+          .and_return({ 'replica_check_interval' => 4 })
+
+        expect(described_class.replica_check_interval).to eq(4)
+      end
+    end
+  end
+
+  describe '.hosts' do
+    it 'returns a list of hosts' do
+      allow(described_class)
+        .to receive(:configuration)
+        .and_return({ 'hosts' => %w(foo bar baz) })
+
+      expect(described_class.hosts).to eq(%w(foo bar baz))
+    end
+  end
+
+  describe '.pool_size' do
+    it 'returns a Fixnum' do
+      expect(described_class.pool_size).to be_a_kind_of(Integer)
+    end
+  end
+
+  describe '.enable?' do
+    before do
+      clear_load_balancing_configuration
+      allow(described_class).to receive(:hosts).and_return(%w(foo))
+    end
+
+    it 'returns false when no hosts are specified' do
+      allow(described_class).to receive(:hosts).and_return([])
+
+      expect(described_class.enable?).to eq(false)
+    end
+
+    it 'returns false when Sidekiq is being used' do
+      allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(true)
+
+      expect(described_class.enable?).to eq(false)
+    end
+
+    it 'returns false when running inside a Rake task' do
+      allow(Gitlab::Runtime).to receive(:rake?).and_return(true)
+
+      expect(described_class.enable?).to eq(false)
+    end
+
+    it 'returns true when load balancing should be enabled' do
+      allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(false)
+
+      expect(described_class.enable?).to eq(true)
+    end
+
+    it 'returns true when service discovery is enabled' do
+      allow(described_class).to receive(:hosts).and_return([])
+      allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(false)
+
+      allow(described_class)
+        .to receive(:service_discovery_enabled?)
+        .and_return(true)
+
+      expect(described_class.enable?).to eq(true)
+    end
+
+    context 'when ENABLE_LOAD_BALANCING_FOR_SIDEKIQ environment variable is set' do
+      before do
+        stub_env('ENABLE_LOAD_BALANCING_FOR_SIDEKIQ', 'true')
+      end
+
+      it 'returns true when Sidekiq is being used' do
+        allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(true)
+
+        expect(described_class.enable?).to eq(true)
+      end
+    end
+  end
+
+  describe '.configured?' do
+    before do
+      clear_load_balancing_configuration
+    end
+
+    it 'returns true when Sidekiq is being used' do
+      allow(described_class).to receive(:hosts).and_return(%w(foo))
+      allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(true)
+      expect(described_class.configured?).to eq(true)
+    end
+
+    it 'returns true when service discovery is enabled in Sidekiq' do
+      allow(described_class).to receive(:hosts).and_return([])
+      allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(true)
+
+      allow(described_class)
+        .to receive(:service_discovery_enabled?)
+        .and_return(true)
+
+      expect(described_class.configured?).to eq(true)
+    end
+
+    it 'returns false when neither service discovery nor hosts are configured' do
+      allow(described_class).to receive(:hosts).and_return([])
+
+      allow(described_class)
+        .to receive(:service_discovery_enabled?)
+        .and_return(false)
+
+      expect(described_class.configured?).to eq(false)
+    end
+  end
+
+  describe '.configure_proxy' do
+    it 'configures the connection proxy' do
+      allow(ActiveRecord::Base.singleton_class).to receive(:prepend)
+
+      described_class.configure_proxy
+
+      expect(ActiveRecord::Base.singleton_class).to have_received(:prepend)
+        .with(Gitlab::Database::LoadBalancing::ActiveRecordProxy)
+    end
+  end
+
+  describe '.active_record_models' do
+    it 'returns an Array' do
+      expect(described_class.active_record_models).to be_an_instance_of(Array)
+    end
+  end
+
+  describe '.service_discovery_enabled?' do
+    it 'returns true if service discovery is enabled' do
+      allow(described_class)
+        .to receive(:configuration)
+        .and_return('discover' => { 'record' => 'foo' })
+
+      expect(described_class.service_discovery_enabled?).to eq(true)
+    end
+
+    it 'returns false if service discovery is disabled' do
+      expect(described_class.service_discovery_enabled?).to eq(false)
+    end
+  end
+
+  describe '.service_discovery_configuration' do
+    context 'when no configuration is provided' do
+      it 'returns a default configuration Hash' do
+        expect(described_class.service_discovery_configuration).to eq(
+          nameserver: 'localhost',
+          port: 8600,
+          record: nil,
+          record_type: 'A',
+          interval: 60,
+          disconnect_timeout: 120,
+          use_tcp: false
+        )
+      end
+    end
+
+    context 'when configuration is provided' do
+      it 'returns a Hash including the custom configuration' do
+        allow(described_class)
+          .to receive(:configuration)
+          .and_return('discover' => { 'record' => 'foo', 'record_type' => 'SRV' })
+
+        expect(described_class.service_discovery_configuration).to eq(
+          nameserver: 'localhost',
+          port: 8600,
+          record: 'foo',
+          record_type: 'SRV',
+          interval: 60,
+          disconnect_timeout: 120,
+          use_tcp: false
+        )
+      end
+    end
+  end
+
+  describe '.start_service_discovery' do
+    it 'does not start if service discovery is disabled' do
+      expect(Gitlab::Database::LoadBalancing::ServiceDiscovery)
+        .not_to receive(:new)
+
+      described_class.start_service_discovery
+    end
+
+    it 'starts service discovery if enabled' do
+      allow(described_class)
+        .to receive(:service_discovery_enabled?)
+        .and_return(true)
+
+      instance = double(:instance)
+
+      expect(Gitlab::Database::LoadBalancing::ServiceDiscovery)
+        .to receive(:new)
+        .with(an_instance_of(Hash))
+        .and_return(instance)
+
+      expect(instance)
+        .to receive(:start)
+
+      described_class.start_service_discovery
+    end
+  end
+
+  describe '.db_role_for_connection' do
+    let(:connection) { double(:conneciton) }
+
+    context 'when the load balancing is not configured' do
+      before do
+        allow(described_class).to receive(:enable?).and_return(false)
+      end
+
+      it 'returns primary' do
+        expect(described_class.db_role_for_connection(connection)).to be(:primary)
+      end
+    end
+
+    context 'when the load balancing is configured' do
+      let(:proxy) { described_class::ConnectionProxy.new(%w(foo)) }
+      let(:load_balancer) { described_class::LoadBalancer.new(%w(foo)) }
+
+      before do
+        allow(ActiveRecord::Base.singleton_class).to receive(:prepend)
+
+        allow(described_class).to receive(:enable?).and_return(true)
+        allow(described_class).to receive(:proxy).and_return(proxy)
+        allow(proxy).to receive(:load_balancer).and_return(load_balancer)
+
+        subject.configure_proxy(proxy)
+      end
+
+      context 'when the load balancer returns :replica' do
+        it 'returns :replica' do
+          allow(load_balancer).to receive(:db_role_for_connection).and_return(:replica)
+
+          expect(described_class.db_role_for_connection(connection)).to be(:replica)
+
+          expect(load_balancer).to have_received(:db_role_for_connection).with(connection)
+        end
+      end
+
+      context 'when the load balancer returns :primary' do
+        it 'returns :primary' do
+          allow(load_balancer).to receive(:db_role_for_connection).and_return(:primary)
+
+          expect(described_class.db_role_for_connection(connection)).to be(:primary)
+
+          expect(load_balancer).to have_received(:db_role_for_connection).with(connection)
+        end
+      end
+
+      context 'when the load balancer returns nil' do
+        it 'returns nil' do
+          allow(load_balancer).to receive(:db_role_for_connection).and_return(nil)
+
+          expect(described_class.db_role_for_connection(connection)).to be(nil)
+
+          expect(load_balancer).to have_received(:db_role_for_connection).with(connection)
+        end
+      end
+    end
+  end
+
+  # For such an important module like LoadBalancing, full mocking is not
+  # enough. This section implements some integration tests to test a full flow
+  # of the load balancer.
+  # - A real model with a table backed behind is defined
+  # - The load balancing module is set up for this module only, as to prevent
+  # breaking other tests. The replica configuration is cloned from the test
+  # configuraiton.
+  # - In each test, we listen to the SQL queries (via sql.active_record
+  # instrumentation) while triggering real queries from the defined model.
+  # - We assert the desinations (replica/primary) of the queries in order.
+  describe 'LoadBalancing integration tests', :delete do
+    before(:all) do
+      ActiveRecord::Schema.define do
+        create_table :load_balancing_test, force: true do |t|
+          t.string :name, null: true
+        end
+      end
+    end
+
+    after(:all) do
+      ActiveRecord::Schema.define do
+        drop_table :load_balancing_test, force: true
+      end
+    end
+
+    shared_context 'LoadBalancing setup' do
+      let(:development_db_config) { ActiveRecord::Base.configurations.configs_for(env_name: 'development').first.configuration_hash }
+      let(:hosts) { [development_db_config[:host]] }
+      let(:model) do
+        Class.new(ApplicationRecord) do
+          self.table_name = "load_balancing_test"
+        end
+      end
+
+      before do
+        # Preloading testing class
+        model.singleton_class.prepend ::Gitlab::Database::LoadBalancing::ActiveRecordProxy
+
+        # Setup load balancing
+        clear_load_balancing_configuration
+        allow(ActiveRecord::Base.singleton_class).to receive(:prepend)
+        subject.configure_proxy(::Gitlab::Database::LoadBalancing::ConnectionProxy.new(hosts))
+
+        original_db_config = Gitlab::Database.config
+        modified_db_config = original_db_config.merge(load_balancing: { hosts: hosts })
+        allow(Gitlab::Database).to receive(:config).and_return(modified_db_config)
+
+        ::Gitlab::Database::LoadBalancing::Session.clear_session
+      end
+    end
+
+    where(:queries, :include_transaction, :expected_results) do
+      [
+        # Read methods
+        [-> { model.first }, false, [:replica]],
+        [-> { model.find_by(id: 123) }, false, [:replica]],
+        [-> { model.where(name: 'hello').to_a }, false, [:replica]],
+
+        # Write methods
+        [-> { model.create!(name: 'test1') }, false, [:primary]],
+        [
+          -> {
+            instance = model.create!(name: 'test1')
+            instance.update!(name: 'test2')
+          },
+          false, [:primary, :primary]
+        ],
+        [-> { model.update_all(name: 'test2') }, false, [:primary]],
+        [
+          -> {
+            instance = model.create!(name: 'test1')
+            instance.destroy!
+          },
+          false, [:primary, :primary]
+        ],
+        [-> { model.delete_all }, false, [:primary]],
+
+        # Custom query
+        [-> { model.connection.exec_query('SELECT 1').to_a }, false, [:primary]],
+
+        # Reads after a write
+        [
+          -> {
+            model.first
+            model.create!(name: 'test1')
+            model.first
+            model.find_by(name: 'test1')
+          },
+          false, [:replica, :primary, :primary, :primary]
+        ],
+
+        # Inside a transaction
+        [
+          -> {
+            model.transaction do
+              model.find_by(name: 'test1')
+              model.create!(name: 'test1')
+              instance = model.find_by(name: 'test1')
+              instance.update!(name: 'test2')
+            end
+            model.find_by(name: 'test1')
+          },
+          true, [:primary, :primary, :primary, :primary, :primary, :primary, :primary]
+        ],
+
+        # Nested transaction
+        [
+          -> {
+            model.transaction do
+              model.transaction do
+                model.create!(name: 'test1')
+              end
+              model.update_all(name: 'test2')
+            end
+            model.find_by(name: 'test1')
+          },
+          true, [:primary, :primary, :primary, :primary, :primary]
+        ],
+
+        # Read-only transaction
+        [
+          -> {
+            model.transaction do
+              model.first
+              model.where(name: 'test1').to_a
+            end
+          },
+          true, [:primary, :primary, :primary, :primary]
+        ],
+
+        # use_primary
+        [
+          -> {
+            ::Gitlab::Database::LoadBalancing::Session.current.use_primary do
+              model.first
+              model.where(name: 'test1').to_a
+            end
+            model.first
+          },
+          false, [:primary, :primary, :replica]
+        ],
+
+        # use_primary!
+        [
+          -> {
+            model.first
+            ::Gitlab::Database::LoadBalancing::Session.current.use_primary!
+            model.where(name: 'test1').to_a
+          },
+          false, [:replica, :primary]
+        ],
+
+        # use_replicas_for_read_queries does not affect read queries
+        [
+          -> {
+            ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
+              model.where(name: 'test1').to_a
+            end
+          },
+          false, [:replica]
+        ],
+
+        # use_replicas_for_read_queries does not affect write queries
+        [
+          -> {
+            ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
+              model.create!(name: 'test1')
+            end
+          },
+          false, [:primary]
+        ],
+
+        # use_replicas_for_read_queries does not affect ambiguous queries
+        [
+          -> {
+            ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
+              model.connection.exec_query("SELECT 1")
+            end
+          },
+          false, [:primary]
+        ],
+
+        # use_replicas_for_read_queries ignores use_primary! for read queries
+        [
+          -> {
+            ::Gitlab::Database::LoadBalancing::Session.current.use_primary!
+            ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
+              model.where(name: 'test1').to_a
+            end
+          },
+          false, [:replica]
+        ],
+
+        # use_replicas_for_read_queries adheres use_primary! for write queries
+        [
+          -> {
+            ::Gitlab::Database::LoadBalancing::Session.current.use_primary!
+            ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
+              model.create!(name: 'test1')
+            end
+          },
+          false, [:primary]
+        ],
+
+        # use_replicas_for_read_queries adheres use_primary! for ambiguous queries
+        [
+          -> {
+            ::Gitlab::Database::LoadBalancing::Session.current.use_primary!
+            ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
+              model.connection.exec_query('SELECT 1')
+            end
+          },
+          false, [:primary]
+        ],
+
+        # use_replicas_for_read_queries ignores use_primary blocks
+        [
+          -> {
+            ::Gitlab::Database::LoadBalancing::Session.current.use_primary do
+              ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
+                model.where(name: 'test1').to_a
+              end
+            end
+          },
+          false, [:replica]
+        ],
+
+        # use_replicas_for_read_queries ignores a session already performed write
+        [
+          -> {
+            ::Gitlab::Database::LoadBalancing::Session.current.write!
+            ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
+              model.where(name: 'test1').to_a
+            end
+          },
+          false, [:replica]
+        ],
+
+        # fallback_to_replicas_for_ambiguous_queries
+        [
+          -> {
+            ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
+              model.first
+              model.where(name: 'test1').to_a
+            end
+          },
+          false, [:replica, :replica]
+        ],
+
+        # fallback_to_replicas_for_ambiguous_queries for read-only transaction
+        [
+          -> {
+            ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
+              model.transaction do
+                model.first
+                model.where(name: 'test1').to_a
+              end
+            end
+          },
+          false, [:replica, :replica]
+        ],
+
+        # A custom read query inside fallback_to_replicas_for_ambiguous_queries
+        [
+          -> {
+            ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
+              model.connection.exec_query("SELECT 1")
+            end
+          },
+          false, [:replica]
+        ],
+
+        # A custom read query inside a transaction fallback_to_replicas_for_ambiguous_queries
+        [
+          -> {
+            ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
+              model.transaction do
+                model.connection.exec_query("SET LOCAL statement_timeout = 5000")
+                model.count
+              end
+            end
+          },
+          true, [:replica, :replica, :replica, :replica]
+        ],
+
+        # fallback_to_replicas_for_ambiguous_queries after a write
+        [
+          -> {
+            model.create!(name: 'Test1')
+            ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
+              model.connection.exec_query("SELECT 1")
+            end
+          },
+          false, [:primary, :primary]
+        ],
+
+        # fallback_to_replicas_for_ambiguous_queries after use_primary!
+        [
+          -> {
+            ::Gitlab::Database::LoadBalancing::Session.current.use_primary!
+            ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
+              model.connection.exec_query("SELECT 1")
+            end
+          },
+          false, [:primary]
+        ],
+
+        # fallback_to_replicas_for_ambiguous_queries inside use_primary
+        [
+          -> {
+            ::Gitlab::Database::LoadBalancing::Session.current.use_primary do
+              ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
+                model.connection.exec_query("SELECT 1")
+              end
+            end
+          },
+          false, [:primary]
+        ],
+
+        # use_primary inside fallback_to_replicas_for_ambiguous_queries
+        [
+          -> {
+            ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
+              ::Gitlab::Database::LoadBalancing::Session.current.use_primary do
+                model.connection.exec_query("SELECT 1")
+              end
+            end
+          },
+          false, [:primary]
+        ],
+
+        # A write query inside fallback_to_replicas_for_ambiguous_queries
+        [
+          -> {
+            ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
+              model.connection.exec_query("SELECT 1")
+              model.delete_all
+              model.connection.exec_query("SELECT 1")
+            end
+          },
+          false, [:replica, :primary, :primary]
+        ],
+
+        # use_replicas_for_read_queries incorporates with fallback_to_replicas_for_ambiguous_queries
+        [
+          -> {
+            ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
+              ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
+                model.connection.exec_query('SELECT 1')
+                model.where(name: 'test1').to_a
+              end
+            end
+          },
+          false, [:replica, :replica]
+        ]
+      ]
+    end
+
+    with_them do
+      include_context 'LoadBalancing setup'
+
+      it 'redirects queries to the right roles' do
+        roles = []
+
+        subscriber = ActiveSupport::Notifications.subscribe('sql.active_record') do |event|
+          payload = event.payload
+
+          assert =
+            if payload[:name] == 'SCHEMA'
+              false
+            elsif payload[:name] == 'SQL' # Custom query
+              true
+            else
+              keywords = %w[load_balancing_test]
+              keywords += %w[begin commit] if include_transaction
+              keywords.any? { |keyword| payload[:sql].downcase.include?(keyword) }
+            end
+
+          if assert
+            db_role = ::Gitlab::Database::LoadBalancing.db_role_for_connection(payload[:connection])
+            roles << db_role
+          end
+        end
+
+        self.instance_exec(&queries)
+
+        expect(roles).to eql(expected_results)
+      ensure
+        ActiveSupport::Notifications.unsubscribe(subscriber) if subscriber
+      end
+    end
+
+    context 'custom connection handling' do
+      where(:queries, :expected_role) do
+        [
+          # Reload cache. The schema loading queries should be handled by
+          # primary.
+          [
+            -> {
+              model.connection.clear_cache!
+              model.connection.schema_cache.add('users')
+              model.connection.pool.release_connection
+            },
+            :primary
+          ],
+
+          # Call model's connection method
+          [
+            -> {
+              connection = model.connection
+              connection.select_one('SELECT 1')
+              connection.pool.release_connection
+            },
+            :replica
+          ],
+
+          # Retrieve connection via #retrieve_connection
+          [
+            -> {
+              connection = model.retrieve_connection
+              connection.select_one('SELECT 1')
+              connection.pool.release_connection
+            },
+            :primary
+          ]
+        ]
+      end
+
+      with_them do
+        include_context 'LoadBalancing setup'
+
+        it 'redirects queries to the right roles' do
+          roles = []
+
+          subscriber = ActiveSupport::Notifications.subscribe('sql.active_record') do |event|
+            role = ::Gitlab::Database::LoadBalancing.db_role_for_connection(event.payload[:connection])
+            roles << role if role.present?
+          end
+
+          self.instance_exec(&queries)
+
+          expect(roles).to all(eql(expected_role))
+        ensure
+          ActiveSupport::Notifications.unsubscribe(subscriber) if subscriber
+        end
+      end
+    end
+
+    context 'a write inside a transaction inside fallback_to_replicas_for_ambiguous_queries block' do
+      include_context 'LoadBalancing setup'
+
+      it 'raises an exception' do
+        expect do
+          ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
+            model.transaction do
+              model.first
+              model.create!(name: 'hello')
+            end
+          end
+        end.to raise_error(Gitlab::Database::LoadBalancing::ConnectionProxy::WriteInsideReadOnlyTransactionError)
+      end
+    end
+  end
+end
diff --git a/spec/lib/gitlab/database/migration_helpers_spec.rb b/spec/lib/gitlab/database/migration_helpers_spec.rb
index 40720628a89..f0ea07646fb 100644
--- a/spec/lib/gitlab/database/migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/migration_helpers_spec.rb
@@ -2001,6 +2001,41 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
     end
   end
 
+  describe '#ensure_batched_background_migration_is_finished' do
+    let(:configuration) do
+      {
+        job_class_name: 'CopyColumnUsingBackgroundMigrationJob',
+        table_name: :events,
+        column_name: :id,
+        job_arguments: [[:id], [:id_convert_to_bigint]]
+      }
+    end
+
+    subject(:ensure_batched_background_migration_is_finished) { model.ensure_batched_background_migration_is_finished(**configuration) }
+
+    it 'raises an error when migration exists and is not marked as finished' do
+      create(:batched_background_migration, configuration.merge(status: :active))
+
+      expect { ensure_batched_background_migration_is_finished }
+        .to raise_error "Expected batched background migration for the given configuration to be marked as 'finished', but it is 'active': #{configuration}"
+    end
+
+    it 'does not raise error when migration exists and is marked as finished' do
+      create(:batched_background_migration, configuration.merge(status: :finished))
+
+      expect { ensure_batched_background_migration_is_finished }
+        .not_to raise_error
+    end
+
+    it 'logs a warning when migration does not exist' do
+      expect(Gitlab::AppLogger).to receive(:warn)
+        .with("Could not find batched background migration for the given configuration: #{configuration}")
+
+      expect { ensure_batched_background_migration_is_finished }
+        .not_to raise_error
+    end
+  end
+
   describe '#index_exists_by_name?' do
     it 'returns true if an index exists' do
       ActiveRecord::Base.connection.execute(
diff --git a/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb b/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb
index c6d456964cf..e096e7f6e91 100644
--- a/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb
@@ -242,6 +242,98 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
     end
   end
 
+  describe '#requeue_background_migration_jobs_by_range_at_intervals' do
+    let!(:job_class_name) { 'TestJob' }
+    let!(:pending_job_1) { create(:background_migration_job, class_name: job_class_name, status: :pending, arguments: [1, 2]) }
+    let!(:pending_job_2) { create(:background_migration_job, class_name: job_class_name, status: :pending, arguments: [3, 4]) }
+    let!(:successful_job_1) { create(:background_migration_job, class_name: job_class_name, status: :succeeded, arguments: [5, 6]) }
+    let!(:successful_job_2) { create(:background_migration_job, class_name: job_class_name, status: :succeeded, arguments: [7, 8]) }
+
+    around do |example|
+      freeze_time do
+        Sidekiq::Testing.fake! do
+          example.run
+        end
+      end
+    end
+
+    subject { model.requeue_background_migration_jobs_by_range_at_intervals(job_class_name, 10.minutes) }
+
+    it 'returns the expected duration' do
+      expect(subject).to eq(20.minutes)
+    end
+
+    context 'when nothing is queued' do
+      subject { model.requeue_background_migration_jobs_by_range_at_intervals('FakeJob', 10.minutes) }
+
+      it 'returns expected duration of zero when nothing gets queued' do
+        expect(subject).to eq(0)
+      end
+    end
+
+    it 'queues pending jobs' do
+      subject
+
+      expect(BackgroundMigrationWorker.jobs[0]['args']).to eq([job_class_name, [1, 2]])
+      expect(BackgroundMigrationWorker.jobs[0]['at']).to be_nil
+      expect(BackgroundMigrationWorker.jobs[1]['args']).to eq([job_class_name, [3, 4]])
+      expect(BackgroundMigrationWorker.jobs[1]['at']).to eq(10.minutes.from_now.to_f)
+    end
+
+    context 'with batch_size option' do
+      subject { model.requeue_background_migration_jobs_by_range_at_intervals(job_class_name, 10.minutes, batch_size: 1) }
+
+      it 'returns the expected duration' do
+        expect(subject).to eq(20.minutes)
+      end
+
+      it 'queues pending jobs' do
+        subject
+
+        expect(BackgroundMigrationWorker.jobs[0]['args']).to eq([job_class_name, [1, 2]])
+        expect(BackgroundMigrationWorker.jobs[0]['at']).to be_nil
+        expect(BackgroundMigrationWorker.jobs[1]['args']).to eq([job_class_name, [3, 4]])
+        expect(BackgroundMigrationWorker.jobs[1]['at']).to eq(10.minutes.from_now.to_f)
+      end
+
+      it 'retrieve jobs in batches' do
+        jobs = double('jobs')
+        expect(Gitlab::Database::BackgroundMigrationJob).to receive(:pending) { jobs }
+        allow(jobs).to receive(:where).with(class_name: job_class_name) { jobs }
+        expect(jobs).to receive(:each_batch).with(of: 1)
+
+        subject
+      end
+    end
+
+    context 'with initial_delay option' do
+      let_it_be(:initial_delay) { 3.minutes }
+
+      subject { model.requeue_background_migration_jobs_by_range_at_intervals(job_class_name, 10.minutes, initial_delay: initial_delay) }
+
+      it 'returns the expected duration' do
+        expect(subject).to eq(23.minutes)
+      end
+
+      it 'queues pending jobs' do
+        subject
+
+        expect(BackgroundMigrationWorker.jobs[0]['args']).to eq([job_class_name, [1, 2]])
+        expect(BackgroundMigrationWorker.jobs[0]['at']).to eq(3.minutes.from_now.to_f)
+        expect(BackgroundMigrationWorker.jobs[1]['args']).to eq([job_class_name, [3, 4]])
+        expect(BackgroundMigrationWorker.jobs[1]['at']).to eq(13.minutes.from_now.to_f)
+      end
+
+      context 'when nothing is queued' do
+        subject { model.requeue_background_migration_jobs_by_range_at_intervals('FakeJob', 10.minutes) }
+
+        it 'returns expected duration of zero when nothing gets queued' do
+          expect(subject).to eq(0)
+        end
+      end
+    end
+  end
+
   describe '#perform_background_migration_inline?' do
     it 'returns true in a test environment' do
       stub_rails_env('test')
@@ -269,6 +361,38 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
       allow(Gitlab::Database::PgClass).to receive(:for_table).and_call_original
     end
 
+    context 'when such migration already exists' do
+      it 'does not create duplicate migration' do
+        create(
+          :batched_background_migration,
+          job_class_name: 'MyJobClass',
+          table_name: :projects,
+          column_name: :id,
+          interval: 10.minutes,
+          min_value: 5,
+          max_value: 1005,
+          batch_class_name: 'MyBatchClass',
+          batch_size: 200,
+          sub_batch_size: 20,
+          job_arguments: [[:id], [:id_convert_to_bigint]]
+        )
+
+        expect do
+          model.queue_batched_background_migration(
+            'MyJobClass',
+            :projects,
+            :id,
+            [:id], [:id_convert_to_bigint],
+            job_interval: 5.minutes,
+            batch_min_value: 5,
+            batch_max_value: 1000,
+            batch_class_name: 'MyBatchClass',
+            batch_size: 100,
+            sub_batch_size: 10)
+        end.not_to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }
+      end
+    end
+
     it 'creates the database record for the migration' do
       expect(Gitlab::Database::PgClass).to receive(:for_table).with(:projects).and_return(pgclass_info)
 
diff --git a/spec/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers_spec.rb b/spec/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers_spec.rb
index 79ddb450d7a..4f1d6302331 100644
--- a/spec/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers_spec.rb
+++ b/spec/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers_spec.rb
@@ -580,7 +580,7 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::TableManagementHe
       it 'idempotently cleans up after failed background migrations' do
         expect(partitioned_model.count).to eq(0)
 
-        partitioned_model.insert!(record2.attributes)
+        partitioned_model.insert(record2.attributes, unique_by: [:id, :created_at])
 
         expect_next_instance_of(Gitlab::Database::PartitioningMigrationHelpers::BackfillPartitionedTable) do |backfill|
           allow(backfill).to receive(:transaction_open?).and_return(false)
diff --git a/spec/lib/gitlab/database/postgresql_adapter/type_map_cache_spec.rb b/spec/lib/gitlab/database/postgresql_adapter/type_map_cache_spec.rb
new file mode 100644
index 00000000000..e9c512f94bb
--- /dev/null
+++ b/spec/lib/gitlab/database/postgresql_adapter/type_map_cache_spec.rb
@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::PostgresqlAdapter::TypeMapCache do
+  let(:db_config) { ActiveRecord::Base.configurations.configs_for(env_name: 'test', name: 'primary').configuration_hash }
+  let(:adapter_class) { ActiveRecord::ConnectionAdapters::PostgreSQLAdapter }
+
+  before do
+    adapter_class.type_map_cache.clear
+  end
+
+  describe '#initialize_type_map' do
+    it 'caches loading of types in memory' do
+      recorder_without_cache = ActiveRecord::QueryRecorder.new(skip_schema_queries: false) { initialize_connection.disconnect! }
+      expect(recorder_without_cache.log).to include(a_string_matching(/FROM pg_type/)).twice
+
+      recorder_with_cache = ActiveRecord::QueryRecorder.new(skip_schema_queries: false) { initialize_connection.disconnect! }
+
+      expect(recorder_with_cache.count).to be < recorder_without_cache.count
+
+      # There's still one pg_type query left here because `#add_pg_decoders` executes another pg_type query
+      # in https://github.com/rails/rails/blob/v6.1.3.2/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb#L912.
+      # This query is much cheaper because it only returns very few records.
+      expect(recorder_with_cache.log).to include(a_string_matching(/FROM pg_type/)).once
+    end
+
+    it 'only reuses the cache if the connection parameters are exactly the same' do
+      initialize_connection.disconnect!
+
+      other_config = db_config.dup
+      other_config[:connect_timeout] = db_config[:connect_timeout].to_i + 10
+
+      recorder = ActiveRecord::QueryRecorder.new(skip_schema_queries: false) { initialize_connection(other_config).disconnect! }
+
+      expect(recorder.log).to include(a_string_matching(/FROM pg_type/)).twice
+    end
+  end
+
+  describe '#reload_type_map' do
+    it 'clears the cache and executes the type map query again' do
+      initialize_connection.disconnect!
+
+      connection = initialize_connection
+      recorder = ActiveRecord::QueryRecorder.new(skip_schema_queries: false) { connection.reload_type_map }
+
+      expect(recorder.log).to include(a_string_matching(/FROM pg_type/)).once
+    end
+  end
+
+  # Based on https://github.com/rails/rails/blob/v6.1.3.2/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb#L36-L41
+  def initialize_connection(config = db_config)
+    conn_params = config.symbolize_keys.compact
+
+    conn_params[:user] = conn_params.delete(:username) if conn_params[:username]
+    conn_params[:dbname] = conn_params.delete(:database) if conn_params[:database]
+
+    valid_conn_param_keys = PG::Connection.conndefaults_hash.keys + [:requiressl]
+    conn_params.slice!(*valid_conn_param_keys)
+
+    adapter_class.new(
+      adapter_class.new_client(conn_params),
+      ActiveRecord::Base.logger,
+      conn_params,
+      config
+    )
+  end
+end
diff --git a/spec/lib/gitlab/database/with_lock_retries_spec.rb b/spec/lib/gitlab/database/with_lock_retries_spec.rb
index b08f39fc92a..df2c506e163 100644
--- a/spec/lib/gitlab/database/with_lock_retries_spec.rb
+++ b/spec/lib/gitlab/database/with_lock_retries_spec.rb
@@ -242,10 +242,10 @@ RSpec.describe Gitlab::Database::WithLockRetries do
     let(:timing_configuration) { [[0.015.seconds, 0.025.seconds], [0.015.seconds, 0.025.seconds]] } # 15ms, 25ms
 
     it 'executes `SET LOCAL lock_timeout` using the configured timeout value in milliseconds' do
-      expect(ActiveRecord::Base.connection).to receive(:execute).with("SAVEPOINT active_record_1").and_call_original
-      expect(ActiveRecord::Base.connection).to receive(:execute).with('RESET idle_in_transaction_session_timeout; RESET lock_timeout').and_call_original
+      expect(ActiveRecord::Base.connection).to receive(:execute).with("RESET idle_in_transaction_session_timeout; RESET lock_timeout").and_call_original
+      expect(ActiveRecord::Base.connection).to receive(:execute).with("SAVEPOINT active_record_1", "TRANSACTION").and_call_original
       expect(ActiveRecord::Base.connection).to receive(:execute).with("SET LOCAL lock_timeout TO '15ms'").and_call_original
-      expect(ActiveRecord::Base.connection).to receive(:execute).with("RELEASE SAVEPOINT active_record_1").and_call_original
+      expect(ActiveRecord::Base.connection).to receive(:execute).with("RELEASE SAVEPOINT active_record_1", "TRANSACTION").and_call_original
 
       subject.run { }
     end