Invalidate transaction as early as possible

This should allow framework to avoid issuing `ROLLBACK` statements
for cases when db has already rolled back the transaction

Author: nvasilevski

activerecord/CHANGELOG.md

@@ -1,3 +1,12 @@
+*   Invalidate transaction as early as possible
+
+    After rescuing a `TransactionRollbackError` exception Rails invalidates transactions earlier in the flow
+    allowing the framework to skip issuing the `ROLLBACK` statement in more cases.
+    Only affects adapters that have `savepoint_errors_invalidate_transactions?` configured as `true`,
+    which at this point is only applicable to the `mysql2` adapter.
+
+    *Nikita Vasilevsky*
+
 *   Allow configuring columns list to be used in SQL queries issued by an `ActiveRecord::Base` object
 
     It is now possible to configure columns list that will be used to build an SQL query clauses when

activerecord/lib/active_record/connection_adapters/abstract/transaction.rb

@@ -457,10 +457,6 @@ def within_new_transaction(isolation: nil, joinable: true)
           ret
         rescue Exception => error
           if transaction
-            if error.is_a?(ActiveRecord::TransactionRollbackError) &&
-                @connection.savepoint_errors_invalidate_transactions?
-              transaction.state.invalidate!
-            end
             rollback_transaction
             after_failure_actions(transaction, error)
           end

activerecord/lib/active_record/connection_adapters/abstract_adapter.rb

@@ -972,6 +972,7 @@ def with_raw_connection(allow_retry: false, uses_transaction: true)
               result
             rescue => original_exception
               translated_exception = translate_exception_class(original_exception, nil, nil)
+              invalidate_transaction(translated_exception)
               retry_deadline_exceeded = deadline && deadline < Process.clock_gettime(Process::CLOCK_MONOTONIC)
 
               if !retry_deadline_exceeded && retries_available > 0
@@ -1007,6 +1008,13 @@ def retryable_connection_error?(exception)
           exception.is_a?(ConnectionNotEstablished) || exception.is_a?(ConnectionFailed)
         end
 
+        def invalidate_transaction(exception)
+          return unless exception.is_a?(TransactionRollbackError)
+          return unless savepoint_errors_invalidate_transactions?
+
+          current_transaction.state.invalidate! if current_transaction
+        end
+
         def retryable_query_error?(exception)
           # We definitely can't retry if we were inside a transaction that was instantly
           # rolled back by this error

activerecord/test/cases/adapter_test.rb

@@ -785,3 +785,27 @@ def test_advisory_locks_enabled?
     end
   end
 end
+
+if ActiveRecord::Base.connection.savepoint_errors_invalidate_transactions?
+  class InvalidateTransactionTest < ActiveRecord::TestCase
+    def test_invalidates_transaction_on_rollback_error
+      @invalidated = false
+      connection = ActiveRecord::Base.connection
+
+      connection.transaction do
+        connection.send(:with_raw_connection) do
+          raise ActiveRecord::Deadlocked, "made-up deadlock"
+        end
+
+      rescue ActiveRecord::Deadlocked => error
+        flunk("Rescuing wrong error") unless error.message == "made-up deadlock"
+
+        @invalidated = connection.current_transaction.state.invalidated?
+      end
+
+      # asserting outside of the transaction to make sure we actually reach the end of the test
+      # and perform the assertion
+      assert @invalidated
+    end
+  end
+end

activerecord/test/cases/adapters/mysql2/nested_deadlock_test.rb

@@ -78,6 +78,57 @@ class Sample < ActiveRecord::Base
       assert_predicate connection, :active?
     end
 
+    test "rollback exception is swallowed after a rollback" do
+      barrier = Concurrent::CyclicBarrier.new(2)
+      deadlocks = 0
+
+      s1 = Sample.create value: 1
+      s2 = Sample.create value: 2
+
+      thread = Thread.new do
+        Sample.transaction(requires_new: false) do
+          make_parent_transaction_dirty
+          Sample.transaction(requires_new: true) do
+            assert_current_transaction_is_savepoint_transaction
+            s1.lock!
+            barrier.wait
+            s2.update value: 4
+
+          rescue ActiveRecord::Deadlocked
+            deadlocks += 1
+
+            # This rollback is actually wrong as mysql automatically rollbacks the transaction
+            # which means we have nothing to rollback on the db side
+            # but we expect the framework to handle our mistake gracefully
+            raise ActiveRecord::Rollback
+          end
+
+          s2.update value: 10
+        end
+      end
+
+      begin
+        Sample.transaction(requires_new: false) do
+          make_parent_transaction_dirty
+          Sample.transaction(requires_new: true) do
+            assert_current_transaction_is_savepoint_transaction
+            s2.lock!
+            barrier.wait
+            s1.update value: 3
+          rescue ActiveRecord::Deadlocked
+            deadlocks += 1
+            raise ActiveRecord::Rollback
+          end
+          s1.update value: 10
+        end
+      ensure
+        thread.join
+      end
+
+      assert_equal 1, deadlocks, "deadlock is required for the test setup"
+      assert_equal [10, 10], Sample.pluck(:value)
+    end
+
     test "deadlock inside nested SavepointTransaction is recoverable" do
       barrier = Concurrent::CyclicBarrier.new(2)
       deadlocks = 0