Use less aggressive sticking for DB load balancing

Previously when performing a write we would stick to the primary for the
duration of a request. This could lead to unnecessary pressure on the
primary database.

To work around this we stop using the primary the moment all secondaries
are in sync with the last write. If we are in an explicit transaction we
will keep using the primary until the transaction finishes.

Fixes https://gitlab.com/gitlab-org/gitlab-ee/issues/6041

ee/changelogs/unreleased/less-aggressive-db-sticking.yml

+---
+title: Use less aggressive sticking for DB load balancing
+merge_request:
+author:
+type: performance

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

@@ -16,7 +16,6 @@ module Gitlab
           delete
           delete_all
           insert
-          transaction
           update
           update_all
         ).freeze
@@ -46,6 +45,19 @@ module Gitlab
           end
         end
 
+        def transaction(*args, &block)
+          Session.current.enter_transaction
+
+          write_using_load_balancer(:transaction, args, sticky: true, &block)
+        ensure
+          Session.current.leave_transaction
+
+          # When the transaction finishes we need to store the last WAL pointer
+          # since individual writes in a transaction don't perform this
+          # operation.
+          record_last_write_location
+        end
+
         # Delegates all unknown messages to a read-write connection.
         def method_missing(name, *args, &block)
           write_using_load_balancer(name, args, &block)
@@ -55,9 +67,7 @@ module Gitlab
         #
         # name - The name of the method to call on a connection object.
         def read_using_load_balancer(name, args, &block)
-          method = Session.current.use_primary? ? :read_write : :read
-
-          @load_balancer.send(method) do |connection|
+          @load_balancer.send(load_balancer_method_for_read) do |connection|
             connection.send(name, *args, &block)
           end
         end
@@ -77,8 +87,45 @@ module Gitlab
             connection.send(name, *args, &block)
           end
 
+          # We only want to record the last write location if we actually
+          # performed a write, and not for all queries sent to the primary.
+          record_last_write_location if sticky
+
           result
         end
+
+        # Returns the method to use for performing a read-only query.
+        def load_balancer_method_for_read
+          session = Session.current
+
+          return :read unless session.use_primary?
+
+          # If we are still inside an explicit transaction we _must_ send the
+          # queries to the primary.
+          return :read_write if session.in_transaction?
+
+          # If we are not in an explicit transaction we are free to return to
+          # using the secondaries once they are all in sync.
+          if @load_balancer.all_caught_up?(session.last_write_location)
+            session.reset!
+
+            :read
+          else
+            :read_write
+          end
+        end
+
+        def record_last_write_location
+          session = Session.current
+
+          # When we are in a transaction it's likely we will perform many
+          # writes. In this case it's pointless to keep retrieving and storing
+          # the WAL location, as we only care about the location once the
+          # transaction finishes.
+          return if session.in_transaction?
+
+          session.last_write_location = @load_balancer.primary_write_location
+        end
       end
     end
   end

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

@@ -9,6 +9,8 @@ module Gitlab
       class Session
         CACHE_KEY = :gitlab_load_balancer_session
 
+        attr_accessor :last_write_location
+
         def self.current
           RequestStore[CACHE_KEY] ||= new
         end
@@ -18,6 +20,13 @@ module Gitlab
         end
 
         def initialize
+          @transaction_nesting = 0
+
+          reset!
+        end
+
+        def reset!
+          @last_write_location = nil
           @use_primary = false
           @performed_write = false
         end
@@ -26,6 +35,18 @@ module Gitlab
           @use_primary
         end
 
+        def enter_transaction
+          @transaction_nesting += 1
+        end
+
+        def leave_transaction
+          @transaction_nesting -= 1
+        end
+
+        def in_transaction?
+          @transaction_nesting.positive?
+        end
+
         def use_primary!
           @use_primary = true
         end

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

@@ -49,8 +49,16 @@ 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) { Gitlab::Database::LoadBalancing::Session.new }
+
+    before do
+      allow(Gitlab::Database::LoadBalancing::Session)
+        .to receive(:current)
+        .and_return(session)
+
+      allow(proxy.load_balancer)
+        .to receive(:primary_write_location)
+        .and_return('123/ABC')
     end
 
     it 'runs the transaction and any nested queries on the primary' do
@@ -60,15 +68,36 @@ describe Gitlab::Database::LoadBalancing::ConnectionProxy do
       allow(primary).to receive(:select)
 
       expect(proxy.load_balancer).to receive(:read_write)
-        .twice.and_yield(primary)
+        .twice
+        .and_yield(primary)
 
       # 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') }
 
-      expect(Gitlab::Database::LoadBalancing::Session.current.use_primary?)
-        .to eq(true)
+      expect(session.use_primary?).to eq(true)
+    end
+
+    it 'tracks the state of the transaction in the session' do
+      expect(proxy)
+        .to receive(:write_using_load_balancer)
+        .with(:transaction, [10], { sticky: true })
+
+      expect(session).to receive(:enter_transaction)
+      expect(session).to receive(:leave_transaction)
+
+      proxy.transaction(10)
+    end
+
+    it 'records the last write location' do
+      allow(proxy)
+        .to receive(:write_using_load_balancer)
+        .with(:transaction, [10], { sticky: true })
+
+      proxy.transaction(10)
+
+      expect(session.last_write_location).to eq('123/ABC')
     end
   end
 
@@ -91,62 +120,136 @@ describe Gitlab::Database::LoadBalancing::ConnectionProxy do
   end
 
   describe '#read_using_load_balancer' do
-    let(:session) { double(:session) }
+    let(:session) { Gitlab::Database::LoadBalancing::Session.new }
     let(:connection) { double(:connection) }
 
     before do
-      allow(Gitlab::Database::LoadBalancing::Session).to receive(:current)
+      allow(Gitlab::Database::LoadBalancing::Session)
+        .to receive(:current)
         .and_return(session)
     end
 
-    describe 'with a regular session' do
-      it 'uses a secondary' do
-        allow(session).to receive(:use_primary?).and_return(false)
+    it 'performs a read-only query' do
+      allow(proxy.load_balancer)
+        .to receive(:load_balancer_method_for_read)
+        .and_return(:read)
 
-        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
+      allow(proxy.load_balancer)
+        .to receive(:read)
+        .and_yield(connection)
 
-    describe 'with a session using the primary' do
-      it 'uses the primary' do
-        allow(session).to receive(:use_primary?).and_return(true)
+      expect(connection)
+        .to receive(:foo)
+        .with('foo')
 
-        expect(connection).to receive(:foo).with('foo')
-
-        expect(proxy.load_balancer).to receive(:read_write)
-          .and_yield(connection)
-
-        proxy.read_using_load_balancer(:foo, ['foo'])
-      end
+      proxy.read_using_load_balancer(:foo, ['foo'])
     end
   end
 
   describe '#write_using_load_balancer' do
-    let(:session) { double(:session) }
+    let(:session) { Gitlab::Database::LoadBalancing::Session.new }
     let(:connection) { double(:connection) }
 
     before do
-      allow(Gitlab::Database::LoadBalancing::Session).to receive(:current)
+      allow(Gitlab::Database::LoadBalancing::Session)
+        .to receive(:current)
         .and_return(session)
+
+      allow(proxy.load_balancer)
+        .to receive(:primary_write_location)
+        .and_return('123/ABC')
+
+      allow(proxy.load_balancer)
+        .to receive(:read_write)
+        .and_yield(connection)
+
+      allow(connection)
+        .to receive(:foo)
+        .with('foo')
     end
 
     it 'it uses but does not stick to the primary when sticking is disabled' do
-      expect(proxy.load_balancer).to receive(:read_write).and_yield(connection)
-      expect(connection).to receive(:foo).with('foo')
       expect(session).not_to receive(:write!)
 
       proxy.write_using_load_balancer(:foo, ['foo'])
     end
 
     it 'sticks to the primary when sticking is enabled' do
-      expect(proxy.load_balancer).to receive(:read_write).and_yield(connection)
-      expect(connection).to receive(:foo).with('foo')
       expect(session).to receive(:write!)
 
       proxy.write_using_load_balancer(:foo, ['foo'], sticky: true)
     end
+
+    it 'tracks the last write location' do
+      proxy.write_using_load_balancer(:foo, ['foo'], sticky: true)
+
+      expect(session.last_write_location).to be_instance_of(String)
+    end
+
+    it 'does not track the last write location inside a transaction' do
+      session.enter_transaction
+
+      proxy.write_using_load_balancer(:foo, ['foo'], sticky: true)
+
+      expect(session.last_write_location).to be_nil
+    end
+
+    it 'does not track the last write location if sticking is not needed' do
+      proxy.write_using_load_balancer(:foo, ['foo'], sticky: false)
+
+      expect(session.last_write_location).to be_nil
+    end
+  end
+
+  describe '#load_balancer_method_for_read' do
+    let(:session) { Gitlab::Database::LoadBalancing::Session.new }
+
+    before do
+      allow(Gitlab::Database::LoadBalancing::Session)
+        .to receive(:current)
+        .and_return(session)
+    end
+
+    context 'when using the primary' do
+      before do
+        session.use_primary!
+      end
+
+      it 'returns :read_write when in a transaction' do
+        session.enter_transaction
+
+        expect(proxy.load_balancer_method_for_read).to eq(:read_write)
+      end
+
+      it 'returns :read_write if the secondaries are not in sync' do
+        session.last_write_location = '123/ABC'
+
+        allow(proxy.load_balancer)
+          .to receive(:all_caught_up?)
+          .with('123/ABC')
+          .and_return(false)
+
+        expect(proxy.load_balancer_method_for_read).to eq(:read_write)
+      end
+
+      it 'returns :read if all secondaries are in sync' do
+        session.last_write_location = '123/ABC'
+
+        allow(proxy.load_balancer)
+          .to receive(:all_caught_up?)
+          .with('123/ABC')
+          .and_return(true)
+
+        expect(proxy.load_balancer_method_for_read).to eq(:read)
+
+        expect(session.use_primary?).to eq(false)
+      end
+    end
+
+    context 'when using a secondary' do
+      it 'returns :read' do
+        expect(proxy.load_balancer_method_for_read).to eq(:read)
+      end
+    end
   end
 end

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

@@ -51,4 +51,33 @@ describe Gitlab::Database::LoadBalancing::Session do
       expect(instance.performed_write?).to eq(true)
     end
   end
+
+  describe '#reset!' do
+    it 'switches the session back to read from a secondary' do
+      instance = described_class.new
+
+      instance.use_primary!
+      instance.last_write_location = 'foo'
+      instance.reset!
+
+      expect(instance.use_primary?).to eq(false)
+      expect(instance.last_write_location).to be_nil
+    end
+  end
+
+  describe 'transaction nesting' do
+    it 'supports tracking of transaction states' do
+      instance = described_class.new
+
+      expect(instance).not_to be_in_transaction
+
+      instance.enter_transaction
+
+      expect(instance).to be_in_transaction
+
+      instance.leave_transaction
+
+      expect(instance).not_to be_in_transaction
+    end
+  end
 end