Merge branch 'less-aggressive-db-sticking' into 'master'

Use less aggressive sticking for DB load balancing

Closes #6041

See merge request gitlab-org/gitlab-ee!5732

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)
-
-        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
-
-    describe 'with a session using the primary' do
-      it 'uses the primary' do
-        allow(session).to receive(:use_primary?).and_return(true)
+    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_write)
+      allow(proxy.load_balancer)
+        .to receive(:read)
         .and_yield(connection)
 
+      expect(connection)
+        .to receive(:foo)
+        .with('foo')
+
       proxy.read_using_load_balancer(:foo, ['foo'])
     end
   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