From 5f7dcb49e2dd7d22300eacc1104722c80adf74b2 Mon Sep 17 00:00:00 2001
From: Jean Boussier <jean.boussier@gmail.com>
Date: Sat, 8 Jun 2024 20:59:29 +0200
Subject: [PATCH] Merge pull request #52037 from
 palkan/feat/active-record/nested-pinning

[ActiveRecord] Support nested connection pinning
---
 .../abstract/connection_pool.rb               | 17 +++++--
 .../test/cases/connection_pool_test.rb        | 50 +++++++++++++++++++
 2 files changed, 62 insertions(+), 5 deletions(-)

diff --git a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
index b1a71df2bf243..14473887711b7 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
+++ b/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,7 +331,10 @@ 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
@@ -338,8 +342,11 @@ def unpin_connection! # :nodoc:
             clean = false
             connection.reset!
           end
-          connection.lock_thread = nil
-          checkin(connection)
+
+          if @pinned_connection.nil?
+            connection.lock_thread = nil
+            checkin(connection)
+          end
         end
 
         clean
diff --git a/activerecord/test/cases/connection_pool_test.rb b/activerecord/test/cases/connection_pool_test.rb
index 11fccfdd2f688..7e7b4bdd35659 100644
--- a/activerecord/test/cases/connection_pool_test.rb
+++ b/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?)