Merge pull request rails#52632 from byroot/3.2-connection-lease-registry

Refactor query cache and connection lease registry for performance

Author: byroot

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

@@ -118,42 +118,24 @@ def dirties_query_cache
     # * private methods that require being called in a +synchronize+ blocks
     #   are now explicitly documented
     class ConnectionPool
-      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
+      class WeakThreadKeyMap # :nodoc:
+        # FIXME: On 3.3 we could use ObjectSpace::WeakKeyMap
+        # but it currently cause GC crashes: https://github.com/byroot/rails/pull/3
+        def initialize
+          @map = {}
+        end
 
-          def []=(key, value)
-            @map[key] = value
-            prune if @map.size != @size
-            value
-          end
+        def clear
+          @map.clear
+        end
 
-          def delete(key)
-            if value = self[key]
-              self[key] = nil
-              prune
-            end
-            value
-          end
+        def [](key)
+          @map[key]
+        end
 
-          private
-            def prune(force = false)
-              @values = @map.values
-              @size = @map.size
-            end
+        def []=(key, value)
+          @map.select! { |c, _| c.alive? }
+          @map[key] = value
         end
       end
 
@@ -186,7 +168,7 @@ def clear(connection)
       class LeaseRegistry # :nodoc:
         def initialize
           @mutex = Mutex.new
-          @map = WeakKeyMap.new
+          @map = WeakThreadKeyMap.new
         end
 
         def [](context)
@@ -197,7 +179,7 @@ def [](context)
 
         def clear
           @mutex.synchronize do
-            @map = WeakKeyMap.new
+            @map.clear
           end
         end
       end

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

@@ -66,6 +66,7 @@ def [](key)
 
         def compute_if_absent(key)
           check_version
+
           return yield unless @enabled
 
           if entry = @map.delete(key)
@@ -93,10 +94,30 @@ def check_version
           end
       end
 
+      class QueryCacheRegistry # :nodoc:
+        def initialize
+          @mutex = Mutex.new
+          @map = ConnectionPool::WeakThreadKeyMap.new
+        end
+
+        def compute_if_absent(context)
+          @map[context] || @mutex.synchronize do
+            @map[context] ||= yield
+          end
+        end
+
+        def clear
+          @map.synchronize do
+            @map.clear
+          end
+        end
+      end
+
       module ConnectionPoolConfiguration # :nodoc:
         def initialize(...)
           super
           @query_cache_version = Concurrent::AtomicFixnum.new
+          @thread_query_caches = QueryCacheRegistry.new
           @query_cache_max_size = \
             case query_cache = db_config&.query_cache
             when 0, false
@@ -108,10 +129,6 @@ def initialize(...)
             end
         end
 
-        def query_cache_version
-          synchronize { @query_cache_version }
-        end
-
         def checkout_and_verify(connection)
           super
           connection.query_cache ||= query_cache
@@ -166,8 +183,9 @@ def clear_query_cache
         end
 
         def query_cache
-          caches = ActiveSupport::IsolatedExecutionState[:active_record_query_caches] ||= ConnectionPool::WeakKeyMap.new
-          caches[self] ||= Store.new(@query_cache_version, @query_cache_max_size)
+          @thread_query_caches.compute_if_absent(ActiveSupport::IsolatedExecutionState.context) do
+            Store.new(@query_cache_version, @query_cache_max_size)
+          end
         end
       end