From bbf34f3b0a9313d26dd666a57977d873ed2f381b Mon Sep 17 00:00:00 2001 From: Jean Boussier Date: Thu, 25 Jan 2024 13:49:48 +0100 Subject: [PATCH] Refactor QueryCache to be owned by the pool Ref: https://github.com/rails/rails/pull/50793 If we want to stop caching the checked out connections, then we must persist the cache in the pool, and assign it to the connection when it's checked out. The pool become responsible for managing the cache lifecycle. This also open the door to sharing the cache between multiple connections, which is valuable for read replicas, etc. This change only really make sense if we go through with no longer caching checked out connections. Otherwise it's just extra complexity. --- .../abstract/connection_pool.rb | 2 +- .../abstract/query_cache.rb | 197 ++++++++++++------ activerecord/lib/active_record/query_cache.rb | 8 +- activerecord/test/cases/query_cache_test.rb | 13 +- 4 files changed, 143 insertions(+), 77 deletions(-) diff --git a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb index ac4a72b5ab018..4e82aafc6836d 100644 --- a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb +++ b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb @@ -113,7 +113,7 @@ def db_config # are now explicitly documented class ConnectionPool include MonitorMixin - include QueryCache::ConnectionPoolConfiguration + prepend QueryCache::ConnectionPoolConfiguration include ConnectionAdapters::AbstractPool attr_accessor :automatic_reconnect, :checkout_timeout diff --git a/activerecord/lib/active_record/connection_adapters/abstract/query_cache.rb b/activerecord/lib/active_record/connection_adapters/abstract/query_cache.rb index 6871300b8c083..dccc0c4ba3c0d 100644 --- a/activerecord/lib/active_record/connection_adapters/abstract/query_cache.rb +++ b/activerecord/lib/active_record/connection_adapters/abstract/query_cache.rb @@ -13,8 +13,7 @@ def included(base) # :nodoc: :truncate_tables, :rollback_to_savepoint, :rollback_db_transaction, :restart_db_transaction, :exec_insert_all - base.set_callback :checkout, :after, :configure_query_cache! - base.set_callback :checkin, :after, :disable_query_cache! + base.set_callback :checkin, :after, :unset_query_cache! end def dirties_query_cache(base, *method_names) @@ -29,60 +28,146 @@ def #{method_name}(...) end end - module ConnectionPoolConfiguration - def initialize(*) + class Store # :nodoc: + attr_reader :enabled + alias_method :enabled?, :enabled + + def initialize(max_size) + @map = {} + @max_size = max_size + @enabled = false + end + + def enabled=(enabled) + clear if @enabled && !enabled + @enabled = enabled + end + + def size + @map.size + end + + def empty? + @map.empty? + end + + def [](key) + return unless @enabled + + if entry = @map.delete(key) + @map[key] = entry + end + end + + def compute_if_absent(key) + return yield unless @enabled + + if entry = @map.delete(key) + return @map[key] = entry + end + + if @max_size && @map.size >= @max_size + @map.shift # evict the oldest entry + end + + @map[key] ||= yield + end + + def clear + @map.clear + self + end + end + + module ConnectionPoolConfiguration # :nodoc: + def initialize(...) super - @query_cache_enabled = Concurrent::Map.new { false } + @thread_query_caches = Concurrent::Map.new(initial_capacity: @size) + @query_cache_max_size = \ + case query_cache = db_config&.query_cache + when 0, false + nil + when Integer + query_cache + when nil + DEFAULT_SIZE + end + end + + def checkout(...) + connection = super + connection.query_cache ||= query_cache + connection + end + + # Disable the query cache within the block. + def disable_query_cache + cache = query_cache + old, cache.enabled = cache.enabled, false + begin + yield + ensure + cache.enabled = old + end + end + + def enable_query_cache + cache = query_cache + old, cache.enabled = cache.enabled, true + begin + yield + ensure + cache.enabled = old + end end def enable_query_cache! - @query_cache_enabled[connection_cache_key(current_thread)] = true - connection.enable_query_cache! if active_connection? + query_cache.enabled = true end def disable_query_cache! - @query_cache_enabled.delete connection_cache_key(current_thread) - connection.disable_query_cache! if active_connection? + query_cache.enabled = false end def query_cache_enabled - @query_cache_enabled[connection_cache_key(current_thread)] + query_cache.enabled end + + private + def query_cache + @thread_query_caches.compute_if_absent(connection_cache_key(current_thread)) do + Store.new(@query_cache_max_size) + end + end end - attr_reader :query_cache, :query_cache_enabled + attr_accessor :query_cache def initialize(*) super - @query_cache = {} - @query_cache_enabled = false - @query_cache_max_size = nil + @query_cache = nil + end + + def query_cache_enabled + @query_cache&.enabled? end # Enable the query cache within the block. - def cache - old, @query_cache_enabled = @query_cache_enabled, true - yield - ensure - @query_cache_enabled = old - clear_query_cache unless @query_cache_enabled + def cache(&) + pool.enable_query_cache(&) end def enable_query_cache! - @query_cache_enabled = true + pool.enable_query_cache! end - def disable_query_cache! - @query_cache_enabled = false - clear_query_cache + # Disable the query cache within the block. + def uncached(&) + pool.disable_query_cache(&) end - # Disable the query cache within the block. - def uncached - old, @query_cache_enabled = @query_cache_enabled, false - yield - ensure - @query_cache_enabled = old + def disable_query_cache! + pool.disable_query_cache! end # Clears the query cache. @@ -93,7 +178,7 @@ def uncached # undermining the randomness you were expecting. def clear_query_cache @lock.synchronize do - @query_cache.clear + @query_cache&.clear end end @@ -102,7 +187,7 @@ def select_all(arel, name = nil, binds = [], preparable: nil, async: false) # :n # If arel is locked this is a SELECT ... FOR UPDATE or somesuch. # Such queries should not be cached. - if @query_cache_enabled && !(arel.respond_to?(:locked) && arel.locked) + if @query_cache&.enabled? && !(arel.respond_to?(:locked) && arel.locked) sql, binds, preparable = to_sql_and_binds(arel, binds, preparable) if async @@ -117,42 +202,37 @@ def select_all(arel, name = nil, binds = [], preparable: nil, async: false) # :n end private + def unset_query_cache! + @query_cache = nil + end + def lookup_sql_cache(sql, name, binds) key = binds.empty? ? sql : [sql, binds] - hit = false - result = nil + result = nil @lock.synchronize do - if (result = @query_cache.delete(key)) - hit = true - @query_cache[key] = result - end + result = @query_cache[key] end - if hit + if result ActiveSupport::Notifications.instrument( "sql.active_record", cache_notification_info(sql, name, binds) ) - - result end + + result end def cache_sql(sql, name, binds) key = binds.empty? ? sql : [sql, binds] result = nil - hit = false + hit = true @lock.synchronize do - if (result = @query_cache.delete(key)) - hit = true - @query_cache[key] = result - else - result = @query_cache[key] = yield - if @query_cache_max_size && @query_cache.size > @query_cache_max_size - @query_cache.shift - end + result = @query_cache.compute_if_absent(key) do + hit = false + yield end end @@ -178,23 +258,6 @@ def cache_notification_info(sql, name, binds) cached: true } end - - def configure_query_cache! - case query_cache = pool.db_config.query_cache - when 0, false - return - when Integer - @query_cache_max_size = query_cache - when nil - @query_cache_max_size = DEFAULT_SIZE - else - @query_cache_max_size = nil # no limit - end - - if pool.query_cache_enabled - enable_query_cache! - end - end end end end diff --git a/activerecord/lib/active_record/query_cache.rb b/activerecord/lib/active_record/query_cache.rb index 08db278880fe1..18b0a7c901f53 100644 --- a/activerecord/lib/active_record/query_cache.rb +++ b/activerecord/lib/active_record/query_cache.rb @@ -8,7 +8,7 @@ module ClassMethods # If it's not, it will execute the given block. def cache(&block) if connected? || !configurations.empty? - connection.cache(&block) + connection_pool.enable_query_cache(&block) else yield end @@ -18,7 +18,7 @@ def cache(&block) # If it's not, it will execute the given block. def uncached(&block) if connected? || !configurations.empty? - connection.uncached(&block) + connection_pool.disable_query_cache(&block) else yield end @@ -26,11 +26,11 @@ def uncached(&block) end def self.run - ActiveRecord::Base.connection_handler.each_connection_pool.reject { |p| p.query_cache_enabled }.each { |p| p.enable_query_cache! } + ActiveRecord::Base.connection_handler.each_connection_pool.reject(&:query_cache_enabled).each(&:enable_query_cache!) end def self.complete(pools) - pools.each { |pool| pool.disable_query_cache! } + pools.each(&:disable_query_cache!) ActiveRecord::Base.connection_handler.each_connection_pool do |pool| pool.release_connection if pool.active_connection? && !pool.connection.transaction_open? diff --git a/activerecord/test/cases/query_cache_test.rb b/activerecord/test/cases/query_cache_test.rb index 89473b73b2732..39a768e7910d8 100644 --- a/activerecord/test/cases/query_cache_test.rb +++ b/activerecord/test/cases/query_cache_test.rb @@ -806,7 +806,7 @@ def test_cache_gets_cleared_after_migration end def test_find - assert_called(Task.connection.query_cache, :clear) do + assert_called(Task.connection.query_cache, :clear, times: 2) do assert_not Task.connection.query_cache_enabled Task.cache do assert Task.connection.query_cache_enabled @@ -920,9 +920,12 @@ def test_cache_is_expired_by_habtm_delete end def test_query_cache_lru_eviction + store = ActiveRecord::ConnectionAdapters::QueryCache::Store.new(2) + store.enabled = true + connection = Post.connection - connection.pool.db_config.stub(:query_cache, 2) do - connection.send(:configure_query_cache!) + old_store, connection.query_cache = connection.query_cache, store + begin Post.cache do assert_queries_count(2) do connection.select_all("SELECT 1") @@ -943,9 +946,9 @@ def test_query_cache_lru_eviction connection.select_all("SELECT 2") end end + ensure + connection.query_cache = old_store end - ensure - connection.send(:configure_query_cache!) end test "threads use the same connection" do