Store query caches in IsolatedExecutionState to avoid memory bloat

Fix: rails#52617
Closes: rails#52618

It isn't per say a leak because the connection reaper eventually prune these,
and also it's not expected that new Threads are constantly making entries into that cache.
But it's true that with Fiber it's probably a bit more common, and the default reaper frequency
isn't adapted to clear this fast enough.

So instead of waiting for the reaper to trigger, which may take a long time, we keep the
caches in IsolatedExcutionState so that when the owning thread or fiber is garbage collected
its cache is too.

However since we need to be able to clear all threads caches, we keep a cache version
on the pool, and bump it to invalidate all caches of a pool at once.

Author: byroot

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

@@ -630,8 +630,6 @@ def reap
             remove conn
           end
         end
-
-        prune_thread_cache
       end
 
       # Disconnect all connections that have been idle for at least

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

@@ -35,22 +35,27 @@ class Store # :nodoc:
         alias_method :enabled?, :enabled
         alias_method :dirties?, :dirties
 
-        def initialize(max_size)
+        def initialize(pool, max_size)
+          @pool = pool
+          @version = pool.query_cache_version
           @map = {}
           @max_size = max_size
           @enabled = false
           @dirties = true
         end
 
         def size
+          check_version
           @map.size
         end
 
         def empty?
+          check_version
           @map.empty?
         end
 
         def [](key)
+          check_version
           return unless @enabled
 
           if entry = @map.delete(key)
@@ -59,6 +64,7 @@ def [](key)
         end
 
         def compute_if_absent(key)
+          check_version
           return yield unless @enabled
 
           if entry = @map.delete(key)
@@ -76,12 +82,21 @@ def clear
           @map.clear
           self
         end
+
+        private
+          def check_version
+            version = @pool.query_cache_version
+            if version != @version
+              @map.clear
+              @version = version
+            end
+          end
       end
 
       module ConnectionPoolConfiguration # :nodoc:
         def initialize(...)
           super
-          @thread_query_caches = Concurrent::Map.new(initial_capacity: @size)
+          @query_cache_version = 0
           @query_cache_max_size = \
             case query_cache = db_config&.query_cache
             when 0, false
@@ -93,6 +108,10 @@ def initialize(...)
             end
         end
 
+        def query_cache_version
+          synchronize { @query_cache_version }
+        end
+
         def checkout_and_verify(connection)
           super
           connection.query_cache ||= query_cache
@@ -141,25 +160,17 @@ def clear_query_cache
             # With transactional fixtures, and especially systems test
             # another thread may use the same connection, but with a different
             # query cache. So we must clear them all.
-            @thread_query_caches.each_value(&:clear)
-          else
-            query_cache.clear
+            synchronize do
+              @query_cache_version += 1
+            end
           end
+          query_cache.clear
         end
 
         def query_cache
-          @thread_query_caches.compute_if_absent(ActiveSupport::IsolatedExecutionState.context) do
-            Store.new(@query_cache_max_size)
-          end
+          key = :"active_record_query_cache_#{object_id}"
+          ActiveSupport::IsolatedExecutionState[key] ||= Store.new(self, @query_cache_max_size)
         end
-
-        private
-          def prune_thread_cache
-            dead_threads = @thread_query_caches.keys.reject(&:alive?)
-            dead_threads.each do |dead_thread|
-              @thread_query_caches.delete(dead_thread)
-            end
-          end
       end
 
       attr_accessor :query_cache

activerecord/test/cases/query_cache_test.rb

@@ -1026,7 +1026,7 @@ def test_cache_is_expired_by_habtm_delete
   end
 
   def test_query_cache_lru_eviction
-    store = ActiveRecord::ConnectionAdapters::QueryCache::Store.new(2)
+    store = ActiveRecord::ConnectionAdapters::QueryCache::Store.new(ActiveRecord::Base.connection_pool, 2)
     store.enabled = true
 
     connection = Post.lease_connection