Implement DB `use_replicas_for_read_queries`

ee/changelogs/unreleased/qmnguyen0711-implement-use-replicas-for-all-reads.yml

+---
+title: Implement use_replicas_for_read_queries
+merge_request: 59167
+author:
+type: added

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

@@ -60,7 +60,7 @@ module Gitlab
         end
 
         def transaction(*args, &block)
-          if ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries?
+          if current_session.fallback_to_replicas_for_ambiguous_queries?
             track_read_only_transaction!
             read_using_load_balancer(:transaction, args, &block)
           else
@@ -73,7 +73,7 @@ module Gitlab
 
         # Delegates all unknown messages to a read-write connection.
         def method_missing(name, *args, &block)
-          if ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries?
+          if current_session.fallback_to_replicas_for_ambiguous_queries?
             read_using_load_balancer(name, args, &block)
           else
             write_using_load_balancer(name, args, &block)
@@ -84,10 +84,15 @@ module Gitlab
         #
         # name - The name of the method to call on a connection object.
         def read_using_load_balancer(name, args, &block)
-          method = ::Gitlab::Database::LoadBalancing::Session.current.use_primary? ? :read_write : :read
-
-          @load_balancer.send(method) do |connection|
-            connection.send(name, *args, &block)
+          if current_session.use_primary? &&
+             !current_session.use_replicas_for_read_queries?
+            @load_balancer.read_write do |connection|
+              connection.send(name, *args, &block)
+            end
+          else
+            @load_balancer.read do |connection|
+              connection.send(name, *args, &block)
+            end
           end
         end
 
@@ -105,7 +110,7 @@ module Gitlab
             # 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
             # secondary instead of on a primary (when necessary).
-            ::Gitlab::Database::LoadBalancing::Session.current.write! if sticky
+            current_session.write! if sticky
 
             connection.send(name, *args, &block)
           end
@@ -115,6 +120,10 @@ module Gitlab
 
         private
 
+        def current_session
+          ::Gitlab::Database::LoadBalancing::Session.current
+        end
+
         def track_read_only_transaction!
           Thread.current[READ_ONLY_TRANSACTION_KEY] = true
         end

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

@@ -27,6 +27,8 @@ module Gitlab
           @use_primary = false
           @performed_write = false
           @ignore_writes = false
+          @fallback_to_replicas_for_ambiguous_queries = false
+          @use_replicas_for_read_queries = false
         end
 
         def use_primary?
@@ -55,6 +57,27 @@ module Gitlab
           @ignore_writes = false
         end
 
+        # Indicates that the read SQL statements from anywhere inside this
+        # blocks should use a replica, regardless of the current primary
+        # stickiness or whether a write query is already performed in the
+        # current session. This interface is reserved mostly for performance
+        # purpose. This is a good tool to push expensive queries, which can
+        # tolerate the replica lags, to the replicas.
+        #
+        # Write and ambiguous queries inside this block are still handled by
+        # the primary.
+        def use_replicas_for_read_queries(&blk)
+          previous_flag = @use_replicas_for_read_queries
+          @use_replicas_for_read_queries = true
+          yield
+        ensure
+          @use_replicas_for_read_queries = previous_flag
+        end
+
+        def use_replicas_for_read_queries?
+          @use_replicas_for_read_queries == true
+        end
+
         # Indicate that the ambiguous SQL statements from anywhere inside this
         # block should use a replica. The ambiguous statements include:
         # - Transactions.

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

@@ -153,6 +153,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
 
       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)
@@ -236,9 +237,34 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
         .and_return(session)
     end
 
-    describe 'with a regular session' do
+    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)
@@ -250,6 +276,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
     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')
 

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

@@ -139,6 +139,78 @@ RSpec.describe Gitlab::Database::LoadBalancing::Session do
     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 }
 

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

@@ -561,6 +561,92 @@ RSpec.describe Gitlab::Database::LoadBalancing do
           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
         [
           -> {
@@ -613,7 +699,7 @@ RSpec.describe Gitlab::Database::LoadBalancing do
           -> {
             model.create!(name: 'Test1')
             ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
-              model.first
+              model.connection.exec_query("SELECT 1")
             end
           },
           false, [:primary, :primary]
@@ -624,7 +710,7 @@ RSpec.describe Gitlab::Database::LoadBalancing do
           -> {
             ::Gitlab::Database::LoadBalancing::Session.current.use_primary!
             ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
-              model.first
+              model.connection.exec_query("SELECT 1")
             end
           },
           false, [:primary]
@@ -635,7 +721,7 @@ RSpec.describe Gitlab::Database::LoadBalancing do
           -> {
             ::Gitlab::Database::LoadBalancing::Session.current.use_primary do
               ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
-                model.first
+                model.connection.exec_query("SELECT 1")
               end
             end
           },
@@ -647,7 +733,7 @@ RSpec.describe Gitlab::Database::LoadBalancing do
           -> {
             ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
               ::Gitlab::Database::LoadBalancing::Session.current.use_primary do
-                model.first
+                model.connection.exec_query("SELECT 1")
               end
             end
           },
@@ -658,12 +744,25 @@ RSpec.describe Gitlab::Database::LoadBalancing do
         [
           -> {
             ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
-              model.first
+              model.connection.exec_query("SELECT 1")
               model.delete_all
-              model.where(name: 'test1').to_a
+              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