Merge pull request rails#51097 from Shopify/test-connection-leak

ActiveRecord::TestCase: reap all leaked connection on teardown

Author: byroot

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

@@ -252,13 +252,15 @@ def release_connection(owner_thread = ActiveSupport::IsolatedExecutionState.cont
       # connection will be properly returned to the pool by the code that checked
       # it out.
       def with_connection
-        unless conn = @thread_cached_conns[ActiveSupport::IsolatedExecutionState.context]
-          conn = connection
-          fresh_connection = true
+        if conn = @thread_cached_conns[ActiveSupport::IsolatedExecutionState.context]
+          yield conn
+        else
+          begin
+            yield connection
+          ensure
+            release_connection
+          end
         end
-        yield conn
-      ensure
-        release_connection if fresh_connection
       end
 
       # Returns true if a connection has already been opened.
@@ -573,6 +575,8 @@ def with_exclusively_acquired_all_connections(raise_on_acquisition_timeout = tru
 
         def attempt_to_checkout_all_existing_connections(raise_on_acquisition_timeout = true)
           collected_conns = synchronize do
+            reap # No need to wait for dead owners
+
             # account for our own connections
             @connections.select { |conn| conn.owner == ActiveSupport::IsolatedExecutionState.context }
           end

activerecord/test/cases/connection_pool_test.rb

@@ -978,6 +978,11 @@ def join(timeout = nil)
         def terminate
           nil
         end
+
+        unless method_defined?(:kill) # RUBY_VERSION <= "3.3"
+          def kill
+          end
+        end
       end
 
       def setup

activerecord/test/cases/test_case.rb

@@ -33,6 +33,56 @@ class TestCase < ActiveSupport::TestCase # :nodoc:
     self.use_instantiated_fixtures = false
     self.use_transactional_tests = true
 
+    def after_teardown
+      super
+      check_connection_leaks
+    end
+
+    def check_connection_leaks
+      return if in_memory_db?
+
+      # Make sure tests didn't leave a connection owned by some background thread
+      # which could lead to some slow wait in a subsequent thread.
+      leaked_conn = []
+      ActiveRecord::Base.connection_handler.each_connection_pool do |pool|
+        # Ensure all in flights tasks are completed.
+        # Otherwise they may still hold a connection.
+        if pool.async_executor
+          if pool.async_executor.scheduled_task_count != pool.async_executor.completed_task_count
+            pool.connections.each do |conn|
+              if conn.in_use? && conn.owner != Fiber.current && conn.owner != Thread.current
+                if conn.owner.respond_to?(:join)
+                  conn.owner&.join(0.5)
+                end
+              end
+            end
+          end
+        end
+
+        pool.reap
+        pool.connections.each do |conn|
+          if conn.in_use?
+            if conn.owner != Fiber.current && conn.owner != Thread.current
+              leaked_conn << [conn.owner, conn.owner.backtrace]
+              conn.owner&.kill
+            end
+            conn.steal!
+            pool.checkin(conn)
+          end
+        end
+      end
+
+      if leaked_conn.size > 0
+        puts "Found #{leaked_conn.size} leaked connections"
+        leaked_conn.each do |owner, backtrace|
+          puts "owner: #{owner}"
+          puts "backtrace:\n#{backtrace}"
+          puts
+        end
+        raise "Found #{leaked_conn.size} leaked connection after #{self.class.name}##{name}"
+      end
+    end
+
     def create_fixtures(*fixture_set_names, &block)
       ActiveRecord::FixtureSet.create_fixtures(ActiveRecord::TestCase.fixture_paths, fixture_set_names, fixture_class_names, &block)
     end