Merge pull request rails#51192 from Shopify/connection-leasing-2

Make `.connection` always return a permanently leased connection

Author: byroot

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

@@ -113,6 +113,90 @@ def db_config
     # * private methods that require being called in a +synchronize+ blocks
     #   are now explicitly documented
     class ConnectionPool
+      class Lease # :nodoc:
+        attr_accessor :connection, :sticky
+
+        def initialize
+          @connection = nil
+          @sticky = false
+        end
+
+        def release
+          conn = @connection
+          @connection = nil
+          @sticky = false
+          conn
+        end
+
+        def clear(connection)
+          if @connection == connection
+            @connection = nil
+            @sticky = false
+            true
+          else
+            false
+          end
+        end
+      end
+
+      if ObjectSpace.const_defined?(:WeakKeyMap) # RUBY_VERSION >= 3.3
+        WeakKeyMap = ::ObjectSpace::WeakKeyMap # :nodoc:
+      else
+        class WeakKeyMap # :nodoc:
+          def initialize
+            @map = ObjectSpace::WeakMap.new
+            @values = nil
+            @size = 0
+          end
+
+          alias_method :clear, :initialize
+
+          def [](key)
+            prune if @map.size != @size
+            @map[key]
+          end
+
+          def []=(key, value)
+            @map[key] = value
+            prune if @map.size != @size
+            value
+          end
+
+          def delete(key)
+            if value = self[key]
+              self[key] = nil
+              prune
+            end
+            value
+          end
+
+          private
+            def prune(force = false)
+              @values = @map.values
+              @size = @map.size
+            end
+        end
+      end
+
+      class LeaseRegistry # :nodoc:
+        def initialize
+          @mutex = Mutex.new
+          @map = WeakKeyMap.new
+        end
+
+        def [](context)
+          @mutex.synchronize do
+            @map[context] ||= Lease.new
+          end
+        end
+
+        def clear
+          @mutex.synchronize do
+            @map = WeakKeyMap.new
+          end
+        end
+      end
+
       include MonitorMixin
       prepend QueryCache::ConnectionPoolConfiguration
       include ConnectionAdapters::AbstractPool
@@ -148,9 +232,9 @@ def initialize(pool_config)
         # then that +thread+ does indeed own that +conn+. However, an absence of such
         # mapping does not mean that the +thread+ doesn't own the said connection. In
         # that case +conn.owner+ attr should be consulted.
-        # Access and modification of <tt>@thread_cached_conns</tt> does not require
+        # Access and modification of <tt>@leases</tt> does not require
         # synchronization.
-        @thread_cached_conns = Concurrent::Map.new(initial_capacity: @size)
+        @leases = LeaseRegistry.new
 
         @connections         = []
         @automatic_reconnect = true
@@ -203,14 +287,18 @@ def internal_metadata # :nodoc:
       #
       # #connection can be called any number of times; the connection is
       # held in a cache keyed by a thread.
-      def connection
-        @thread_cached_conns[ActiveSupport::IsolatedExecutionState.context] ||= checkout
+      def lease_connection
+        lease = connection_lease
+        lease.sticky = true
+        lease.connection ||= checkout
       end
 
+      alias_method :connection, :lease_connection # TODO: deprecate
+
       def pin_connection!(lock_thread) # :nodoc:
         raise "There is already a pinned connection" if @pinned_connection
 
-        @pinned_connection = (@thread_cached_conns[ActiveSupport::IsolatedExecutionState.context] || checkout)
+        @pinned_connection = (connection_lease&.connection || checkout)
         # Any leased connection must be in @connections otherwise
         # some methods like #connected? won't behave correctly
         unless @connections.include?(@pinned_connection)
@@ -252,7 +340,7 @@ def connection_class # :nodoc:
       # #connection or #with_connection methods. Connections obtained through
       # #checkout will not be detected by #active_connection?
       def active_connection?
-        @thread_cached_conns[ActiveSupport::IsolatedExecutionState.context]
+        connection_lease.connection
       end
 
       # Signal that the thread is finished with the current connection.
@@ -262,10 +350,12 @@ def active_connection?
       # This method only works for connections that have been obtained through
       # #connection or #with_connection methods, connections obtained through
       # #checkout will not be automatically released.
-      def release_connection(owner_thread = ActiveSupport::IsolatedExecutionState.context)
-        if conn = @thread_cached_conns.delete(owner_thread)
+      def release_connection(existing_lease = nil)
+        if conn = connection_lease.release
           checkin conn
+          return true
         end
+        false
       end
 
       # Yields a connection from the connection pool to the block. If no connection
@@ -278,13 +368,14 @@ 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
-        if conn = @thread_cached_conns[ActiveSupport::IsolatedExecutionState.context]
-          yield conn
+        lease = connection_lease
+        if lease.connection
+          yield lease.connection
         else
           begin
-            yield connection
+            yield lease.connection = checkout
           ensure
-            release_connection
+            release_connection(lease) unless lease.sticky
           end
         end
       end
@@ -326,7 +417,7 @@ def disconnect(raise_on_acquisition_timeout = true)
               conn.disconnect!
             end
             @connections = []
-            @thread_cached_conns.clear
+            @leases.clear
             @available.clear
           end
         end
@@ -353,7 +444,7 @@ def discard! # :nodoc:
           @connections.each do |conn|
             conn.discard!
           end
-          @connections = @available = @thread_cached_conns = nil
+          @connections = @available = @leases = nil
         end
       end
 
@@ -436,7 +527,7 @@ def checkin(conn)
 
         conn.lock.synchronize do
           synchronize do
-            remove_connection_from_thread_cache conn
+            connection_lease.clear(conn)
 
             conn._run_checkin_callbacks do
               conn.expire
@@ -560,6 +651,10 @@ def schedule_query(future_result) # :nodoc:
       end
 
       private
+        def connection_lease
+          @leases[ActiveSupport::IsolatedExecutionState.context]
+        end
+
         def build_async_executor
           case ActiveRecord.async_query_executor
           when :multi_thread_pool
@@ -734,17 +829,10 @@ def acquire_connection(checkout_timeout)
         #--
         # if owner_thread param is omitted, this must be called in synchronize block
         def remove_connection_from_thread_cache(conn, owner_thread = conn.owner)
-          @thread_cached_conns.delete_pair(owner_thread, conn)
+          @leases[owner_thread].clear(conn)
         end
         alias_method :release, :remove_connection_from_thread_cache
 
-        def prune_thread_cache
-          dead_threads = @thread_cached_conns.keys.reject(&:alive?)
-          dead_threads.each do |dead_thread|
-            @thread_cached_conns.delete(dead_thread)
-          end
-        end
-
         def new_connection
           connection = db_config.new_connection
           connection.pool = self
@@ -788,6 +876,12 @@ def try_to_checkout_new_connection
         def adopt_connection(conn)
           conn.pool = self
           @connections << conn
+
+          # We just created the first connection, it's time to load the schema
+          # cache if that wasn't eagerly done before
+          if @schema_cache.nil? && ActiveRecord.lazily_load_schema_cache
+            schema_cache.load!
+          end
         end
 
         def checkout_new_connection

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

@@ -89,7 +89,7 @@ def initialize(...)
             end
         end
 
-        def checkout(...)
+        def lease_connection
           connection = super
           connection.query_cache ||= query_cache
           connection
@@ -141,7 +141,6 @@ def clear_query_cache
 
         private
           def prune_thread_cache
-            super
             dead_threads = @thread_query_caches.keys.reject(&:alive?)
             dead_threads.each do |dead_thread|
               @thread_query_caches.delete(dead_thread)

activerecord/lib/active_record/connection_adapters/abstract_adapter.rb

@@ -49,10 +49,6 @@ def pool=(value)
         return if value.eql?(@pool)
         @schema_cache = nil
         @pool = value
-
-        if @pool && ActiveRecord.lazily_load_schema_cache
-          @pool.schema_reflection.load!(@pool)
-        end
       end
 
       set_callback :checkin, :after, :enable_lazy_transactions!

activerecord/lib/active_record/connection_handling.rb

@@ -250,8 +250,17 @@ def clear_query_caches_for_current_thread
     # Returns the connection currently associated with the class. This can
     # also be used to "borrow" the connection to do database work unrelated
     # to any of the specific Active Records.
-    def connection
-      connection_pool.connection
+    # The connection will remain leased for the entire duration of the request
+    # or job, or until +#release_connection+ is called.
+    def lease_connection
+      connection_pool.lease_connection
+    end
+
+    alias_method :connection, :lease_connection
+
+    # Return the currently leased connection into the pool
+    def release_connection
+      connection.release_connection
     end
 
     # Checkouts a connection from the pool, yield it and then check it back in.

activerecord/test/cases/connection_handling_test.rb

@@ -11,12 +11,28 @@ class ConnectionHandlingTest < ActiveRecord::TestCase
 
         ActiveRecord::Base.with_connection do |connection|
           assert_predicate ActiveRecord::Base.connection_pool, :active_connection?
-          assert_same connection, ActiveRecord::Base.connection
         end
 
         assert_not_predicate ActiveRecord::Base.connection_pool, :active_connection?
       end
 
+      test "#connection makes the lease permanent even inside #with_connection" do
+        ActiveRecord::Base.connection_pool.release_connection
+        assert_not_predicate ActiveRecord::Base.connection_pool, :active_connection?
+
+        conn = nil
+        ActiveRecord::Base.with_connection do |connection|
+          conn = connection
+          assert_predicate ActiveRecord::Base.connection_pool, :active_connection?
+          2.times do
+            assert_same connection, ActiveRecord::Base.connection
+          end
+        end
+
+        assert_predicate ActiveRecord::Base.connection_pool, :active_connection?
+        assert_same conn, ActiveRecord::Base.connection
+      end
+
       test "#with_connection use the already leased connection if available" do
         leased_connection = ActiveRecord::Base.connection
         assert_predicate ActiveRecord::Base.connection_pool, :active_connection?

activerecord/test/cases/connection_pool_test.rb

@@ -46,10 +46,6 @@ def teardown
         ActiveSupport::IsolatedExecutionState.isolation_level = @previous_isolation_level
       end
 
-      def active_connections(pool)
-        pool.connections.find_all(&:in_use?)
-      end
-
       def test_checkout_after_close
         connection = pool.connection
         assert_predicate connection, :in_use?
@@ -90,6 +86,16 @@ def test_with_connection
             assert_equal 2, active_connections(pool).size
           end
           assert_equal 1, active_connections(pool).size
+
+          pool.with_connection do |conn|
+            assert conn
+            assert_equal 2, active_connections(pool).size
+            pool.connection # lease
+          end
+
+          assert_equal 2, active_connections(pool).size
+          pool.release_connection
+          assert_equal 1, active_connections(pool).size
         }.join
 
         main_thread.close
@@ -697,15 +703,24 @@ def test_disconnect_and_clear_reloadable_connections_attempt_to_wait_for_threads
 
       def test_bang_versions_of_disconnect_and_clear_reloadable_connections_if_unable_to_acquire_all_connections_proceed_anyway
         @pool.checkout_timeout = 0.001 # no need to delay test suite by waiting the whole full default timeout
-        [:disconnect!, :clear_reloadable_connections!].each do |group_action_method|
-          @pool.with_connection do |connection|
-            new_thread { @pool.send(group_action_method) }.join
-            # assert connection has been forcefully taken away from us
-            assert_not_predicate @pool, :active_connection?
 
-            # make a new connection for with_connection to clean up
-            @pool.connection
-          end
+        @pool.with_connection do |connection|
+          new_thread { @pool.disconnect! }.join
+          # assert connection has been forcefully taken away from us
+          assert_not_predicate @pool, :active_connection?
+
+          # make a new connection for with_connection to clean up
+          @pool.connection
+        end
+        @pool.release_connection
+
+        @pool.with_connection do |connection|
+          new_thread { @pool.clear_reloadable_connections! }.join
+          # assert connection has been forcefully taken away from us
+          assert_not_predicate @pool, :active_connection?
+
+          # make a new connection for with_connection to clean up
+          @pool.connection
         end
       end
 
@@ -920,6 +935,10 @@ def test_unpin_connection_returns_whether_transaction_has_been_rolledback
       end
 
       private
+        def active_connections(pool)
+          pool.connections.find_all(&:in_use?)
+        end
+
         def with_single_connection_pool
           config = @db_config.configuration_hash.merge(pool: 1)
           db_config = ActiveRecord::DatabaseConfigurations::HashConfig.new("arunit", "primary", config)

activerecord/test/cases/pooled_connections_test.rb

@@ -36,7 +36,7 @@ def test_pooled_connection_remove
       old_connection = ActiveRecord::Base.connection
       extra_connection = ActiveRecord::Base.connection_pool.checkout
       ActiveRecord::Base.connection_pool.remove(extra_connection)
-      assert_equal ActiveRecord::Base.connection, old_connection
+      assert_equal ActiveRecord::Base.connection.object_id, old_connection.object_id
     end
 
     private