Merge pull request #36618 from engwan/fix-query-cache-with-shared-ar-connection

Fix query cache when using shared connections

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

@@ -373,7 +373,7 @@ def lock_thread=(lock_thread)
       # #connection can be called any number of times; the connection is
       # held in a cache keyed by a thread.
       def connection
-        @thread_cached_conns[connection_cache_key(@lock_thread || Thread.current)] ||= checkout
+        @thread_cached_conns[connection_cache_key(current_thread)] ||= checkout
       end
 
       # Returns true if there is an open connection being used for the current thread.
@@ -382,7 +382,7 @@ def connection
       # #connection or #with_connection methods. Connections obtained through
       # #checkout will not be detected by #active_connection?
       def active_connection?
-        @thread_cached_conns[connection_cache_key(Thread.current)]
+        @thread_cached_conns[connection_cache_key(current_thread)]
       end
 
       # Signal that the thread is finished with the current connection.
@@ -618,6 +618,10 @@ def connection_cache_key(thread)
           thread
         end
 
+        def current_thread
+          @lock_thread || Thread.current
+        end
+
         # Take control of all existing connections so a "group" action such as
         # reload/disconnect can be performed safely. It is no longer enough to
         # wrap it in +synchronize+ because some pool's actions are allowed

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

@@ -28,17 +28,17 @@ def initialize(*)
         end
 
         def enable_query_cache!
-          @query_cache_enabled[connection_cache_key(Thread.current)] = true
+          @query_cache_enabled[connection_cache_key(current_thread)] = true
           connection.enable_query_cache! if active_connection?
         end
 
         def disable_query_cache!
-          @query_cache_enabled.delete connection_cache_key(Thread.current)
+          @query_cache_enabled.delete connection_cache_key(current_thread)
           connection.disable_query_cache! if active_connection?
         end
 
         def query_cache_enabled
-          @query_cache_enabled[connection_cache_key(Thread.current)]
+          @query_cache_enabled[connection_cache_key(current_thread)]
         end
       end
 

activerecord/test/cases/query_cache_test.rb

@@ -465,6 +465,32 @@ def test_query_caching_is_local_to_the_current_thread
     end
   end
 
+  def test_query_cache_is_enabled_on_all_connection_pools
+    middleware {
+      ActiveRecord::Base.connection_handler.connection_pool_list.each do |pool|
+        assert pool.query_cache_enabled
+        assert pool.connection.query_cache_enabled
+      end
+    }.call({})
+  end
+
+  def test_query_cache_is_enabled_in_threads_with_shared_connection
+    ActiveRecord::Base.connection_pool.lock_thread = true
+
+    assert_cache :off
+
+    thread_a = Thread.new do
+      middleware { |env|
+        assert_cache :clean
+        [200, {}, nil]
+      }.call({})
+    end
+
+    thread_a.join
+
+    ActiveRecord::Base.connection_pool.lock_thread = false
+  end
+
   private
     def middleware(&app)
       executor = Class.new(ActiveSupport::Executor)