Merge pull request rails#52037 from palkan/feat/active-record/nested-pinning

[ActiveRecord] Support nested connection pinning

Author: byroot

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

@@ -253,6 +253,7 @@ def initialize(pool_config)
 
         @available = ConnectionLeasingQueue.new self
         @pinned_connection = nil
+        @pinned_connections_depth = 0
 
         @async_executor = build_async_executor
 
@@ -311,9 +312,9 @@ def connection
       end
 
       def pin_connection!(lock_thread) # :nodoc:
-        raise "There is already a pinned connection" if @pinned_connection
+        @pinned_connection ||= (connection_lease&.connection || checkout)
+        @pinned_connections_depth += 1
 
-        @pinned_connection = (connection_lease&.connection || checkout)
         # Any leased connection must be in @connections otherwise
         # some methods like #connected? won't behave correctly
         unless @connections.include?(@pinned_connection)
@@ -330,16 +331,22 @@ def unpin_connection! # :nodoc:
 
         clean = true
         @pinned_connection.lock.synchronize do
-          connection, @pinned_connection = @pinned_connection, nil
+          @pinned_connections_depth -= 1
+          connection = @pinned_connection
+          @pinned_connection = nil if @pinned_connections_depth.zero?
+
           if connection.transaction_open?
             connection.rollback_transaction
           else
             # Something committed or rolled back the transaction
             clean = false
             connection.reset!
           end
-          connection.lock_thread = nil
-          checkin(connection)
+
+          if @pinned_connection.nil?
+            connection.lock_thread = nil
+            checkin(connection)
+          end
         end
 
         clean

activerecord/test/cases/connection_pool_test.rb

@@ -934,6 +934,56 @@ def test_unpin_connection_returns_whether_transaction_has_been_rolledback
         assert_equal false, @pool.unpin_connection!
       end
 
+      def test_pin_connection_nesting
+        assert_instance_of NullTransaction, @pool.lease_connection.current_transaction
+        @pool.pin_connection!(true)
+        assert_instance_of RealTransaction, @pool.lease_connection.current_transaction
+        @pool.pin_connection!(true)
+        assert_instance_of SavepointTransaction, @pool.lease_connection.current_transaction
+        @pool.unpin_connection!
+        assert_instance_of RealTransaction, @pool.lease_connection.current_transaction
+        @pool.unpin_connection!
+        assert_instance_of NullTransaction, @pool.lease_connection.current_transaction
+
+        assert_raises(RuntimeError, match: /There isn't a pinned connection/) do
+          @pool.unpin_connection!
+        end
+      end
+
+      def test_pin_connection_nesting_lock
+        assert_equal ActiveSupport::Concurrency::NullLock, @pool.lease_connection.lock
+
+        @pool.pin_connection!(true)
+        actual_lock = @pool.lease_connection.lock
+        assert_not_equal ActiveSupport::Concurrency::NullLock, actual_lock
+
+        @pool.pin_connection!(false)
+        assert_same actual_lock, @pool.lease_connection.lock
+
+        @pool.unpin_connection!
+        assert_same actual_lock, @pool.lease_connection.lock
+
+        @pool.unpin_connection!
+        assert_equal ActiveSupport::Concurrency::NullLock, @pool.lease_connection.lock
+      end
+
+      def test_pin_connection_nesting_lock_inverse
+        assert_equal ActiveSupport::Concurrency::NullLock, @pool.lease_connection.lock
+
+        @pool.pin_connection!(false)
+        assert_equal ActiveSupport::Concurrency::NullLock, @pool.lease_connection.lock
+
+        @pool.pin_connection!(true)
+        actual_lock = @pool.lease_connection.lock
+        assert_not_equal ActiveSupport::Concurrency::NullLock, actual_lock
+
+        @pool.unpin_connection!
+        assert_same actual_lock, @pool.lease_connection.lock # The lock persist until full unpin
+
+        @pool.unpin_connection!
+        assert_equal ActiveSupport::Concurrency::NullLock, @pool.lease_connection.lock
+      end
+
       private
         def active_connections(pool)
           pool.connections.find_all(&:in_use?)