Make it possible to ignore load balancing sticky write

This commit makes it possible to avoid upgrading ActiveRecord connection
proxy and switching to the primary database when a write has happened.

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

@@ -19,9 +19,14 @@ module Gitlab
           RequestStore.delete(CACHE_KEY)
         end
 
+        def self.without_sticky_writes(&block)
+          current.ignore_writes(&block)
+        end
+
         def initialize
           @use_primary = false
           @performed_write = false
+          @ignore_writes = false
         end
 
         def use_primary?
@@ -42,8 +47,20 @@ module Gitlab
           @use_primary = used_primary || @performed_write
         end
 
+        def ignore_writes(&block)
+          @ignore_writes = true
+
+          yield
+        ensure
+          @ignore_writes = false
+        end
+
+
         def write!
           @performed_write = true
+
+          return if @ignore_writes
+
           use_primary!
         end
 

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

@@ -22,6 +22,28 @@ RSpec.describe Gitlab::Database::LoadBalancing::Session do
     end
   end
 
+  describe '.without_sticky_writes' do
+    it 'ignores sticky write events sent by a connection proxy' do
+      described_class.without_sticky_writes do
+        described_class.current.write!
+      end
+
+      session = described_class.current
+
+      expect(session).not_to be_using_primary
+    end
+
+    it 'still is aware of write that happened' do
+      described_class.without_sticky_writes do
+        described_class.current.write!
+      end
+
+      session = described_class.current
+
+      expect(session.performed_write?).to be true
+    end
+  end
+
   describe '#use_primary?' do
     it 'returns true when the primary should be used' do
       instance = described_class.new
@@ -96,4 +118,24 @@ RSpec.describe Gitlab::Database::LoadBalancing::Session do
       expect(instance.performed_write?).to eq(true)
     end
   end
+
+  describe '#ignore_writes' do
+    it 'ignores write events' do
+      instance = described_class.new
+
+      instance.ignore_writes { instance.write! }
+
+      expect(instance).not_to be_using_primary
+      expect(instance.performed_write?).to eq true
+    end
+
+    it 'does not prevent using primary if an exception is raised' do
+      instance = described_class.new
+
+      instance.ignore_writes { raise ArgumentError } rescue ArgumentError
+      instance.write!
+
+      expect(instance).to be_using_primary
+    end
+  end
 end