Merge pull request #14360 from matthewd/thread_aware_reaper

Track owning thread for connection pool

8 files changed, 103 insertions, 69 deletions

diff --git a/activerecord/CHANGELOG.md b/activerecord/CHANGELOG.md
index b648eed06a..f5b8a3145d 100644
--- a/activerecord/CHANGELOG.md
+++ b/activerecord/CHANGELOG.md
@@ -1,3 +1,11 @@
+*   Reap connections that were checked out by now-dead threads, instead
+    of waiting until they disconnect by themselves. Before this change,
+    a suitably constructed series of short-lived threads could starve
+    the connection pool, without ever having more than a couple alive at
+    the same time.
+
+    *Matthew Draper*
+
 *   `where.not` adds `references` for `includes` like normal `where` calls do.
 
     Fixes #14406.
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 fd185de589..db80c0faee 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
@@ -58,13 +58,11 @@ module ActiveRecord
     # * +checkout_timeout+: number of seconds to block and wait for a connection
     #   before giving up and raising a timeout error (default 5 seconds).
     # * +reaping_frequency+: frequency in seconds to periodically run the
-    #   Reaper, which attempts to find and close dead connections, which can
-    #   occur if a programmer forgets to close a connection at the end of a
-    #   thread or a thread dies unexpectedly. (Default nil, which means don't
-    #   run the Reaper).
-    # * +dead_connection_timeout+: number of seconds from last checkout
-    #   after which the Reaper will consider a connection reapable. (default
-    #   5 seconds).
+    #   Reaper, which attempts to find and recover connections from dead
+    #   threads, which can occur if a programmer forgets to close a
+    #   connection at the end of a thread or a thread dies unexpectedly.
+    #   Regardless of this setting, the Reaper will be invoked before every
+    #   blocking wait. (Default nil, which means don't schedule the Reaper).
     class ConnectionPool
       # Threadsafe, fair, FIFO queue.  Meant to be used by ConnectionPool
       # with which it shares a Monitor.  But could be a generic Queue.
@@ -222,7 +220,7 @@ module ActiveRecord
 
       include MonitorMixin
 
-      attr_accessor :automatic_reconnect, :checkout_timeout, :dead_connection_timeout
+      attr_accessor :automatic_reconnect, :checkout_timeout
       attr_reader :spec, :connections, :size, :reaper
 
       # Creates a new ConnectionPool object. +spec+ is a ConnectionSpecification
@@ -237,7 +235,6 @@ module ActiveRecord
         @spec = spec
 
         @checkout_timeout = spec.config[:checkout_timeout] || 5
-        @dead_connection_timeout = spec.config[:dead_connection_timeout] || 5
         @reaper  = Reaper.new self, spec.config[:reaping_frequency]
         @reaper.run
 
@@ -361,11 +358,13 @@ module ActiveRecord
       # calling +checkout+ on this pool.
       def checkin(conn)
         synchronize do
+          owner = conn.owner
+
           conn.run_callbacks :checkin do
             conn.expire
           end
 
-          release conn
+          release conn, owner
 
           @available.add conn
         end
@@ -378,22 +377,28 @@ module ActiveRecord
           @connections.delete conn
           @available.delete conn
 
-          # FIXME: we might want to store the key on the connection so that removing
-          # from the reserved hash will be a little easier.
-          release conn
+          release conn, conn.owner
 
           @available.add checkout_new_connection if @available.any_waiting?
         end
       end
 
-      # Removes dead connections from the pool.  A dead connection can occur
-      # if a programmer forgets to close a connection at the end of a thread
+      # Recover lost connections for the pool.  A lost connection can occur if
+      # a programmer forgets to checkin a connection at the end of a thread
       # or a thread dies unexpectedly.
       def reap
-        synchronize do
-          stale = Time.now - @dead_connection_timeout
-          connections.dup.each do |conn|
-            if conn.in_use? && stale > conn.last_use && !conn.active_threadsafe?
+        stale_connections = synchronize do
+          @connections.select do |conn|
+            conn.in_use? && !conn.owner.alive?
+          end
+        end
+
+        stale_connections.each do |conn|
+          synchronize do
+            if conn.active?
+              conn.reset!
+              checkin conn
+            else
               remove conn
             end
           end
@@ -415,20 +420,15 @@ module ActiveRecord
         elsif @connections.size < @size
           checkout_new_connection
         else
+          reap
           @available.poll(@checkout_timeout)
         end
       end
 
-      def release(conn)
-        thread_id = if @reserved_connections[current_connection_id] == conn
-          current_connection_id
-        else
-          @reserved_connections.keys.find { |k|
-            @reserved_connections[k] == conn
-          }
-        end
+      def release(conn, owner)
+        thread_id = owner.object_id
 
-        @reserved_connections.delete thread_id if thread_id
+        @reserved_connections.delete thread_id
       end
 
       def new_connection
diff --git a/activerecord/lib/active_record/connection_adapters/abstract_adapter.rb b/activerecord/lib/active_record/connection_adapters/abstract_adapter.rb
index 496153af83..c7e5a18f02 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract_adapter.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract_adapter.rb
@@ -71,8 +71,8 @@ module ActiveRecord
       define_callbacks :checkout, :checkin
 
       attr_accessor :visitor, :pool
-      attr_reader :schema_cache, :last_use, :in_use, :logger
-      alias :in_use? :in_use
+      attr_reader :schema_cache, :owner, :logger
+      alias :in_use? :owner
 
       def self.type_cast_config_to_integer(config)
         if config =~ SIMPLE_INT
@@ -94,9 +94,8 @@ module ActiveRecord
         super()
 
         @connection          = connection
-        @in_use              = false
+        @owner               = nil
         @instrumenter        = ActiveSupport::Notifications.instrumenter
-        @last_use            = false
         @logger              = logger
         @pool                = pool
         @schema_cache        = SchemaCache.new self
@@ -114,9 +113,8 @@ module ActiveRecord
 
       def lease
         synchronize do
-          unless in_use
-            @in_use   = true
-            @last_use = Time.now
+          unless in_use?
+            @owner = Thread.current
           end
         end
       end
@@ -127,7 +125,7 @@ module ActiveRecord
       end
 
       def expire
-        @in_use = false
+        @owner = nil
       end
 
       def unprepared_visitor
@@ -262,12 +260,6 @@ module ActiveRecord
       def active?
       end
 
-      # Adapter should redefine this if it needs a threadsafe way to approximate
-      # if the connection is active
-      def active_threadsafe?
-        active?
-      end
-
       # Disconnects from the database if already connected, and establishes a
       # new connection with the database. Implementors should call super if they
       # override the default implementation.
diff --git a/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb b/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb
index e5f7913c70..bcad9f30d7 100644
--- a/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb
+++ b/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb
@@ -603,10 +603,6 @@ module ActiveRecord
         false
       end
 
-      def active_threadsafe?
-        @connection.connect_poll != PG::PGRES_POLLING_FAILED
-      end
-
       # Close then reopen the connection.
       def reconnect!
         super
@@ -616,7 +612,12 @@ module ActiveRecord
 
       def reset!
         clear_cache!
-        super
+        reset_transaction
+        unless @connection.transaction_status == ::PG::PQTRANS_IDLE
+          @connection.query 'ROLLBACK'
+        end
+        @connection.query 'DISCARD ALL'
+        configure_connection
       end
 
       # Disconnects from the database if already connected. Otherwise, this
diff --git a/activerecord/test/cases/adapters/postgresql/connection_test.rb b/activerecord/test/cases/adapters/postgresql/connection_test.rb
index 4715fa002d..aa4996133f 100644
--- a/activerecord/test/cases/adapters/postgresql/connection_test.rb
+++ b/activerecord/test/cases/adapters/postgresql/connection_test.rb
@@ -45,6 +45,37 @@ module ActiveRecord
       assert_equal 'off', expect
     end
 
+    def test_reset
+      @connection.query('ROLLBACK')
+      @connection.query('SET geqo TO off')
+
+      # Verify the setting has been applied.
+      expect = @connection.query('show geqo').first.first
+      assert_equal 'off', expect
+
+      @connection.reset!
+
+      # Verify the setting has been cleared.
+      expect = @connection.query('show geqo').first.first
+      assert_equal 'on', expect
+    end
+
+    def test_reset_with_transaction
+      @connection.query('ROLLBACK')
+      @connection.query('SET geqo TO off')
+
+      # Verify the setting has been applied.
+      expect = @connection.query('show geqo').first.first
+      assert_equal 'off', expect
+
+      @connection.query('BEGIN')
+      @connection.reset!
+
+      # Verify the setting has been cleared.
+      expect = @connection.query('show geqo').first.first
+      assert_equal 'on', expect
+    end
+
     def test_tables_logs_name
       @connection.tables('hello')
       assert_equal 'SCHEMA', @subscriber.logged[0][1]
diff --git a/activerecord/test/cases/connection_adapters/abstract_adapter_test.rb b/activerecord/test/cases/connection_adapters/abstract_adapter_test.rb
index eb2fe5639b..deed226eab 100644
--- a/activerecord/test/cases/connection_adapters/abstract_adapter_test.rb
+++ b/activerecord/test/cases/connection_adapters/abstract_adapter_test.rb
@@ -29,12 +29,6 @@ module ActiveRecord
         assert_not adapter.lease, 'should not lease adapter'
       end
 
-      def test_last_use
-        assert_not adapter.last_use
-        adapter.lease
-        assert adapter.last_use
-      end
-
       def test_expire_mutates_in_use
         assert adapter.lease, 'lease adapter'
         assert adapter.in_use?, 'adapter is in use'
diff --git a/activerecord/test/cases/connection_pool_test.rb b/activerecord/test/cases/connection_pool_test.rb
index c4108ad7f6..6300011f4a 100644
--- a/activerecord/test/cases/connection_pool_test.rb
+++ b/activerecord/test/cases/connection_pool_test.rb
@@ -124,7 +124,6 @@ module ActiveRecord
         @pool.checkout
         @pool.checkout
         @pool.checkout
-        @pool.dead_connection_timeout = 0
 
         connections = @pool.connections.dup
 
@@ -134,21 +133,25 @@ module ActiveRecord
       end
 
       def test_reap_inactive
+        ready = false
         @pool.checkout
-        @pool.checkout
-        @pool.checkout
-        @pool.dead_connection_timeout = 0
-
-        connections = @pool.connections.dup
-        connections.each do |conn|
-          conn.extend(Module.new { def active_threadsafe?; false; end; })
+        child = Thread.new do
+          @pool.checkout
+          @pool.checkout
+          ready = true
+          Thread.stop
         end
+        Thread.pass until ready
+
+        assert_equal 3, active_connections(@pool).size
 
+        child.terminate
+        child.join
         @pool.reap
 
-        assert_equal 0, @pool.connections.length
+        assert_equal 1, active_connections(@pool).size
       ensure
-        connections.each(&:close)
+        @pool.connections.each(&:close)
       end
 
       def test_remove_connection
diff --git a/activerecord/test/cases/reaper_test.rb b/activerecord/test/cases/reaper_test.rb
index 015c37cce8..f52fd22489 100644
--- a/activerecord/test/cases/reaper_test.rb
+++ b/activerecord/test/cases/reaper_test.rb
@@ -63,17 +63,22 @@ module ActiveRecord
         spec.config[:reaping_frequency] = 0.0001
 
         pool = ConnectionPool.new spec
-        pool.dead_connection_timeout = 0
 
-        conn = pool.checkout
-        count = pool.connections.length
+        conn = nil
+        child = Thread.new do
+          conn = pool.checkout
+          Thread.stop
+        end
+        Thread.pass while conn.nil?
+
+        assert conn.in_use?
 
-        conn.extend(Module.new { def active_threadsafe?; false; end; })
+        child.terminate
 
-        while count == pool.connections.length
+        while conn.in_use?
           Thread.pass
         end
-        assert_equal(count - 1, pool.connections.length)
+        assert !conn.in_use?
       end
     end
   end