Make connection stealing more explicit

2 files changed, 23 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 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