Don't run WAL queries when not using replicas

This changes the load balancer to not run redundant WAL related queries
when no replicas are in use. These queries are not needed as a single
database is always in sync with itself. It should also ensure the new
load balancer setup works on Aurora, at least for the time being.

Changelog: fixed

lib/gitlab/database/load_balancing/primary_host.rb

@@ -11,6 +11,12 @@ module Gitlab
       # balancing is enabled, but no replicas have been configured (= the
       # default case).
       class PrimaryHost
+        WAL_ERROR_MESSAGE = <<~MSG.strip
+          Obtaining WAL information when not using any replicas results in
+          redundant queries, and may break installations that don't support
+          streaming replication (e.g. AWS' Aurora database).
+        MSG
+
         def initialize(load_balancer)
           @load_balancer = load_balancer
         end
@@ -51,30 +57,16 @@ module Gitlab
         end
 
         def primary_write_location
-          @load_balancer.primary_write_location
+          raise NotImplementedError, WAL_ERROR_MESSAGE
         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 *Host::CONNECTION_ERRORS
-          nil
+          raise NotImplementedError, WAL_ERROR_MESSAGE
         end
 
         def caught_up?(_location)
           true
         end
-
-        def query_and_release(sql)
-          connection.select_all(sql).first || {}
-        rescue StandardError
-          {}
-        ensure
-          release_connection
-        end
       end
     end
   end

lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb

@@ -42,6 +42,9 @@ module Gitlab
         end
 
         def wal_location_for(load_balancer)
+          # When only using the primary there's no need for any WAL queries.
+          return if load_balancer.primary_only?
+
           if ::Gitlab::Database::LoadBalancing::Session.current.use_primary?
             load_balancer.primary_write_location
           else

lib/gitlab/database/load_balancing/sticking.rb

@@ -104,6 +104,10 @@ module Gitlab
         end
 
         def with_primary_write_location
+          # When only using the primary, there's no point in getting write
+          # locations, as the primary is always in sync with itself.
+          return if @load_balancer.primary_only?
+
           location = @load_balancer.primary_write_location
 
           return if location.blank?

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

@@ -63,9 +63,8 @@ RSpec.describe Gitlab::Database::LoadBalancing::PrimaryHost do
   end
 
   describe '#primary_write_location' do
-    it 'returns the write location of the primary' do
-      expect(host.primary_write_location).to be_an_instance_of(String)
-      expect(host.primary_write_location).not_to be_empty
+    it 'raises NotImplementedError' do
+      expect { host.primary_write_location }.to raise_error(NotImplementedError)
     end
   end
 
@@ -76,51 +75,8 @@ RSpec.describe Gitlab::Database::LoadBalancing::PrimaryHost do
   end
 
   describe '#database_replica_location' do
-    let(:connection) { double(:connection) }
-
-    it 'returns the write ahead location of the replica', :aggregate_failures do
-      expect(host)
-        .to receive(:query_and_release)
-        .and_return({ 'location' => '0/D525E3A8' })
-
-      expect(host.database_replica_location).to be_an_instance_of(String)
-    end
-
-    it 'returns nil when the database query returned no rows' do
-      expect(host).to receive(:query_and_release).and_return({})
-
-      expect(host.database_replica_location).to be_nil
-    end
-
-    it 'returns nil when the database connection fails' do
-      allow(host).to receive(:connection).and_raise(PG::Error)
-
-      expect(host.database_replica_location).to be_nil
-    end
-  end
-
-  describe '#query_and_release' do
-    it 'executes a SQL query' do
-      results = host.query_and_release('SELECT 10 AS number')
-
-      expect(results).to be_an_instance_of(Hash)
-      expect(results['number'].to_i).to eq(10)
-    end
-
-    it 'releases the connection after running the query' do
-      expect(host)
-        .to receive(:release_connection)
-        .once
-
-      host.query_and_release('SELECT 10 AS number')
-    end
-
-    it 'returns an empty Hash in the event of an error' do
-      expect(host.connection)
-        .to receive(:select_all)
-        .and_raise(RuntimeError, 'kittens')
-
-      expect(host.query_and_release('SELECT 10 AS number')).to eq({})
+    it 'raises NotImplementedError' do
+      expect { host.database_replica_location }.to raise_error(NotImplementedError)
     end
   end
 end

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

@@ -21,7 +21,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
     middleware.call(worker_class, job, nil, nil) {}
   end
 
-  describe '#call' do
+  describe '#call', :database_replica do
     shared_context 'data consistency worker class' do |data_consistency, feature_flag|
       let(:expected_consistency) { data_consistency }
       let(:worker_class) do

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

@@ -188,6 +188,10 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
     end
 
     it 'sticks an entity to the primary', :aggregate_failures do
+      allow(ActiveRecord::Base.connection.load_balancer)
+        .to receive(:primary_only?)
+        .and_return(false)
+
       ids.each do |id|
         expect(sticking)
           .to receive(:set_write_location_for)
@@ -199,6 +203,12 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
 
       subject
     end
+
+    it 'does not update the write location when no replicas are used' do
+      expect(sticking).not_to receive(:set_write_location_for)
+
+      subject
+    end
   end
 
   describe '#stick' do
@@ -221,12 +231,22 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
         .to receive(:primary_write_location)
         .and_return('foo')
 
+      allow(ActiveRecord::Base.connection.load_balancer)
+        .to receive(:primary_only?)
+        .and_return(false)
+
       expect(sticking)
         .to receive(:set_write_location_for)
         .with(:user, 42, 'foo')
 
       sticking.mark_primary_write_location(:user, 42)
     end
+
+    it 'does nothing when no replicas are used' do
+      expect(sticking).not_to receive(:set_write_location_for)
+
+      sticking.mark_primary_write_location(:user, 42)
+    end
   end
 
   describe '#unstick' do

spec/models/ci/pipeline_spec.rb

@@ -2791,17 +2791,7 @@ RSpec.describe Ci::Pipeline, :mailer, factory_default: :keep do
         extra_update_queries = 4 # transition ... => :canceled, queue pop
         extra_generic_commit_status_validation_queries = 2 # name_uniqueness_across_types
 
-        # The number of extra load balancing queries depends on whether or not
-        # we use a load balancer for CI. That in turn depends on the contents of
-        # database.yml, so here we support both cases.
-        extra_load_balancer_queries =
-          if Gitlab::Database.has_config?(:ci)
-            6
-          else
-            3
-          end
-
-        expect(control2.count).to eq(control1.count + extra_update_queries + extra_generic_commit_status_validation_queries + extra_load_balancer_queries)
+        expect(control2.count).to eq(control1.count + extra_update_queries + extra_generic_commit_status_validation_queries)
       end
     end