From 8530a6e23c40a6771a969e1e2aa5ec726e28ddb3 Mon Sep 17 00:00:00 2001
From: Jean Boussier <jean.boussier@gmail.com>
Date: Mon, 30 Sep 2024 11:16:37 -0400
Subject: [PATCH] Make the query cache executor transactional fixtures aware

Fix: https://github.com/rails/rails/issues/52973

Checking for `transaction_open?` only cause the Puma thread to
not release the connection if we're still inside a fixture transaction.

We need to check `joinable?` too.
---
 activerecord/CHANGELOG.md                     | 10 ++++++++
 .../abstract/connection_pool.rb               | 25 +++++++++++++++++++
 activerecord/lib/active_record/query_cache.rb |  4 ---
 activerecord/lib/active_record/railtie.rb     |  1 +
 .../test/cases/connection_management_test.rb  | 14 +++++++++++
 5 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/activerecord/CHANGELOG.md b/activerecord/CHANGELOG.md
index 06a857de75278..b9b2594d9197a 100644
--- a/activerecord/CHANGELOG.md
+++ b/activerecord/CHANGELOG.md
@@ -1,3 +1,13 @@
+*   Properly release pinned connections with non joinable connections
+
+    Fixes #52973
+
+    When running system tests with transactional fixtures on, it could happen that
+    the connection leased by the Puma thread wouldn't be properly released back to the pool,
+    causing "Cannot expire connection, it is owned by a different thread" errors in later tests.
+
+    *Jean Boussier*
+
 *   Make Float distinguish between `float4` and `float8` in PostgreSQL.
 
     Fixes #52742
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 97d7f482deb81..226c79dee30eb 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
@@ -184,6 +184,31 @@ def clear
         end
       end
 
+      module ExecutorHooks # :nodoc:
+        class << self
+          def run
+            # noop
+          end
+
+          def complete(_)
+            ActiveRecord::Base.connection_handler.each_connection_pool do |pool|
+              if (connection = pool.active_connection?)
+                transaction = connection.current_transaction
+                if transaction.closed? || !transaction.joinable?
+                  pool.release_connection
+                end
+              end
+            end
+          end
+        end
+      end
+
+      class << self
+        def install_executor_hooks(executor = ActiveSupport::Executor)
+          executor.register_hook(ExecutorHooks)
+        end
+      end
+
       include MonitorMixin
       prepend QueryCache::ConnectionPoolConfiguration
       include ConnectionAdapters::AbstractPool
diff --git a/activerecord/lib/active_record/query_cache.rb b/activerecord/lib/active_record/query_cache.rb
index 89a022848c7bb..db0b2f3c5a3cc 100644
--- a/activerecord/lib/active_record/query_cache.rb
+++ b/activerecord/lib/active_record/query_cache.rb
@@ -43,10 +43,6 @@ def self.complete(pools)
         pool.disable_query_cache!
         pool.clear_query_cache
       end
-
-      ActiveRecord::Base.connection_handler.each_connection_pool do |pool|
-        pool.release_connection if pool.active_connection? && !pool.lease_connection.transaction_open?
-      end
     end
 
     def self.install_executor_hooks(executor = ActiveSupport::Executor)
diff --git a/activerecord/lib/active_record/railtie.rb b/activerecord/lib/active_record/railtie.rb
index e4ae6bb51c626..5dae1d95bb6f2 100644
--- a/activerecord/lib/active_record/railtie.rb
+++ b/activerecord/lib/active_record/railtie.rb
@@ -312,6 +312,7 @@ class Railtie < Rails::Railtie # :nodoc:
     initializer "active_record.set_executor_hooks" do
       ActiveRecord::QueryCache.install_executor_hooks
       ActiveRecord::AsynchronousQueriesTracker.install_executor_hooks
+      ActiveRecord::ConnectionAdapters::ConnectionPool.install_executor_hooks
     end
 
     initializer "active_record.add_watchable_files" do |app|
diff --git a/activerecord/test/cases/connection_management_test.rb b/activerecord/test/cases/connection_management_test.rb
index 49aa9257488bb..5885dde168b12 100644
--- a/activerecord/test/cases/connection_management_test.rb
+++ b/activerecord/test/cases/connection_management_test.rb
@@ -50,6 +50,19 @@ def test_connections_are_cleared_after_body_close
         assert_not ActiveRecord::Base.connection_handler.active_connections?(:all)
       end
 
+      test "connections are cleared even if inside a non-joinable transaction" do
+        ActiveRecord::Base.connection_pool.pin_connection!(Thread.current)
+        Thread.new do
+          assert ActiveRecord::Base.lease_connection
+          assert ActiveRecord::Base.connection_handler.active_connections?(:all)
+          _, _, body = @management.call(@env)
+          body.close
+          assert_not ActiveRecord::Base.connection_handler.active_connections?(:all)
+        end.join
+      ensure
+        ActiveRecord::Base.connection_pool.unpin_connection!
+      end
+
       def test_active_connections_are_not_cleared_on_body_close_during_transaction
         ActiveRecord::Base.transaction do
           _, _, body = @management.call(@env)
@@ -123,6 +136,7 @@ def executor
           @executor ||= Class.new(ActiveSupport::Executor).tap do |exe|
             ActiveRecord::QueryCache.install_executor_hooks(exe)
             ActiveRecord::AsynchronousQueriesTracker.install_executor_hooks(exe)
+            ActiveRecord::ConnectionAdapters::ConnectionPool.install_executor_hooks(exe)
           end
         end