Refactor query cache and connection lease registry for performance

Fix: rails#52617
Followup: rails#52622

Previous fixes solved the memory issues, but our fallback implementation
of WeakKeyMap actually have terrible performance, and I can't find a way
to do it in a performent way.

So instead we replace it by a specialized weak map that only accept
Thread or Fiber as keys, and simple purge dead threads on insertion.

This gives us reasonable performance on Ruby 3.1 and 3.2.

Author: byroot

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

@@ -119,41 +119,25 @@ def dirties_query_cache
     #   are now explicitly documented
     class ConnectionPool
       if ObjectSpace.const_defined?(:WeakKeyMap) # RUBY_VERSION >= 3.3
-        WeakKeyMap = ::ObjectSpace::WeakKeyMap # :nodoc:
+        WeakThreadKeyMap = ::ObjectSpace::WeakKeyMap # :nodoc:
       else
-        class WeakKeyMap # :nodoc:
+        class WeakThreadKeyMap # :nodoc:
           def initialize
-            @map = ObjectSpace::WeakMap.new
-            @values = nil
-            @size = 0
+            @map = {}
           end
 
-          alias_method :clear, :initialize
+          def clear
+            @map.clear
+          end
 
           def [](key)
-            prune if @map.size != @size
             @map[key]
           end
 
           def []=(key, value)
+            @map.select! { |c, _| c.alive? }
             @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
 
@@ -186,7 +170,7 @@ def clear(connection)
       class LeaseRegistry # :nodoc:
         def initialize
           @mutex = Mutex.new
-          @map = WeakKeyMap.new
+          @map = WeakThreadKeyMap.new
         end
 
         def [](context)
@@ -197,7 +181,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