Hold query caches in a WeakKeyMap

Followup: rails#52622

Avoid leaking query caches when a ConnectionPool is gargabe collected.

Author: byroot

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

@@ -118,6 +118,45 @@ 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
+
+          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 Lease # :nodoc:
         attr_accessor :connection, :sticky
 
@@ -145,45 +184,6 @@ def clear(connection)
       end
 
       class LeaseRegistry # :nodoc:
-        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
-
         def initialize
           @mutex = Mutex.new
           @map = WeakKeyMap.new

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

@@ -166,8 +166,8 @@ def clear_query_cache
         end
 
         def query_cache
-          key = :"active_record_query_cache_#{object_id}"
-          ActiveSupport::IsolatedExecutionState[key] ||= Store.new(@query_cache_version, @query_cache_max_size)
+          caches = ActiveSupport::IsolatedExecutionState[:active_record_query_caches] ||= ConnectionPool::WeakKeyMap.new
+          caches[self] ||= Store.new(@query_cache_version, @query_cache_max_size)
         end
       end