Merge branch...

Merge branch '322133-applicationrecord-with_fast_statement_timeout-forces-known-expensive-queries-to-run-on-the' into 'master'

Force ApplicationRecord#with_fast_statement_timeout to use replicas

See merge request gitlab-org/gitlab!56476

app/models/application_record.rb

@@ -52,9 +52,9 @@ class ApplicationRecord < ActiveRecord::Base
 
   # Start a new transaction with a shorter-than-usual statement timeout. This is
   # currently one third of the default 15-second timeout
-  def self.with_fast_statement_timeout
+  def self.with_fast_read_statement_timeout(timeout_ms = 5000)
     transaction(requires_new: true) do
-      connection.exec_query("SET LOCAL statement_timeout = 5000")
+      connection.exec_query("SET LOCAL statement_timeout = #{timeout_ms}")
 
       yield
     end
@@ -79,3 +79,5 @@ class ApplicationRecord < ActiveRecord::Base
     enum(enum_mod.key => values)
   end
 end
+
+ApplicationRecord.prepend_if_ee('EE::ApplicationRecordHelpers')

ee/app/models/ee/application_record_helpers.rb

+# frozen_string_literal: true
+
+module EE
+  # Intentionally pick a different name, to prevent naming conflicts
+  module ApplicationRecordHelpers
+    extend ActiveSupport::Concern
+
+    class_methods do
+      extend ::Gitlab::Utils::Override
+
+      override :with_fast_read_statement_timeout
+      def with_fast_read_statement_timeout(timeout_ms = 5000)
+        ::Gitlab::Database::LoadBalancing::Session.current.use_replica_if_possible do
+          super
+        end
+      end
+    end
+  end
+end

ee/changelogs/unreleased/322133-applicationrecord-with_fast_statement_timeout-forces-known-expe.yml

+---
+title: Force ApplicationRecord#with_fast_statement_timeout to use replicas
+merge_request: 56476
+author:
+type: performance

ee/lib/gitlab/database/load_balancing.rb

@@ -140,7 +140,7 @@ module Gitlab
       # Clear configuration
       def self.clear_configuration
         @proxy = nil
-        remove_instance_variable(:@feature_available)
+        remove_instance_variable(:@feature_available) if defined?(@feature_available)
       end
 
       def self.active_record_models

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

@@ -10,6 +10,9 @@ module Gitlab
       # The ConnectionProxy class redirects ActiveRecord connection requests to
       # the right load balancer pool, depending on the type of query.
       class ConnectionProxy
+        WriteInsideReadOnlyTransactionError = Class.new(StandardError)
+        READ_ONLY_TRANSACTION_KEY = :load_balacing_read_only_transaction
+
         attr_reader :load_balancer
 
         # These methods perform writes after which we need to stick to the
@@ -18,7 +21,6 @@ module Gitlab
           delete
           delete_all
           insert
-          transaction
           update
           update_all
         ).freeze
@@ -57,9 +59,25 @@ module Gitlab
           end
         end
 
+        def transaction(*args, &block)
+          if ::Gitlab::Database::LoadBalancing::Session.current.use_replica?
+            track_read_only_transaction!
+            read_using_load_balancer(:transaction, args, &block)
+          else
+            write_using_load_balancer(:transaction, args, sticky: true, &block)
+          end
+
+        ensure
+          untrack_read_only_transaction!
+        end
+
         # Delegates all unknown messages to a read-write connection.
         def method_missing(name, *args, &block)
-          write_using_load_balancer(name, args, &block)
+          if ::Gitlab::Database::LoadBalancing::Session.current.use_replica?
+            read_using_load_balancer(name, args, &block)
+          else
+            write_using_load_balancer(name, args, &block)
+          end
         end
 
         # Performs a read using the load balancer.
@@ -79,6 +97,10 @@ module Gitlab
         # sticky - If set to true the session will stick to the master after
         #          the write.
         def write_using_load_balancer(name, args, sticky: false, &block)
+          if read_only_transaction?
+            raise WriteInsideReadOnlyTransactionError, 'A write query is performed inside a read-only transaction'
+          end
+
           result = @load_balancer.read_write do |connection|
             # Sticking has to be enabled before calling the method. Not doing so
             # could lead to methods called in a block still being performed on a
@@ -90,6 +112,20 @@ module Gitlab
 
           result
         end
+
+        private
+
+        def track_read_only_transaction!
+          Thread.current[READ_ONLY_TRANSACTION_KEY] = true
+        end
+
+        def untrack_read_only_transaction!
+          Thread.current[READ_ONLY_TRANSACTION_KEY] = nil
+        end
+
+        def read_only_transaction?
+          Thread.current[READ_ONLY_TRANSACTION_KEY] == true
+        end
       end
     end
   end

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

@@ -19,6 +19,7 @@ module Gitlab
         def initialize(hosts = [])
           @host_list = HostList.new(hosts.map { |addr| Host.new(addr, self) })
           @connection_db_roles = {}.compare_by_identity
+          @connection_db_roles_count = {}.compare_by_identity
         end
 
         # Yields a connection that can be used for reads.
@@ -34,11 +35,11 @@ module Gitlab
 
             begin
               connection = host.connection
-              @connection_db_roles[connection] = ROLE_REPLICA
+              track_connection_role(connection, ROLE_REPLICA)
 
               return yield connection
             rescue => error
-              @connection_db_roles.delete(connection) if connection.present?
+              untrack_connection_role(connection)
 
               if serialization_failure?(error)
                 # This error can occur when a query conflicts. See
@@ -83,7 +84,7 @@ module Gitlab
 
           read_write(&block)
         ensure
-          @connection_db_roles.delete(connection) if connection.present?
+          untrack_connection_role(connection)
         end
 
         # Yields a connection that can be used for both reads and writes.
@@ -94,12 +95,12 @@ module Gitlab
           # a few times.
           retry_with_backoff do
             connection = ActiveRecord::Base.retrieve_connection
-            @connection_db_roles[connection] = ROLE_PRIMARY
+            track_connection_role(connection, ROLE_PRIMARY)
 
             yield connection
           end
         ensure
-          @connection_db_roles.delete(connection) if connection.present?
+          untrack_connection_role(connection)
         end
 
         # Recognize the role (primary/replica) of the database this connection
@@ -203,6 +204,22 @@ module Gitlab
         def ensure_caching!
           host.enable_query_cache! unless host.query_cache_enabled
         end
+
+        def track_connection_role(connection, role)
+          @connection_db_roles[connection] = role
+          @connection_db_roles_count[connection] ||= 0
+          @connection_db_roles_count[connection] += 1
+        end
+
+        def untrack_connection_role(connection)
+          return if connection.blank? || @connection_db_roles_count[connection].blank?
+
+          @connection_db_roles_count[connection] -= 1
+          if @connection_db_roles_count[connection] <= 0
+            @connection_db_roles.delete(connection)
+            @connection_db_roles_count.delete(connection)
+          end
+        end
       end
     end
   end

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

@@ -55,6 +55,24 @@ module Gitlab
           @ignore_writes = false
         end
 
+        # Indicate that all the SQL statements from anywhere inside this block
+        # should use a replica. This is a weak enforcement. The queries
+        # prioritize using a primary over a replica in those cases:
+        # - If the queries are about to write
+        # - The current session already performed writes
+        # - It prefers to use primary, aka, use_primary or use_primary! were called
+        def use_replica_if_possible(&blk)
+          used_replica = @use_replica
+          @use_replica = true
+          yield
+        ensure
+          @use_replica = used_replica
+        end
+
+        def use_replica?
+          @use_replica == true && !use_primary? && !performed_write?
+        end
+
         def write!
           @performed_write = true
 

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

@@ -108,26 +108,78 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
   # We have an extra test for #transaction here to make sure that nested queries
   # are also sent to a primary.
   describe '#transaction' do
-    after do
-      Gitlab::Database::LoadBalancing::Session.clear_session
+    let(:session) { double(:session) }
+
+    before do
+      allow(Gitlab::Database::LoadBalancing::Session).to receive(:current)
+        .and_return(session)
+    end
+
+    context 'session prefers to use a replica' do
+      let(:replica) { double(:connection) }
+
+      before do
+        allow(session).to receive(:use_replica?).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
 
-    it 'runs the transaction and any nested queries on the primary' do
-      primary = double(:connection)
+    context 'session does not prefer to use a replica' do
+      let(:primary) { double(:connection) }
 
-      allow(primary).to receive(:transaction).and_yield
-      allow(primary).to receive(:select)
+      before do
+        allow(session).to receive(:use_replica?).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
 
-      expect(proxy.load_balancer).to receive(:read_write)
-        .twice.and_yield(primary)
+      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!)
 
-      # This expectation is put in place to ensure no read is performed.
-      expect(proxy.load_balancer).not_to receive(:read)
+          proxy.transaction { proxy.select('true') }
+        end
+      end
 
-      proxy.transaction { proxy.select('true') }
+      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
 
-      expect(Gitlab::Database::LoadBalancing::Session.current.use_primary?)
-        .to eq(true)
+          proxy.transaction { proxy.insert('something') }
+        end
+      end
     end
   end
 
@@ -147,6 +199,32 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
       expect { proxy.case_sensitive_comparison(:table, :attribute, :column, { value: :value, format: :format }) }
         .not_to raise_error
     end
+
+    context 'current session prefers to use a replica' do
+      let(:session) { double(:session) }
+
+      before do
+        allow(Gitlab::Database::LoadBalancing::Session).to receive(:current)
+          .and_return(session)
+        allow(session).to receive(:use_replica?).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

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

@@ -147,6 +147,20 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
       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
@@ -158,6 +172,20 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
       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)

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

@@ -138,4 +138,144 @@ RSpec.describe Gitlab::Database::LoadBalancing::Session do
       expect(instance).to be_using_primary
     end
   end
+
+  describe '#use_replica_if_possible' do
+    let(:instance) { described_class.new }
+
+    it 'uses replica during block' do
+      expect do |blk|
+        instance.use_replica_if_possible do
+          expect(instance.use_replica?).to eq(true)
+
+          # call yield probe
+          blk.to_proc.call
+        end
+      end.to yield_control
+
+      expect(instance.use_replica?).to eq(false)
+    end
+
+    it 'restores state after use' do
+      expect do |blk|
+        instance.use_replica_if_possible do
+          instance.use_replica_if_possible do
+            expect(instance.use_replica?).to eq(true)
+
+            # call yield probe
+            blk.to_proc.call
+          end
+
+          expect(instance.use_replica?).to eq(true)
+        end
+      end.to yield_control
+
+      expect(instance.use_replica?).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.use_replica?).to eq(false)
+
+        expect do |blk|
+          instance.use_replica_if_possible do
+            expect(instance.use_replica?).to eq(false)
+
+            # call yield probe
+            blk.to_proc.call
+          end
+        end.to yield_control
+
+        expect(instance.use_replica?).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.use_replica?).to eq(false)
+
+        expect do |blk|
+          instance.use_replica_if_possible do
+            expect(instance.use_replica?).to eq(false)
+
+            # call yield probe
+            blk.to_proc.call
+          end
+        end.to yield_control
+
+        expect(instance.use_replica?).to eq(false)
+      end
+    end
+
+    context 'when primary was used inside the block' do
+      it 'uses primary aterward' do
+        expect(instance.use_replica?).to eq(false)
+
+        instance.use_replica_if_possible do
+          expect(instance.use_replica?).to eq(true)
+
+          instance.use_primary!
+
+          expect(instance.use_replica?).to eq(false)
+        end
+
+        expect(instance.use_replica?).to eq(false)
+      end
+
+      it 'restores state after use' do
+        instance.use_replica_if_possible do
+          instance.use_replica_if_possible do
+            expect(instance.use_replica?).to eq(true)
+
+            instance.use_primary!
+
+            expect(instance.use_replica?).to eq(false)
+          end
+
+          expect(instance.use_replica?).to eq(false)
+        end
+
+        expect(instance.use_replica?).to eq(false)
+      end
+    end
+
+    context 'when a write was performed inside the block' do
+      it 'uses primary aterward' do
+        expect(instance.use_replica?).to eq(false)
+
+        instance.use_replica_if_possible do
+          expect(instance.use_replica?).to eq(true)
+
+          instance.write!
+
+          expect(instance.use_replica?).to eq(false)
+        end
+
+        expect(instance.use_replica?).to eq(false)
+      end
+
+      it 'restores state after use' do
+        instance.use_replica_if_possible do
+          instance.use_replica_if_possible do
+            expect(instance.use_replica?).to eq(true)
+
+            instance.write!
+
+            expect(instance.use_replica?).to eq(false)
+          end
+
+          expect(instance.use_replica?).to eq(false)
+        end
+
+        expect(instance.use_replica?).to eq(false)
+      end
+    end
+  end
 end

ee/spec/lib/gitlab/database/load_balancing_spec.rb

@@ -403,4 +403,316 @@ RSpec.describe Gitlab::Database::LoadBalancing do
       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(:hosts) { [ActiveRecord::Base.configurations["development"]['host']] }
+      let(:model) do
+        Class.new(ApplicationRecord) do
+          self.table_name = "load_balancing_test"
+        end
+      end
+
+      before do
+        stub_licensed_features(db_load_balancing: true)
+        # Preloading testing class
+        model.singleton_class.prepend ::Gitlab::Database::LoadBalancing::ActiveRecordProxy
+
+        # Setup load balancing
+        subject.clear_configuration
+        allow(ActiveRecord::Base.singleton_class).to receive(:prepend)
+        subject.configure_proxy(::Gitlab::Database::LoadBalancing::ConnectionProxy.new(hosts))
+        allow(ActiveRecord::Base.configurations[Rails.env])
+          .to receive(:[])
+          .with('load_balancing')
+          .and_return('hosts' => hosts)
+        ::Gitlab::Database::LoadBalancing::Session.clear_session
+      end
+
+      after do
+        subject.clear_configuration
+      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_replica_if_possible
+        [
+          -> {
+            ::Gitlab::Database::LoadBalancing::Session.current.use_replica_if_possible do
+              model.first
+              model.where(name: 'test1').to_a
+            end
+          },
+          false, [:replica, :replica]
+        ],
+
+        # use_replica_if_possible for read-only transaction
+        [
+          -> {
+            ::Gitlab::Database::LoadBalancing::Session.current.use_replica_if_possible do
+              model.transaction do
+                model.first
+                model.where(name: 'test1').to_a
+              end
+            end
+          },
+          false, [:replica, :replica]
+        ],
+
+        # A custom read query inside use_replica_if_possible
+        [
+          -> {
+            ::Gitlab::Database::LoadBalancing::Session.current.use_replica_if_possible do
+              model.connection.exec_query("SELECT 1")
+            end
+          },
+          false, [:replica]
+        ],
+
+        # A custom read query inside a transaction use_replica_if_possible
+        [
+          -> {
+            ::Gitlab::Database::LoadBalancing::Session.current.use_replica_if_possible do
+              model.transaction do
+                model.connection.exec_query("SET LOCAL statement_timeout = 5000")
+                model.count
+              end
+            end
+          },
+          true, [:replica, :replica, :replica, :replica]
+        ],
+
+        # use_replica_if_possible after a write
+        [
+          -> {
+            model.create!(name: 'Test1')
+            ::Gitlab::Database::LoadBalancing::Session.current.use_replica_if_possible do
+              model.first
+            end
+          },
+          false, [:primary, :primary]
+        ],
+
+        # use_replica_if_possible after use_primary!
+        [
+          -> {
+            ::Gitlab::Database::LoadBalancing::Session.current.use_primary!
+            ::Gitlab::Database::LoadBalancing::Session.current.use_replica_if_possible do
+              model.first
+            end
+          },
+          false, [:primary]
+        ],
+
+        # use_replica_if_possible inside use_primary
+        [
+          -> {
+            ::Gitlab::Database::LoadBalancing::Session.current.use_primary do
+              ::Gitlab::Database::LoadBalancing::Session.current.use_replica_if_possible do
+                model.first
+              end
+            end
+          },
+          false, [:primary]
+        ],
+
+        # use_primary inside use_replica_if_possible
+        [
+          -> {
+            ::Gitlab::Database::LoadBalancing::Session.current.use_replica_if_possible do
+              ::Gitlab::Database::LoadBalancing::Session.current.use_primary do
+                model.first
+              end
+            end
+          },
+          false, [:primary]
+        ],
+
+        # A write query inside use_replica_if_possible
+        [
+          -> {
+            ::Gitlab::Database::LoadBalancing::Session.current.use_replica_if_possible do
+              model.first
+              model.delete_all
+              model.where(name: 'test1').to_a
+            end
+          },
+          false, [:replica, :primary, :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|
+          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 'a write inside a transaction inside use_replica_if_possible block' do
+      include_context 'LoadBalancing setup'
+
+      it 'raises an exception' do
+        expect do
+          ::Gitlab::Database::LoadBalancing::Session.current.use_replica_if_possible 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

ee/spec/models/ee/application_record_helpers_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe ApplicationRecord do
+  describe '.with_fast_read_statement_timeout' do
+    let(:session) { double(:session) }
+
+    before do
+      allow(::Gitlab::Database::LoadBalancing::Session).to receive(:current).and_return(session)
+      allow(session).to receive(:use_replica_if_possible).and_yield
+    end
+
+    it 'yields control' do
+      expect do |blk|
+        described_class.with_fast_read_statement_timeout(&blk)
+      end.to yield_control.once
+    end
+
+    context 'when the query runs faster than configured timeout' do
+      it 'executes the query without error' do
+        result = nil
+
+        expect do
+          described_class.with_fast_read_statement_timeout(100) do
+            result = described_class.connection.exec_query('SELECT 1')
+          end
+        end.not_to raise_error
+
+        expect(result).not_to be_nil
+      end
+    end
+
+    # This query hangs for 10ms and then gets cancelled.  As there is no
+    # other way to test the timeout for sure, 10ms of waiting seems to be
+    # reasonable!
+    context 'when the query runs longer than configured timeout' do
+      it 'cancels the query and raiss an exception' do
+        expect do
+          described_class.with_fast_read_statement_timeout(10) do
+            described_class.connection.exec_query('SELECT pg_sleep(0.1)')
+          end
+        end.to raise_error(ActiveRecord::QueryCanceled)
+      end
+    end
+  end
+end

lib/gitlab/issuables_count_for_state.rb

@@ -78,7 +78,7 @@ module Gitlab
       # to perform the calculation more efficiently. Until then, use a shorter
       # timeout and return -1 as a sentinel value if it is triggered
       begin
-        ApplicationRecord.with_fast_statement_timeout do
+        ApplicationRecord.with_fast_read_statement_timeout do
           finder.count_by_state
         end
       rescue ActiveRecord::QueryCanceled => err

spec/models/application_record_spec.rb

@@ -100,4 +100,33 @@ RSpec.describe ApplicationRecord do
       expect(User.where_exists(User.limit(1))).to eq([user])
     end
   end
+
+  describe '.with_fast_read_statement_timeout' do
+    context 'when the query runs faster than configured timeout' do
+      it 'executes the query without error' do
+        result = nil
+
+        expect do
+          described_class.with_fast_read_statement_timeout(100) do
+            result = described_class.connection.exec_query('SELECT 1')
+          end
+        end.not_to raise_error
+
+        expect(result).not_to be_nil
+      end
+    end
+
+    # This query hangs for 10ms and then gets cancelled.  As there is no
+    # other way to test the timeout for sure, 10ms of waiting seems to be
+    # reasonable!
+    context 'when the query runs longer than configured timeout' do
+      it 'cancels the query and raises an exception' do
+        expect do
+          described_class.with_fast_read_statement_timeout(10) do
+            described_class.connection.exec_query('SELECT pg_sleep(0.1)')
+          end
+        end.to raise_error(ActiveRecord::QueryCanceled)
+      end
+    end
+  end
 end