Merge branch 'yorick/remove-redundant-wal-queries' into 'master'

Don't run WAL queries when not using replicas

See merge request gitlab-org/gitlab!71736

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