Make connection stealing more explicit

3 files changed, 27 insertions, 6 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 98cf97910e..bf70d4dc59 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
@@ -415,7 +415,10 @@ module ActiveRecord
         with_exclusively_acquired_all_connections(raise_on_acquisition_timeout) do
           synchronize do
             @connections.each do |conn|
-              checkin conn
+              if conn.in_use?
+                conn.steal!
+                checkin conn
+              end
               conn.disconnect!
             end
             @connections = []
@@ -447,7 +450,10 @@ module ActiveRecord
         with_exclusively_acquired_all_connections(raise_on_acquisition_timeout) do
           synchronize do
             @connections.each do |conn|
-              checkin conn
+              if conn.in_use?
+                conn.steal!
+                checkin conn
+              end
               conn.disconnect! if conn.requires_reloading?
             end
             @connections.delete_if(&:requires_reloading?)
@@ -557,8 +563,7 @@ module ActiveRecord
           @connections.select do |conn|
             conn.in_use? && !conn.owner.alive?
           end.each do |conn|
-            conn.expire
-            conn.lease
+            conn.steal!
           end
         end
 
diff --git a/activerecord/lib/active_record/connection_adapters/abstract_adapter.rb b/activerecord/lib/active_record/connection_adapters/abstract_adapter.rb
index 5140238571..5747e4d1ee 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract_adapter.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract_adapter.rb
@@ -185,7 +185,7 @@ module ActiveRecord
       # this method must only be called while holding connection pool's mutex
       def expire
         if in_use?
-          if @owner != Thread.current && @owner.alive?
+          if @owner != Thread.current
             raise ActiveRecordError, "Cannot expire connection, " <<
               "it is owned by a different thread: #{@owner}. " <<
               "Current thread: #{Thread.current}."
@@ -197,6 +197,19 @@ module ActiveRecord
         end
       end
 
+      # this method must only be called while holding connection pool's mutex (and a desire for segfaults)
+      def steal! # :nodoc:
+        if in_use?
+          if @owner != Thread.current
+            pool.send :remove_connection_from_thread_cache, self, @owner
+
+            @owner = Thread.current
+          end
+        else
+          raise ActiveRecordError, 'Cannot steal connection, it is not currently leased.'
+        end
+      end
+
       def unprepared_statement
         old_prepared_statements, @prepared_statements = @prepared_statements, false
         yield
diff --git a/activerecord/test/cases/connection_pool_test.rb b/activerecord/test/cases/connection_pool_test.rb
index a45ee281c7..09e7848bda 100644
--- a/activerecord/test/cases/connection_pool_test.rb
+++ b/activerecord/test/cases/connection_pool_test.rb
@@ -151,7 +151,7 @@ module ActiveRecord
 
         assert_equal 1, active_connections(@pool).size
       ensure
-        @pool.connections.each(&:close)
+        @pool.connections.each { |conn| conn.close if conn.in_use? }
       end
 
       def test_remove_connection
@@ -432,6 +432,9 @@ module ActiveRecord
             Thread.new { @pool.send(group_action_method) }.join
             # assert connection has been forcefully taken away from us
             assert_not @pool.active_connection?
+
+            # make a new connection for with_connection to clean up
+            @pool.connection
           end
         end
       end