Add database_replica_location

- Fix specs for sidekiq middleware
- Add query to the host

ee/lib/gitlab/database/load_balancing/host.rb

@@ -163,6 +163,16 @@ module Gitlab
           load_balancer.release_primary_connection
         end
 
+        def database_replica_location
+          row = query_and_release(<<-SQL.squish)
+            SELECT pg_last_wal_replay_lsn()::text AS location
+          SQL
+
+          row['location'] if row.any?
+        rescue *CONNECTION_ERRORS
+          nil
+        end
+
         # Returns true if this host has caught up to the given transaction
         # write location.
         #

ee/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb

@@ -27,11 +27,7 @@ module Gitlab
           if Session.current.performed_write?
             job['database_write_location'] = load_balancer.primary_write_location
           else
-            # It is possible that the current replica has a different write-ahead
-            # replication log location from the sidekiq server replica.
-            # In the follow-up issue https://gitlab.com/gitlab-org/gitlab/-/issues/325519,
-            # we want to pass database replica location as well
-            job['database_replica_location'] = true
+            job['database_replica_location'] = load_balancer.host.database_replica_location
           end
         end
 

ee/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb

@@ -28,7 +28,9 @@ module Gitlab
           return true if worker_class.get_data_consistency == :always
           return true unless worker_class.get_data_consistency_feature_flag_enabled?
 
-          if job['database_replica_location'] || replica_caught_up?(job['database_write_location'])
+          location = job['database_write_location'] || job['database_replica_location']
+
+          if replica_caught_up?(location)
             job[:database_chosen] = 'replica'
             false
           elsif worker_class.get_data_consistency == :delayed && job['retry_count'].to_i == 0

ee/spec/lib/gitlab/database/load_balancing/host_spec.rb

@@ -382,6 +382,36 @@ RSpec.describe Gitlab::Database::LoadBalancing::Host do
     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')

ee/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb

@@ -31,16 +31,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
       end
     end
 
-    shared_examples_for 'mark database_replica_location' do
-      it 'passes database_replica_location' do
-        expect(middleware).not_to receive(:load_balancer)
-
-        middleware.call(worker_class, job, double(:queue), redis_pool) { 10 }
-
-        expect(job['database_replica_location']).to be_truthy
-      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 }
@@ -68,7 +58,13 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
           allow(Gitlab::Database::LoadBalancing::Session.current).to receive(:performed_write?).and_return(false)
         end
 
-        include_examples 'mark database_replica_location'
+        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

ee/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb

@@ -31,6 +31,14 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware do
       end
     end
 
+    shared_examples_for 'job marked with chosen database' do
+      it 'yields and sets 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
+
     shared_examples_for 'stick to the primary' do
       it 'sticks to the primary' do
         middleware.call(worker, job, double(:queue)) do
@@ -39,6 +47,18 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware do
       end
     end
 
+    shared_examples_for 'replica is up to date' do |location|
+      it 'do 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
+    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
 
@@ -50,43 +70,37 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware do
         include_examples 'stick to the primary'
       end
 
-      context 'database replica location is set' do
-        let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e', 'database_replica_location' => 'true' } }
+      context 'when database replica location is set' do
+        let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e', 'database_replica_location' => '0/D525E3A8' } }
 
-        it 'do not stick to the primary' do
-          middleware.call(worker, job, double(:queue)) do
-            expect(Gitlab::Database::LoadBalancing::Session.current.use_primary?).not_to be_truthy
-          end
+        before do
+          allow(middleware).to receive(:replica_caught_up?).and_return(true)
         end
-      end
-
-      context 'write was not performed' do
-        let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e' } }
 
-        it 'do not stick to the primary' do
-          middleware.call(worker, job, double(:queue)) do
-            expect(Gitlab::Database::LoadBalancing::Session.current.use_primary?).not_to be_truthy
-          end
-        end
+        it_behaves_like 'replica is up to date', '0/D525E3A8'
       end
 
-      context 'replica is up to date' do
+      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 'do not stick to the primary' do
-          middleware.call(worker, job, double(:queue)) do
-            expect(Gitlab::Database::LoadBalancing::Session.current.use_primary?).not_to be_truthy
-          end
-        end
+        it_behaves_like 'replica is up to date', '0/D525E3A8'
+      end
+
+      context 'when database location is not set' do
+        let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e' } }
+
+        it_behaves_like 'replica is up to date', nil
       end
     end
 
     let(:queue) { 'default' }
     let(:redis_pool) { Sidekiq.redis_pool }
     let(:worker) { worker_class.new }
-    let(:job) { { "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", 'primary_write_location' => '0/D525E3A8' } }
+    let(:job) { { "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e" } }
     let(:block) { 10 }
 
     before do
@@ -128,6 +142,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware do
           end
 
           include_examples 'stick to the primary'
+          include_examples 'job marked with chosen database'
         end
       end
     end
@@ -141,6 +156,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware do
         end
 
         include_examples 'stick to the primary'
+        include_examples 'job marked with chosen database'
       end
     end
   end