From 3751e2774d63ebb7ccb0b67fdf92009195c200e3 Mon Sep 17 00:00:00 2001 From: Igor Ohrimenko Date: Sun, 26 Jul 2026 23:51:33 +0300 Subject: [PATCH 1/3] feat(config): add query_cache_memory_limit and query_cache_ttl The AST query cache was bounded only by entry count (query_cache_limit), not by memory, and had no idle expiry. Heavy/complex queries produce large parse trees, so a full cache can hold GBs of live RSS that only drop once 1000 newer distinct queries evict them (or on RESET/restart). Adds two [general] options (also PGDOG_QUERY_CACHE_MEMORY_LIMIT / PGDOG_QUERY_CACHE_TTL env), both default 0 (off): - query_cache_memory_limit: evict LRU until the summed entry size is under the byte budget. Entry size is measured via jemalloc per-thread allocation counters (clamped to 0, since cross-thread frees under a work-stealing runtime can make the delta negative); falls back to query length on non-jemalloc builds. - query_cache_ttl: drop entries not accessed within the window, via the existing 1s maintenance sweep. Refs #1261. --- .schema/pgdog.schema.json | 16 + Cargo.lock | 18 + pgdog-config/src/general.rs | 36 ++ pgdog/Cargo.toml | 1 + pgdog/src/backend/databases.rs | 16 +- pgdog/src/frontend/prepared_statements/mod.rs | 3 + pgdog/src/frontend/router/parser/cache/ast.rs | 24 +- .../router/parser/cache/cache_impl.rs | 349 ++++++++++++++++-- 8 files changed, 415 insertions(+), 48 deletions(-) diff --git a/.schema/pgdog.schema.json b/.schema/pgdog.schema.json index 3900e97cd..f632cf11d 100644 --- a/.schema/pgdog.schema.json +++ b/.schema/pgdog.schema.json @@ -84,6 +84,8 @@ "prepared_statements_limit": 9223372036854775807, "pub_sub_channel_size": 0, "query_cache_limit": 1000, + "query_cache_memory_limit": 0, + "query_cache_ttl": 0, "query_log": null, "query_log_stdout": false, "query_parser": "auto", @@ -978,6 +980,20 @@ "default": 1000, "minimum": 0 }, + "query_cache_memory_limit": { + "description": "Approximate memory budget (bytes) for the query (AST) cache; entries are evicted (LRU) once the sum of their sizes exceeds it. Bounds RAM regardless of query complexity, unlike the count-based `query_cache_limit`. `0` disables the memory cap.\n\n**Note:** Entry sizes are measured with jemalloc; on builds without it the budget is approximated from query text length.\n\n_Default:_ `0`\n\n", + "type": "integer", + "format": "uint", + "default": 0, + "minimum": 0 + }, + "query_cache_ttl": { + "description": "Time-to-idle (seconds) for query (AST) cache entries; an entry not accessed within this window is dropped by the maintenance sweep. `0` disables idle expiry.\n\n_Default:_ `0`\n\n", + "type": "integer", + "format": "uint", + "default": 0, + "minimum": 0 + }, "query_log": { "description": "Path to a file where all queries are logged. Logging every query is slow; do not use in production.", "type": [ diff --git a/Cargo.lock b/Cargo.lock index 662b1f8ee..b48888766 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3225,6 +3225,12 @@ dependencies = [ "windows-link", ] +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + [[package]] name = "pem-rfc7468" version = "0.7.0" @@ -3396,6 +3402,7 @@ dependencies = [ "stats_alloc", "tempfile", "thiserror 2.0.18", + "tikv-jemalloc-ctl", "tikv-jemallocator", "tokio", "tokio-rustls", @@ -5284,6 +5291,17 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "tikv-jemalloc-ctl" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "661f1f6a57b3a36dc9174a2c10f19513b4866816e13425d3e418b11cc37bc24c" +dependencies = [ + "libc", + "paste", + "tikv-jemalloc-sys", +] + [[package]] name = "tikv-jemalloc-sys" version = "0.6.1+5.3.0-1-ge13ca993e8ccb9ba9847cc330696e02839f328f7" diff --git a/pgdog-config/src/general.rs b/pgdog-config/src/general.rs index 5ad4cc018..76e669d09 100644 --- a/pgdog-config/src/general.rs +++ b/pgdog-config/src/general.rs @@ -378,6 +378,24 @@ pub struct General { #[serde(default = "General::query_cache_limit")] pub query_cache_limit: usize, + /// Approximate memory budget (bytes) for the query (AST) cache; entries are evicted (LRU) once the sum of their sizes exceeds it. Bounds RAM regardless of query complexity, unlike the count-based `query_cache_limit`. `0` disables the memory cap. + /// + /// **Note:** Entry sizes are measured with jemalloc; on builds without it the budget is approximated from query text length. + /// + /// _Default:_ `0` + /// + /// + #[serde(default = "General::query_cache_memory_limit")] + pub query_cache_memory_limit: usize, + + /// Time-to-idle (seconds) for query (AST) cache entries; an entry not accessed within this window is dropped by the maintenance sweep. `0` disables idle expiry. + /// + /// _Default:_ `0` + /// + /// + #[serde(default = "General::query_cache_ttl")] + pub query_cache_ttl: usize, + /// Toggle automatic creation of connection pools given the user name, database and password. /// /// _Default:_ `disabled` @@ -854,6 +872,8 @@ impl Default for General { query_parser_engine: QueryParserEngine::default(), prepared_statements_limit: Self::prepared_statements_limit(), query_cache_limit: Self::query_cache_limit(), + query_cache_memory_limit: Self::query_cache_memory_limit(), + query_cache_ttl: Self::query_cache_ttl(), passthrough_auth: Self::default_passthrough_auth(), connect_timeout: Self::default_connect_timeout(), connect_attempt_delay: Self::default_connect_attempt_delay(), @@ -1325,6 +1345,14 @@ impl General { Self::env_or_default("PGDOG_QUERY_CACHE_LIMIT", 1_000) } + pub fn query_cache_memory_limit() -> usize { + Self::env_or_default("PGDOG_QUERY_CACHE_MEMORY_LIMIT", 0) + } + + pub fn query_cache_ttl() -> usize { + Self::env_or_default("PGDOG_QUERY_CACHE_TTL", 0) + } + pub fn log_format() -> LogFormat { Self::env_enum_or_default("PGDOG_LOG_FORMAT") } @@ -1743,6 +1771,8 @@ mod tests { let _guard = set_env_var("PGDOG_OPENMETRICS_PORT", "9090"); let _guard = set_env_var("PGDOG_PREPARED_STATEMENTS_LIMIT", "1000"); let _guard = set_env_var("PGDOG_QUERY_CACHE_LIMIT", "500"); + let _guard = set_env_var("PGDOG_QUERY_CACHE_MEMORY_LIMIT", "1048576"); + let _guard = set_env_var("PGDOG_QUERY_CACHE_TTL", "600"); let _guard = set_env_var("PGDOG_CONNECT_ATTEMPTS", "3"); let _guard = set_env_var("PGDOG_MIRROR_QUEUE", "256"); let _guard = set_env_var("PGDOG_MIRROR_EXPOSURE", "0.5"); @@ -1755,6 +1785,8 @@ mod tests { assert_eq!(General::openmetrics_port(), Some(9090)); assert_eq!(General::prepared_statements_limit(), 1000); assert_eq!(General::query_cache_limit(), 500); + assert_eq!(General::query_cache_memory_limit(), 1048576); + assert_eq!(General::query_cache_ttl(), 600); assert_eq!(General::connect_attempts(), 3); assert_eq!(General::mirror_queue(), 256); assert_eq!(General::mirror_exposure(), 0.5); @@ -1767,6 +1799,8 @@ mod tests { let _guard = remove_env_var("PGDOG_OPENMETRICS_PORT"); let _guard = remove_env_var("PGDOG_PREPARED_STATEMENTS_LIMIT"); let _guard = remove_env_var("PGDOG_QUERY_CACHE_LIMIT"); + let _guard = remove_env_var("PGDOG_QUERY_CACHE_MEMORY_LIMIT"); + let _guard = remove_env_var("PGDOG_QUERY_CACHE_TTL"); let _guard = remove_env_var("PGDOG_CONNECT_ATTEMPTS"); let _guard = remove_env_var("PGDOG_MIRROR_QUEUE"); let _guard = remove_env_var("PGDOG_MIRROR_EXPOSURE"); @@ -1779,6 +1813,8 @@ mod tests { assert_eq!(General::openmetrics_port(), None); assert_eq!(General::prepared_statements_limit(), i64::MAX as usize); assert_eq!(General::query_cache_limit(), 1_000); + assert_eq!(General::query_cache_memory_limit(), 0); + assert_eq!(General::query_cache_ttl(), 0); assert_eq!(General::connect_attempts(), 1); assert_eq!(General::mirror_queue(), 128); assert_eq!(General::mirror_exposure(), 1.0); diff --git a/pgdog/Cargo.toml b/pgdog/Cargo.toml index 725aa94f3..20bc624e2 100644 --- a/pgdog/Cargo.toml +++ b/pgdog/Cargo.toml @@ -92,6 +92,7 @@ libc = "0.2" [target.'cfg(not(target_env = "msvc"))'.dependencies] tikv-jemallocator = "0.6" +tikv-jemalloc-ctl = "0.6" [build-dependencies] diff --git a/pgdog/src/backend/databases.rs b/pgdog/src/backend/databases.rs index 11d553d94..cc7da3800 100644 --- a/pgdog/src/backend/databases.rs +++ b/pgdog/src/backend/databases.rs @@ -103,8 +103,12 @@ pub fn init() -> Result<(), Error> { let config = config(); replace_databases(from_config(&config), false)?; - // Resize query cache - Cache::resize(config.config.general.query_cache_limit); + // Configure query cache limits + Cache::configure( + config.config.general.query_cache_limit, + config.config.general.query_cache_memory_limit, + config.config.general.query_cache_ttl, + ); // Start two-pc manager. let _monitor = Manager::get(); @@ -151,8 +155,12 @@ pub fn reload() -> Result<(), Error> { .write() .close_unused(new_config.config.general.prepared_statements_limit); - // Resize query cache. - Cache::resize(new_config.config.general.query_cache_limit); + // Configure query cache limits. + Cache::configure( + new_config.config.general.query_cache_limit, + new_config.config.general.query_cache_memory_limit, + new_config.config.general.query_cache_ttl, + ); Ok(()) } diff --git a/pgdog/src/frontend/prepared_statements/mod.rs b/pgdog/src/frontend/prepared_statements/mod.rs index 7ee696be1..d36a0c254 100644 --- a/pgdog/src/frontend/prepared_statements/mod.rs +++ b/pgdog/src/frontend/prepared_statements/mod.rs @@ -9,6 +9,7 @@ use tracing::debug; use crate::{ config::{PreparedStatements as PreparedStatementsLevel, config}, + frontend::router::parser::Cache, net::{Parse, ProtocolMessage}, }; @@ -189,6 +190,8 @@ pub fn start_maintenance() { pub fn run_maintenance() { let capacity = config().config.general.prepared_statements_limit; PreparedStatements::global().write().close_unused(capacity); + // Drop AST cache entries idle beyond their time-to-idle (no-op when disabled). + Cache::sweep(); } #[cfg(test)] diff --git a/pgdog/src/frontend/router/parser/cache/ast.rs b/pgdog/src/frontend/router/parser/cache/ast.rs index 12fc0a86d..b41fee389 100644 --- a/pgdog/src/frontend/router/parser/cache/ast.rs +++ b/pgdog/src/frontend/router/parser/cache/ast.rs @@ -55,22 +55,22 @@ pub struct AstInner { impl AstInner { /// Create new AST record, with no rewrite or comment routing. #[cfg(feature = "new_parser")] - pub(crate) fn new(ast: Owned) -> Self { + pub(crate) fn new(ast: Owned, query_without_comment: Arc) -> Self { Self { ast, stats: Mutex::new(Stats::new()), rewrite_plan: RewritePlan::default(), - query_without_comment: "".into(), + query_without_comment, } } #[cfg(not(feature = "new_parser"))] - pub(crate) fn old(ast: ParseResult) -> Self { + pub(crate) fn old(ast: ParseResult, query_without_comment: Arc) -> Self { Self { ast, stats: Mutex::new(Stats::new()), rewrite_plan: RewritePlan::default(), - query_without_comment: "".into(), + query_without_comment, } } } @@ -84,6 +84,14 @@ impl Deref for Ast { } impl Ast { + /// Rough byte footprint of this cache entry, used as a fallback when the + /// jemalloc allocation measurement is unavailable. The query text length + /// tracks entry weight well enough for the memory budget: a wide report + /// query has a long body, `SELECT 1` a short one. + pub fn approx_size(&self) -> usize { + self.query_without_comment.len() + } + /// Parse statement and run the rewrite engine, if necessary. pub(super) fn new( query: &AstQuery, @@ -195,7 +203,7 @@ impl Ast { comment_role: None, comment_shard: None, query_parser_engine, - inner: Arc::new(AstInner::new(ast.into_inner())), + inner: Arc::new(AstInner::new(ast.into_inner(), query.into())), }) } @@ -213,7 +221,7 @@ impl Ast { comment_role: None, comment_shard: None, query_parser_engine, - inner: Arc::new(AstInner::old(ast)), + inner: Arc::new(AstInner::old(ast, query.into())), }) } } @@ -228,7 +236,7 @@ impl Ast { comment_role: None, comment_shard: None, query_parser_engine: QueryParserEngine::default(), - inner: Arc::new(AstInner::new(stmts)), + inner: Arc::new(AstInner::new(stmts, "".into())), } } @@ -240,7 +248,7 @@ impl Ast { comment_role: None, comment_shard: None, query_parser_engine: QueryParserEngine::default(), - inner: Arc::new(AstInner::old(parse_result)), + inner: Arc::new(AstInner::old(parse_result, "".into())), } } diff --git a/pgdog/src/frontend/router/parser/cache/cache_impl.rs b/pgdog/src/frontend/router/parser/cache/cache_impl.rs index 107a395e0..d55e5b470 100644 --- a/pgdog/src/frontend/router/parser/cache/cache_impl.rs +++ b/pgdog/src/frontend/router/parser/cache/cache_impl.rs @@ -6,7 +6,7 @@ use pg_query::normalize; use pg_raw_parse::normalize::normalize; use pgdog_config::QueryParserEngine; use std::collections::HashMap; -use std::time::Duration; +use std::time::{Duration, Instant}; use parking_lot::Mutex; use std::sync::Arc; @@ -45,15 +45,121 @@ impl Stats { } } +/// Query cache entry: the AST plus bookkeeping for the memory budget +/// (`size`) and the idle-expiry sweep (`accessed`). +#[derive(Debug, Clone)] +struct Entry { + ast: Ast, + accessed: Instant, + size: usize, +} + /// Mutex-protected query cache. #[derive(Debug)] pub(super) struct Inner { - /// Least-recently-used cache. - queries: LruCache, Ast>, + /// Least-recently-used cache. Kept unbounded; the count and memory limits + /// are enforced together in `enforce`. + queries: LruCache, Entry>, + /// Maximum number of cached entries (0 = unlimited). + count_limit: usize, + /// Approximate memory budget in bytes (0 = unlimited). + byte_limit: usize, + /// Sum of `Entry::size` across all cached entries. + bytes: usize, + /// Idle expiry: entries untouched for longer than this are swept (None = off). + ttl: Option, /// Cache global stats. pub(super) stats: Stats, } +impl Inner { + /// Insert (or replace) an entry and enforce the count and memory limits. + /// `size` is the measured byte footprint (0 falls back to an estimate). + fn insert(&mut self, key: Arc, ast: Ast, size: usize) { + let size = if size > 0 { size } else { ast.approx_size() }; + if let Some(old) = self.queries.put( + key, + Entry { + ast, + accessed: Instant::now(), + size, + }, + ) { + self.bytes = self.bytes.saturating_sub(old.size); + } + self.bytes = self.bytes.saturating_add(size); + self.enforce(); + } + + /// Evict least-recently-used entries until both limits are satisfied. + fn enforce(&mut self) { + while (self.count_limit > 0 && self.queries.len() > self.count_limit) + || (self.byte_limit > 0 && self.bytes > self.byte_limit) + { + match self.queries.pop_lru() { + Some((_, evicted)) => self.bytes = self.bytes.saturating_sub(evicted.size), + None => break, + } + } + } + + /// Drop entries not accessed within the time-to-idle window. + fn sweep(&mut self) { + let Some(ttl) = self.ttl else { + return; + }; + let now = Instant::now(); + // Two-pass: the LRU cache can't be mutated while iterating. `collect` + // into an empty Vec does not allocate, so a sweep with nothing idle + // (the common case) is allocation-free. + let stale: Vec> = self + .queries + .iter() + .filter(|(_, e)| now.duration_since(e.accessed) >= ttl) + .map(|(k, _)| k.clone()) + .collect(); + for key in stale { + if let Some(e) = self.queries.pop(&key) { + self.bytes = self.bytes.saturating_sub(e.size); + } + } + } +} + +/// Measure the net bytes a build closure keeps allocated, using jemalloc's +/// per-thread allocation counters. The parse is synchronous (no `.await` +/// between the reads), so the delta is attributable to the entry it builds. +/// Returns 0 when the counters are unavailable, so callers fall back to the +/// `approx_size` estimate. +#[cfg(all(not(test), not(target_env = "msvc")))] +fn measure_build(f: impl FnOnce() -> T) -> (T, usize) { + use tikv_jemalloc_ctl::thread::{allocatedp, allocatedp_mib, deallocatedp, deallocatedp_mib}; + // Cache the MIBs once so each measurement skips the name-to-MIB lookup. + static ALLOCATED: Lazy> = Lazy::new(|| allocatedp::mib().ok()); + static DEALLOCATED: Lazy> = Lazy::new(|| deallocatedp::mib().ok()); + let net = || -> Option { + let a = ALLOCATED.as_ref()?.read().ok()?.get(); + let d = DEALLOCATED.as_ref()?.read().ok()?.get(); + Some(a.wrapping_sub(d)) + }; + match net() { + Some(before) => { + let r = f(); + let after = net().unwrap_or(before); + // `deallocated` counts frees of memory allocated on other threads, + // so the delta can be negative under a work-stealing runtime. + let delta = after.wrapping_sub(before) as i64; + (r, delta.max(0) as usize) + } + None => (f(), 0), + } +} + +#[cfg(any(test, target_env = "msvc"))] +fn measure_build(f: impl FnOnce() -> T) -> (T, usize) { + (f(), 0) +} + /// AST cache. #[derive(Clone, Debug)] pub struct Cache { @@ -66,26 +172,46 @@ impl Cache { Self { inner: Arc::new(Mutex::new(Inner { queries: LruCache::unbounded(), + count_limit: 0, + byte_limit: 0, + bytes: 0, + ttl: None, stats: Stats::default(), })), } } - /// Resize cache to capacity, evicting any statements exceeding the capacity. - /// - /// Minimum capacity is 1. - pub fn resize(capacity: usize) { - let capacity = if capacity == 0 { 1 } else { capacity }; - - CACHE - .inner - .lock() - .queries - .resize(capacity.try_into().unwrap()); + /// Apply cache limits from configuration, evicting anything over the new + /// caps. A `count`, `bytes`, or `ttl_seconds` of 0 disables that limit. + pub fn configure(count: usize, bytes: usize, ttl_seconds: usize) { + let mut guard = CACHE.inner.lock(); + guard.count_limit = count; + guard.byte_limit = bytes; + guard.ttl = if ttl_seconds == 0 { + None + } else { + Some(Duration::from_secs(ttl_seconds as u64)) + }; + guard.enforce(); + debug!( + "ast cache limits: count={} bytes={} ttl={}s", + count, bytes, ttl_seconds + ); + } + /// Resize cache to a count capacity, keeping the memory and TTL limits. + pub fn resize(capacity: usize) { + let mut guard = CACHE.inner.lock(); + guard.count_limit = capacity.max(1); + guard.enforce(); debug!("ast cache size set to {}", capacity); } + /// Run the idle-expiry sweep. Called periodically by maintenance. + pub fn sweep() { + CACHE.inner.lock().sweep(); + } + /// Handle parsing a query. pub fn query( &self, @@ -116,9 +242,11 @@ impl Cache { { let mut guard = self.inner.lock(); + let now = Instant::now(); let ast = guard.queries.get_mut(query_and_comment.query).map(|entry| { - entry.stats.lock().hits += 1; // No contention on this. - entry.clone() + entry.accessed = now; + entry.ast.stats.lock().hits += 1; // No contention on this. + entry.ast.clone() }); if let Some(mut ast) = ast { guard.stats.hits += 1; @@ -129,15 +257,18 @@ impl Cache { } } - // Parse query without holding lock. - let mut entry = Ast::with_context( - &AstQuery { - original_query: query, - query_without_comment: query_and_comment.query, - }, - ctx, - prepared_statements, - )?; + // Parse query without holding lock, measuring the entry's footprint. + let (built, size) = measure_build(|| { + Ast::with_context( + &AstQuery { + original_query: query, + query_without_comment: query_and_comment.query, + }, + ctx, + prepared_statements, + ) + }); + let mut entry = built?; entry.comment_role = query_and_comment.role; entry.comment_shard = query_and_comment.shard.clone(); let parse_time = entry.stats.lock().parse_time; @@ -150,9 +281,7 @@ impl Cache { // (direct-shard) variant. let cacheable = entry.comment_shard.is_none() || entry.rewrite_plan.is_empty(); if cacheable { - guard - .queries - .put(entry.query_without_comment.clone(), entry.clone()); + guard.insert(entry.query_without_comment.clone(), entry.clone(), size); } guard.stats.misses += 1; guard.stats.parse_time += parse_time; @@ -205,18 +334,21 @@ impl Cache { { let mut guard = self.inner.lock(); - if let Some(entry) = guard.queries.get(normalized.as_str()) { - entry.update_stats(route); + let now = Instant::now(); + if let Some(entry) = guard.queries.get_mut(normalized.as_str()) { + entry.accessed = now; + entry.ast.update_stats(route); guard.stats.hits += 1; return Ok(()); } } - let entry = Ast::new_record(&normalized, query_parser_engine)?; + let (built, size) = measure_build(|| Ast::new_record(&normalized, query_parser_engine)); + let entry = built?; entry.update_stats(route); let mut guard = self.inner.lock(); - guard.queries.put(normalized.into(), entry); + guard.insert(normalized.into(), entry, size); guard.stats.misses += 1; Ok(()) @@ -237,7 +369,7 @@ impl Cache { guard .queries .iter() - .map(|c| *c.1.stats.lock()) + .map(|c| *c.1.ast.stats.lock()) .collect::>(), guard.stats, ) @@ -256,17 +388,162 @@ impl Cache { .lock() .queries .iter() - .map(|i| (i.0.clone(), i.1.clone())) + .map(|i| (i.0.clone(), i.1.ast.clone())) .collect() } - /// Reset cache, removing all statements - /// and setting stats to 0. + /// Reset cache, removing all statements and setting stats to 0. The + /// configured count/memory/TTL limits are kept. pub fn reset() { let cache = Self::get(); let mut guard = cache.inner.lock(); guard.queries.clear(); + guard.bytes = 0; guard.stats.hits = 0; guard.stats.misses = 0; } } + +#[cfg(test)] +mod tests { + use super::*; + + /// A minimal valid AST entry; its content is irrelevant to the limit + /// logic, which is driven by the explicit `size` passed to `insert`. + fn ast() -> Ast { + Ast::new_record("SELECT 1", QueryParserEngine::PgQueryProtobuf).expect("parse") + } + + fn inner(count_limit: usize, byte_limit: usize, ttl: Option) -> Inner { + Inner { + queries: LruCache::unbounded(), + count_limit, + byte_limit, + bytes: 0, + ttl, + stats: Stats::default(), + } + } + + #[test] + fn count_limit_caps_and_evicts_lru() { + let mut c = inner(3, 0, None); + for i in 0..5 { + c.insert(format!("q{i}").into(), ast(), 10); + } + assert_eq!(c.queries.len(), 3); + assert!(c.queries.peek("q4").is_some(), "newest kept"); + assert!(c.queries.peek("q2").is_some()); + assert!(c.queries.peek("q1").is_none(), "oldest evicted"); + assert_eq!(c.bytes, 30, "byte total tracks surviving entries"); + } + + #[test] + fn count_limit_zero_is_unlimited() { + let mut c = inner(0, 0, None); + for i in 0..1000 { + c.insert(format!("q{i}").into(), ast(), 1); + } + assert_eq!(c.queries.len(), 1000); + assert_eq!(c.bytes, 1000); + } + + #[test] + fn byte_limit_caps_and_evicts_lru() { + let mut c = inner(0, 25, None); + for i in 0..5 { + c.insert(format!("q{i}").into(), ast(), 10); + } + assert!(c.bytes <= 25, "stays within byte budget"); + assert_eq!(c.queries.len(), 2); + assert!(c.queries.peek("q4").is_some(), "newest kept"); + assert!(c.queries.peek("q0").is_none(), "oldest evicted"); + } + + #[test] + fn entry_larger_than_byte_budget_is_not_cached() { + let mut c = inner(0, 25, None); + c.insert("big".into(), ast(), 100); + assert_eq!( + c.queries.len(), + 0, + "an entry over the whole budget is evicted" + ); + assert_eq!(c.bytes, 0); + } + + #[test] + fn replacing_a_key_updates_byte_total() { + let mut c = inner(0, 0, None); + c.insert("q".into(), ast(), 10); + assert_eq!(c.bytes, 10); + c.insert("q".into(), ast(), 30); + assert_eq!(c.queries.len(), 1); + assert_eq!(c.bytes, 30, "old size subtracted, new size added"); + } + + #[test] + fn byte_total_saturates_instead_of_overflowing() { + let mut c = inner(0, 0, None); + c.insert("q1".into(), ast(), usize::MAX); + c.insert("q2".into(), ast(), usize::MAX); + assert_eq!(c.bytes, usize::MAX); + assert_eq!(c.queries.len(), 2); + } + + #[test] + fn zero_size_falls_back_to_query_length() { + let mut c = inner(0, 0, None); + c.insert("q".into(), ast(), 0); + assert_eq!( + c.bytes, + "SELECT 1".len(), + "record entries carry their query text" + ); + } + + #[test] + fn no_limits_keeps_everything() { + let mut c = inner(0, 0, None); + for i in 0..50 { + c.insert(format!("q{i}").into(), ast(), 7); + } + assert_eq!(c.queries.len(), 50); + assert_eq!(c.bytes, 350); + } + + #[test] + fn sweep_is_noop_when_ttl_disabled() { + let mut c = inner(0, 0, None); + for i in 0..3 { + c.insert(format!("q{i}").into(), ast(), 5); + } + c.sweep(); + assert_eq!(c.queries.len(), 3); + assert_eq!(c.bytes, 15); + } + + #[test] + fn sweep_keeps_fresh_entries() { + // TTL far larger than the entries' age: nothing is idle yet. + let mut c = inner(0, 0, Some(Duration::from_secs(3600))); + for i in 0..3 { + c.insert(format!("q{i}").into(), ast(), 5); + } + c.sweep(); + assert_eq!(c.queries.len(), 3); + assert_eq!(c.bytes, 15); + } + + #[test] + fn sweep_drops_idle_entries_and_updates_bytes() { + // TTL of zero: every entry is at or past its idle window. + let mut c = inner(0, 0, Some(Duration::ZERO)); + for i in 0..3 { + c.insert(format!("q{i}").into(), ast(), 5); + } + c.sweep(); + assert_eq!(c.queries.len(), 0); + assert_eq!(c.bytes, 0); + } +} From b78543c9b7a29b9341368dc42a706a3978b2c7b8 Mon Sep 17 00:00:00 2001 From: Igor Ohrimenko Date: Sat, 1 Aug 2026 17:00:58 +0300 Subject: [PATCH 2/3] refactor(config): rename query_cache_ttl to query_cache_idle_timeout, use milliseconds MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The behaviour is time-to-idle (expiry counted from last access), not time-to-live, so the ttl name over-promised. The config already has an idle-timeout vocabulary (idle_timeout, client_idle_timeout) — reuse it, and switch the unit from seconds to milliseconds to match the other *_timeout options. Env variable becomes PGDOG_QUERY_CACHE_IDLE_TIMEOUT. Schema regenerated. --- .schema/pgdog.schema.json | 16 ++++----- pgdog-config/src/general.rs | 22 ++++++------ pgdog/src/backend/databases.rs | 4 +-- pgdog/src/frontend/prepared_statements/mod.rs | 2 +- .../router/parser/cache/cache_impl.rs | 36 +++++++++---------- 5 files changed, 40 insertions(+), 40 deletions(-) diff --git a/.schema/pgdog.schema.json b/.schema/pgdog.schema.json index f632cf11d..08c65b71b 100644 --- a/.schema/pgdog.schema.json +++ b/.schema/pgdog.schema.json @@ -83,9 +83,9 @@ "prepared_statements": "extended", "prepared_statements_limit": 9223372036854775807, "pub_sub_channel_size": 0, + "query_cache_idle_timeout": 0, "query_cache_limit": 1000, "query_cache_memory_limit": 0, - "query_cache_ttl": 0, "query_log": null, "query_log_stdout": false, "query_parser": "auto", @@ -973,6 +973,13 @@ "default": 0, "minimum": 0 }, + "query_cache_idle_timeout": { + "description": "Idle timeout (milliseconds) for query (AST) cache entries; an entry not accessed within this window is dropped by the maintenance sweep. `0` disables idle expiry.\n\n_Default:_ `0`\n\n", + "type": "integer", + "format": "uint", + "default": 0, + "minimum": 0 + }, "query_cache_limit": { "description": "Limit on the number of statements saved in the statement cache used to accelerate query parsing.\n\n_Default:_ `50000`\n\n", "type": "integer", @@ -987,13 +994,6 @@ "default": 0, "minimum": 0 }, - "query_cache_ttl": { - "description": "Time-to-idle (seconds) for query (AST) cache entries; an entry not accessed within this window is dropped by the maintenance sweep. `0` disables idle expiry.\n\n_Default:_ `0`\n\n", - "type": "integer", - "format": "uint", - "default": 0, - "minimum": 0 - }, "query_log": { "description": "Path to a file where all queries are logged. Logging every query is slow; do not use in production.", "type": [ diff --git a/pgdog-config/src/general.rs b/pgdog-config/src/general.rs index 76e669d09..b23807ab6 100644 --- a/pgdog-config/src/general.rs +++ b/pgdog-config/src/general.rs @@ -388,13 +388,13 @@ pub struct General { #[serde(default = "General::query_cache_memory_limit")] pub query_cache_memory_limit: usize, - /// Time-to-idle (seconds) for query (AST) cache entries; an entry not accessed within this window is dropped by the maintenance sweep. `0` disables idle expiry. + /// Idle timeout (milliseconds) for query (AST) cache entries; an entry not accessed within this window is dropped by the maintenance sweep. `0` disables idle expiry. /// /// _Default:_ `0` /// - /// - #[serde(default = "General::query_cache_ttl")] - pub query_cache_ttl: usize, + /// + #[serde(default = "General::query_cache_idle_timeout")] + pub query_cache_idle_timeout: usize, /// Toggle automatic creation of connection pools given the user name, database and password. /// @@ -873,7 +873,7 @@ impl Default for General { prepared_statements_limit: Self::prepared_statements_limit(), query_cache_limit: Self::query_cache_limit(), query_cache_memory_limit: Self::query_cache_memory_limit(), - query_cache_ttl: Self::query_cache_ttl(), + query_cache_idle_timeout: Self::query_cache_idle_timeout(), passthrough_auth: Self::default_passthrough_auth(), connect_timeout: Self::default_connect_timeout(), connect_attempt_delay: Self::default_connect_attempt_delay(), @@ -1349,8 +1349,8 @@ impl General { Self::env_or_default("PGDOG_QUERY_CACHE_MEMORY_LIMIT", 0) } - pub fn query_cache_ttl() -> usize { - Self::env_or_default("PGDOG_QUERY_CACHE_TTL", 0) + pub fn query_cache_idle_timeout() -> usize { + Self::env_or_default("PGDOG_QUERY_CACHE_IDLE_TIMEOUT", 0) } pub fn log_format() -> LogFormat { @@ -1772,7 +1772,7 @@ mod tests { let _guard = set_env_var("PGDOG_PREPARED_STATEMENTS_LIMIT", "1000"); let _guard = set_env_var("PGDOG_QUERY_CACHE_LIMIT", "500"); let _guard = set_env_var("PGDOG_QUERY_CACHE_MEMORY_LIMIT", "1048576"); - let _guard = set_env_var("PGDOG_QUERY_CACHE_TTL", "600"); + let _guard = set_env_var("PGDOG_QUERY_CACHE_IDLE_TIMEOUT", "600000"); let _guard = set_env_var("PGDOG_CONNECT_ATTEMPTS", "3"); let _guard = set_env_var("PGDOG_MIRROR_QUEUE", "256"); let _guard = set_env_var("PGDOG_MIRROR_EXPOSURE", "0.5"); @@ -1786,7 +1786,7 @@ mod tests { assert_eq!(General::prepared_statements_limit(), 1000); assert_eq!(General::query_cache_limit(), 500); assert_eq!(General::query_cache_memory_limit(), 1048576); - assert_eq!(General::query_cache_ttl(), 600); + assert_eq!(General::query_cache_idle_timeout(), 600000); assert_eq!(General::connect_attempts(), 3); assert_eq!(General::mirror_queue(), 256); assert_eq!(General::mirror_exposure(), 0.5); @@ -1800,7 +1800,7 @@ mod tests { let _guard = remove_env_var("PGDOG_PREPARED_STATEMENTS_LIMIT"); let _guard = remove_env_var("PGDOG_QUERY_CACHE_LIMIT"); let _guard = remove_env_var("PGDOG_QUERY_CACHE_MEMORY_LIMIT"); - let _guard = remove_env_var("PGDOG_QUERY_CACHE_TTL"); + let _guard = remove_env_var("PGDOG_QUERY_CACHE_IDLE_TIMEOUT"); let _guard = remove_env_var("PGDOG_CONNECT_ATTEMPTS"); let _guard = remove_env_var("PGDOG_MIRROR_QUEUE"); let _guard = remove_env_var("PGDOG_MIRROR_EXPOSURE"); @@ -1814,7 +1814,7 @@ mod tests { assert_eq!(General::prepared_statements_limit(), i64::MAX as usize); assert_eq!(General::query_cache_limit(), 1_000); assert_eq!(General::query_cache_memory_limit(), 0); - assert_eq!(General::query_cache_ttl(), 0); + assert_eq!(General::query_cache_idle_timeout(), 0); assert_eq!(General::connect_attempts(), 1); assert_eq!(General::mirror_queue(), 128); assert_eq!(General::mirror_exposure(), 1.0); diff --git a/pgdog/src/backend/databases.rs b/pgdog/src/backend/databases.rs index cc7da3800..0bd7cce51 100644 --- a/pgdog/src/backend/databases.rs +++ b/pgdog/src/backend/databases.rs @@ -107,7 +107,7 @@ pub fn init() -> Result<(), Error> { Cache::configure( config.config.general.query_cache_limit, config.config.general.query_cache_memory_limit, - config.config.general.query_cache_ttl, + config.config.general.query_cache_idle_timeout, ); // Start two-pc manager. @@ -159,7 +159,7 @@ pub fn reload() -> Result<(), Error> { Cache::configure( new_config.config.general.query_cache_limit, new_config.config.general.query_cache_memory_limit, - new_config.config.general.query_cache_ttl, + new_config.config.general.query_cache_idle_timeout, ); Ok(()) diff --git a/pgdog/src/frontend/prepared_statements/mod.rs b/pgdog/src/frontend/prepared_statements/mod.rs index d36a0c254..ebdc852b3 100644 --- a/pgdog/src/frontend/prepared_statements/mod.rs +++ b/pgdog/src/frontend/prepared_statements/mod.rs @@ -190,7 +190,7 @@ pub fn start_maintenance() { pub fn run_maintenance() { let capacity = config().config.general.prepared_statements_limit; PreparedStatements::global().write().close_unused(capacity); - // Drop AST cache entries idle beyond their time-to-idle (no-op when disabled). + // Drop AST cache entries idle beyond the idle timeout (no-op when disabled). Cache::sweep(); } diff --git a/pgdog/src/frontend/router/parser/cache/cache_impl.rs b/pgdog/src/frontend/router/parser/cache/cache_impl.rs index d55e5b470..7481a5585 100644 --- a/pgdog/src/frontend/router/parser/cache/cache_impl.rs +++ b/pgdog/src/frontend/router/parser/cache/cache_impl.rs @@ -67,7 +67,7 @@ pub(super) struct Inner { /// Sum of `Entry::size` across all cached entries. bytes: usize, /// Idle expiry: entries untouched for longer than this are swept (None = off). - ttl: Option, + idle_timeout: Option, /// Cache global stats. pub(super) stats: Stats, } @@ -103,9 +103,9 @@ impl Inner { } } - /// Drop entries not accessed within the time-to-idle window. + /// Drop entries not accessed within the idle window. fn sweep(&mut self) { - let Some(ttl) = self.ttl else { + let Some(idle_timeout) = self.idle_timeout else { return; }; let now = Instant::now(); @@ -115,7 +115,7 @@ impl Inner { let stale: Vec> = self .queries .iter() - .filter(|(_, e)| now.duration_since(e.accessed) >= ttl) + .filter(|(_, e)| now.duration_since(e.accessed) >= idle_timeout) .map(|(k, _)| k.clone()) .collect(); for key in stale { @@ -175,31 +175,31 @@ impl Cache { count_limit: 0, byte_limit: 0, bytes: 0, - ttl: None, + idle_timeout: None, stats: Stats::default(), })), } } /// Apply cache limits from configuration, evicting anything over the new - /// caps. A `count`, `bytes`, or `ttl_seconds` of 0 disables that limit. - pub fn configure(count: usize, bytes: usize, ttl_seconds: usize) { + /// caps. A `count`, `bytes`, or `idle_timeout_ms` of 0 disables that limit. + pub fn configure(count: usize, bytes: usize, idle_timeout_ms: usize) { let mut guard = CACHE.inner.lock(); guard.count_limit = count; guard.byte_limit = bytes; - guard.ttl = if ttl_seconds == 0 { + guard.idle_timeout = if idle_timeout_ms == 0 { None } else { - Some(Duration::from_secs(ttl_seconds as u64)) + Some(Duration::from_millis(idle_timeout_ms as u64)) }; guard.enforce(); debug!( - "ast cache limits: count={} bytes={} ttl={}s", - count, bytes, ttl_seconds + "ast cache limits: count={} bytes={} idle_timeout={}ms", + count, bytes, idle_timeout_ms ); } - /// Resize cache to a count capacity, keeping the memory and TTL limits. + /// Resize cache to a count capacity, keeping the memory and idle limits. pub fn resize(capacity: usize) { let mut guard = CACHE.inner.lock(); guard.count_limit = capacity.max(1); @@ -393,7 +393,7 @@ impl Cache { } /// Reset cache, removing all statements and setting stats to 0. The - /// configured count/memory/TTL limits are kept. + /// configured count/memory/idle limits are kept. pub fn reset() { let cache = Self::get(); let mut guard = cache.inner.lock(); @@ -414,13 +414,13 @@ mod tests { Ast::new_record("SELECT 1", QueryParserEngine::PgQueryProtobuf).expect("parse") } - fn inner(count_limit: usize, byte_limit: usize, ttl: Option) -> Inner { + fn inner(count_limit: usize, byte_limit: usize, idle_timeout: Option) -> Inner { Inner { queries: LruCache::unbounded(), count_limit, byte_limit, bytes: 0, - ttl, + idle_timeout, stats: Stats::default(), } } @@ -513,7 +513,7 @@ mod tests { } #[test] - fn sweep_is_noop_when_ttl_disabled() { + fn sweep_is_noop_when_idle_timeout_disabled() { let mut c = inner(0, 0, None); for i in 0..3 { c.insert(format!("q{i}").into(), ast(), 5); @@ -525,7 +525,7 @@ mod tests { #[test] fn sweep_keeps_fresh_entries() { - // TTL far larger than the entries' age: nothing is idle yet. + // Idle window far larger than the entries' age: nothing is idle yet. let mut c = inner(0, 0, Some(Duration::from_secs(3600))); for i in 0..3 { c.insert(format!("q{i}").into(), ast(), 5); @@ -537,7 +537,7 @@ mod tests { #[test] fn sweep_drops_idle_entries_and_updates_bytes() { - // TTL of zero: every entry is at or past its idle window. + // Zero idle window: every entry is at or past it. let mut c = inner(0, 0, Some(Duration::ZERO)); for i in 0..3 { c.insert(format!("q{i}").into(), ast(), 5); From 102fb3e899105f61ff05f48bd9d5665aa4a8d48b Mon Sep 17 00:00:00 2001 From: Igor Ohrimenko Date: Mon, 3 Aug 2026 09:59:13 +0300 Subject: [PATCH 3/3] test(cache): cover Cache::configure and Cache::resize The existing tests exercise the limit logic on a standalone Inner, so the public entry points that apply configuration to the global cache were never called: configure() with a non-zero idle timeout and resize() (flagged by codecov on the patch). Cover both, including the zero-capacity floor of resize(). --- .../router/parser/cache/cache_impl.rs | 57 +++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/pgdog/src/frontend/router/parser/cache/cache_impl.rs b/pgdog/src/frontend/router/parser/cache/cache_impl.rs index 7481a5585..5881818e0 100644 --- a/pgdog/src/frontend/router/parser/cache/cache_impl.rs +++ b/pgdog/src/frontend/router/parser/cache/cache_impl.rs @@ -546,4 +546,61 @@ mod tests { assert_eq!(c.queries.len(), 0); assert_eq!(c.bytes, 0); } + + #[test] + fn configure_and_resize_apply_limits_to_global_cache() { + // The cache is a process-wide singleton, so save the live limits and + // put them back at the end; assertions on entry counts are `<=` since + // other tests may share the cache. + let (count, bytes, idle) = { + let guard = CACHE.inner.lock(); + (guard.count_limit, guard.byte_limit, guard.idle_timeout) + }; + let keys: Vec> = (0..5).map(|i| format!("__configure_q{i}").into()).collect(); + + Cache::configure(3, 500, 30_000); + { + let mut guard = CACHE.inner.lock(); + assert_eq!(guard.count_limit, 3); + assert_eq!(guard.byte_limit, 500); + assert_eq!(guard.idle_timeout, Some(Duration::from_millis(30_000))); + for key in &keys { + guard.insert(key.clone(), ast(), 10); + } + assert!(guard.queries.len() <= 3, "configure() caps are enforced"); + } + + Cache::resize(1); + { + let guard = CACHE.inner.lock(); + assert_eq!(guard.count_limit, 1); + assert!( + guard.queries.len() <= 1, + "resize() evicts down to the new capacity" + ); + } + + // Zero means "unlimited" in configure(), but resize() keeps at least + // one entry. + Cache::resize(0); + assert_eq!(CACHE.inner.lock().count_limit, 1); + + Cache::configure( + count, + bytes, + idle.map(|d| d.as_millis() as usize).unwrap_or(0), + ); + { + let mut guard = CACHE.inner.lock(); + assert_eq!( + guard.idle_timeout, idle, + "restored, including the None branch" + ); + for key in &keys { + if let Some(entry) = guard.queries.pop(key.as_ref()) { + guard.bytes = guard.bytes.saturating_sub(entry.size); + } + } + } + } }