Handle read-only transaction state

Issue https://gitlab.com/gitlab-org/gitlab/-/issues/322133

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
@@ -58,10 +61,14 @@ module Gitlab
 
         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.
@@ -90,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
@@ -101,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/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb

@@ -116,44 +116,69 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
     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
 
-      it 'runs the transaction and any nested queries on the replica' do
-        replica = double(:connection)
+      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!)
 
-        allow(replica).to receive(:transaction).and_yield
-        allow(replica).to receive(:select)
+          proxy.transaction { proxy.select('true') }
+        end
+      end
 
-        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!)
+      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)
 
-        proxy.transaction { proxy.select('true') }
+          expect do
+            proxy.transaction { proxy.insert('something') }
+          end.to raise_error(Gitlab::Database::LoadBalancing::ConnectionProxy::WriteInsideReadOnlyTransactionError)
+        end
       end
     end
 
     context 'session does not prefer to use a replica' do
+      let(:primary) { double(:connection) }
+
       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
 
-      it 'runs the transaction and any nested queries on the primary and stick to it' do
-        primary = double(:connection)
+      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!)
 
-        allow(primary).to receive(:transaction).and_yield
-        allow(primary).to receive(:select)
+          proxy.transaction { proxy.select('true') }
+        end
+      end
 
-        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!)
+      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
 
-        proxy.transaction { proxy.select('true') }
+          proxy.transaction { proxy.insert('something') }
+        end
       end
     end
   end

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

@@ -415,6 +415,43 @@ RSpec.describe Gitlab::Database::LoadBalancing do
   # instrumentaiton) 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
+    shared_context 'LoadBalancing setup' do
+      let!(:license) { create(:license, plan: ::License::PREMIUM_PLAN) }
+      let(:hosts) { [ActiveRecord::Base.configurations["development"]['host']] }
+      let(:model) do
+        Class.new(ApplicationRecord) do
+          self.table_name = "load_balancing_test"
+        end
+      end
+
+      before do
+        ActiveRecord::Schema.define do
+          create_table :load_balancing_test, force: true do |t|
+            t.string :name, null: true
+          end
+        end
+        # 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
+        ActiveRecord::Schema.define do
+          drop_table :load_balancing_test, force: true
+        end
+      end
+    end
+
     where(:queries, :include_transaction, :expected_results) do
       [
         # Read methods
@@ -573,7 +610,7 @@ RSpec.describe Gitlab::Database::LoadBalancing do
           false, [:primary]
         ],
 
-        # use_replica_if_possible inside use_primary!
+        # use_replica_if_possible inside use_primary
         [
           -> {
             ::Gitlab::Database::LoadBalancing::Session.current.use_primary do
@@ -583,45 +620,36 @@ RSpec.describe Gitlab::Database::LoadBalancing do
             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
-      let!(:license) { create(:license, plan: ::License::PREMIUM_PLAN) }
-      let(:hosts) { [ActiveRecord::Base.configurations["development"]['host']] }
-      let(:model) do
-        Class.new(ApplicationRecord) do
-          self.table_name = "load_balancing_test"
-        end
-      end
-
-      before do
-        ActiveRecord::Schema.define do
-          create_table :load_balancing_test, force: true do |t|
-            t.string :name, null: true
-          end
-        end
-        # 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
-        ActiveRecord::Schema.define do
-          drop_table :load_balancing_test, force: true
-        end
-      end
+      include_context 'LoadBalancing setup'
 
       it 'redirects queries to the right roles' do
         roles = []
@@ -653,5 +681,20 @@ RSpec.describe Gitlab::Database::LoadBalancing do
         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