From b72a45049359b97c87fe74f64645fccbce475e12 Mon Sep 17 00:00:00 2001 From: Igor Ohrimenko Date: Wed, 29 Jul 2026 11:01:08 +0300 Subject: [PATCH 01/10] fix(stats): count hash table capacity in prepared statements memory accounting The global prepared statements cache reports memory usage by summing live entries only. Hash tables allocate capacity, not len: after a spike of unique prepared statements the spare capacity left behind dominates actual memory use but is invisible to the metric, which reports near zero while the process holds gigabytes. Count allocated slots (plus control bytes) in the MemoryUsage impls for HashMap and HashSet, make them generic over the hasher so FnvHashSet is covered, and export the statements table capacity as a new prepared_statements_capacity gauge: a capacity far above prepared_statements signals memory held from a past spike. --- .../prepared_statements/global_cache.rs | 37 +++++++++++- pgdog/src/stats/memory.rs | 56 +++++++++++++++++-- pgdog/src/stats/query_cache.rs | 14 ++++- 3 files changed, 98 insertions(+), 9 deletions(-) diff --git a/pgdog/src/frontend/prepared_statements/global_cache.rs b/pgdog/src/frontend/prepared_statements/global_cache.rs index 9d50e84fc..328f43038 100644 --- a/pgdog/src/frontend/prepared_statements/global_cache.rs +++ b/pgdog/src/frontend/prepared_statements/global_cache.rs @@ -122,7 +122,7 @@ impl MemoryUsage for GlobalCache { + self.names.memory_usage() + self.counter.memory_usage() + self.versions.memory_usage() - + self.unused.capacity() * 1usize.memory_usage() + + self.unused.memory_usage() } } @@ -286,6 +286,12 @@ impl GlobalCache { self.statements.len() } + /// Number of slots allocated by the statements table. A capacity far + /// above `len` means the cache is holding on to memory from a past spike. + pub fn capacity(&self) -> usize { + self.statements.capacity() + } + /// True if the local cache is empty. pub fn is_empty(&self) -> bool { self.len() == 0 @@ -650,4 +656,33 @@ mod test { assert!(cache.names.is_empty()); assert!(cache.unused.is_empty()); } + + #[test] + fn test_memory_usage_counts_table_capacity() { + let mut cache = GlobalCache::default(); + for i in 0..10_000 { + let parse = Parse::named("s", format!("SELECT {}", i)); + cache.insert(&parse); + } + let spike_capacity = cache.capacity(); + assert!(spike_capacity >= 10_000); + + for i in 1..=10_000 { + cache.close(&global_name(i)); + } + cache.close_unused(100); + assert_eq!(cache.len(), 100); + + // The table holds on to its allocation after the spike (capacity() + // may dip slightly due to tombstones); the accounting must report + // that memory. + assert!(cache.capacity() * 2 >= spike_capacity); + let table_floor = cache.capacity() * (std::mem::size_of::<(CacheKey, CachedStmt)>() + 1); + assert!( + cache.memory_usage() >= table_floor, + "memory_usage {} must include table capacity {}", + cache.memory_usage(), + table_floor + ); + } } diff --git a/pgdog/src/stats/memory.rs b/pgdog/src/stats/memory.rs index d10645d07..afd40a5ce 100644 --- a/pgdog/src/stats/memory.rs +++ b/pgdog/src/stats/memory.rs @@ -54,12 +54,17 @@ impl MemoryUsage for Vec { } } -impl MemoryUsage for HashMap { +impl MemoryUsage for HashMap { #[inline(always)] fn memory_usage(&self) -> usize { - self.iter() - .map(|(k, v)| k.memory_usage() + v.memory_usage()) - .sum::() + // The table allocates capacity() slots (plus one control byte each), + // not len(): spare capacity left behind by removed entries still + // occupies memory and has to be counted. + self.capacity() * (std::mem::size_of::<(K, V)>() + 1) + + self + .iter() + .map(|(k, v)| k.memory_usage() + v.memory_usage()) + .sum::() } } @@ -72,10 +77,12 @@ impl MemoryUsage for BTreeMap { } } -impl MemoryUsage for HashSet { +impl MemoryUsage for HashSet { #[inline(always)] fn memory_usage(&self) -> usize { - self.iter().map(|v| v.memory_usage()).sum::() + // Same as HashMap: count allocated slots, not just live entries. + self.capacity() * (std::mem::size_of::() + 1) + + self.iter().map(|v| v.memory_usage()).sum::() } } @@ -101,3 +108,40 @@ impl MemoryUsage for Bytes { 0 } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn hash_map_counts_spare_capacity() { + let mut map: HashMap = HashMap::new(); + for i in 0..1000 { + map.insert(i, i); + } + let capacity = map.capacity(); + for i in 0..1000 { + map.remove(&i); + } + assert!(map.is_empty()); + // The allocation survives removals; capacity() may dip slightly + // due to tombstones but stays the same order of magnitude. + assert!(map.capacity() * 2 >= capacity); + assert!( + map.memory_usage() >= map.capacity() * (std::mem::size_of::<(usize, usize)>() + 1), + "spare capacity left by removed entries must be counted" + ); + } + + #[test] + fn hash_set_counts_spare_capacity() { + let mut set: HashSet = HashSet::new(); + for i in 0..1000 { + set.insert(i); + } + let capacity = set.capacity(); + set.clear(); + assert_eq!(set.capacity(), capacity); + assert!(set.memory_usage() >= capacity * (std::mem::size_of::() + 1)); + } +} diff --git a/pgdog/src/stats/query_cache.rs b/pgdog/src/stats/query_cache.rs index 472b60939..a9ca9459d 100644 --- a/pgdog/src/stats/query_cache.rs +++ b/pgdog/src/stats/query_cache.rs @@ -19,15 +19,16 @@ pub struct QueryCache { stats: Stats, len: usize, prepared_statements: usize, + prepared_statements_capacity: usize, prepared_statements_memory: usize, } impl QueryCache { pub(crate) fn load() -> Self { - let (prepared_statements, prepared_statements_memory) = { + let (prepared_statements, prepared_statements_capacity, prepared_statements_memory) = { let global = PreparedStatements::global(); let guard = global.read(); - (guard.len(), guard.memory_usage()) + (guard.len(), guard.capacity(), guard.memory_usage()) }; let (stats, len) = Cache::stats(); @@ -36,6 +37,7 @@ impl QueryCache { stats, len, prepared_statements, + prepared_statements_capacity, prepared_statements_memory, } } @@ -90,6 +92,12 @@ impl QueryCache { value: self.prepared_statements, gauge: true, }), + Metric::new(QueryCacheMetric { + name: "prepared_statements_capacity".into(), + help: "Number of slots allocated by the prepared statements cache".into(), + value: self.prepared_statements_capacity, + gauge: true, + }), Metric::new(QueryCacheMetric { name: "prepared_statements_memory_used".into(), help: "Amount of bytes used for the prepared statements cache".into(), @@ -173,6 +181,7 @@ mod tests { }, len: 5, prepared_statements: 6, + prepared_statements_capacity: 8, prepared_statements_memory: 7, }; @@ -189,6 +198,7 @@ mod tests { "query_cache_parse_time".to_string(), "query_cache_fingerprints".to_string(), "prepared_statements".to_string(), + "prepared_statements_capacity".to_string(), "prepared_statements_memory_used".to_string(), ] ); From 23699705f7c8d3a22b706e7636340eeb60ce018b Mon Sep 17 00:00:00 2001 From: Igor Ohrimenko Date: Wed, 29 Jul 2026 13:57:34 +0300 Subject: [PATCH 02/10] fix(memory): shrink prepared statement cache tables after mass removal Hash tables never return capacity to the allocator: after a spike of unique prepared statements from long-lived clients, the global cache tables keep gigabytes of empty buckets for the lifetime of the process. Shrink the tables in the maintenance sweep once they are mostly empty (capacity > 8x len) and large enough to matter (> 4096 slots), leaving 2x headroom to avoid rehashing on every sweep. Also shrink on reset(). --- .../prepared_statements/global_cache.rs | 96 +++++++++++++++++-- 1 file changed, 87 insertions(+), 9 deletions(-) diff --git a/pgdog/src/frontend/prepared_statements/global_cache.rs b/pgdog/src/frontend/prepared_statements/global_cache.rs index 328f43038..4803f2859 100644 --- a/pgdog/src/frontend/prepared_statements/global_cache.rs +++ b/pgdog/src/frontend/prepared_statements/global_cache.rs @@ -338,9 +338,28 @@ impl GlobalCache { // unused will hold the remaining elements that was not extracted above self.unused = unused; + self.maybe_shrink(); + removed } + /// Return table memory to the allocator after a spike of unique + /// statements. Hysteresis (mostly-empty table, above a minimum size) + /// avoids rehashing on every sweep; shrinking to twice the current + /// len leaves headroom for new statements. + fn maybe_shrink(&mut self) { + const SHRINK_FACTOR: usize = 8; + const MIN_CAPACITY: usize = 4096; + + if self.statements.capacity() > MIN_CAPACITY + && self.statements.capacity() / SHRINK_FACTOR > self.statements.len() + { + self.statements.shrink_to(self.statements.len() * 2); + self.names.shrink_to(self.names.len() * 2); + self.unused.shrink_to(self.unused.len() * 2); + } + } + /// Remove statement from global cache. fn remove(&mut self, name: &str) { if let Some(stmt) = self.names.remove(name) { @@ -667,22 +686,81 @@ mod test { let spike_capacity = cache.capacity(); assert!(spike_capacity >= 10_000); + // The table allocates capacity, not len; the accounting + // must report that memory. + let table_floor = spike_capacity * (std::mem::size_of::<(CacheKey, CachedStmt)>() + 1); + assert!( + cache.memory_usage() >= table_floor, + "memory_usage {} must include table capacity {}", + cache.memory_usage(), + table_floor + ); + } + + #[test] + fn test_close_unused_shrinks_tables_after_spike() { + let mut cache = GlobalCache::default(); + for i in 0..10_000 { + let parse = Parse::named("s", format!("SELECT {}", i)); + cache.insert(&parse); + } + let spike_capacity = cache.capacity(); + let spike_memory = cache.memory_usage(); + for i in 1..=10_000 { cache.close(&global_name(i)); } cache.close_unused(100); assert_eq!(cache.len(), 100); - // The table holds on to its allocation after the spike (capacity() - // may dip slightly due to tombstones); the accounting must report - // that memory. - assert!(cache.capacity() * 2 >= spike_capacity); - let table_floor = cache.capacity() * (std::mem::size_of::<(CacheKey, CachedStmt)>() + 1); + // The sweep returns table memory to the allocator. assert!( - cache.memory_usage() >= table_floor, - "memory_usage {} must include table capacity {}", - cache.memory_usage(), - table_floor + cache.capacity() < spike_capacity / 8, + "capacity {} must shrink well below spike {}", + cache.capacity(), + spike_capacity ); + assert!(cache.memory_usage() < spike_memory / 8); + + // Statements that survived the sweep are still usable. + let survivors: Vec = cache.names().keys().cloned().collect(); + assert_eq!(survivors.len(), 100); + for name in survivors { + assert!(cache.parse(&name).is_some()); + } + } + + #[test] + fn test_no_shrink_below_min_capacity() { + let mut cache = GlobalCache::default(); + for i in 0..1_000 { + let parse = Parse::named("s", format!("SELECT {}", i)); + cache.insert(&parse); + } + let capacity = cache.capacity(); + + for i in 1..=1_000 { + cache.close(&global_name(i)); + } + cache.close_unused(10); + + // Small tables are not worth rehashing. + assert!(cache.capacity() >= capacity / 2); + } + + #[test] + fn test_no_shrink_when_mostly_full() { + let mut cache = GlobalCache::default(); + for i in 0..10_000 { + let parse = Parse::named("s", format!("SELECT {}", i)); + cache.insert(&parse); + } + let capacity = cache.capacity(); + + cache.close_unused(20_000); + + // Nothing was removed; the table must not shrink. + assert_eq!(cache.len(), 10_000); + assert!(cache.capacity() >= capacity / 2); } } From 4bf73f56e2e8bb202746867c58256c2bcb503708 Mon Sep 17 00:00:00 2001 From: Igor Ohrimenko Date: Wed, 29 Jul 2026 15:59:50 +0300 Subject: [PATCH 03/10] test: drop comments that restate the assertions --- pgdog/src/frontend/prepared_statements/global_cache.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/pgdog/src/frontend/prepared_statements/global_cache.rs b/pgdog/src/frontend/prepared_statements/global_cache.rs index 4803f2859..8bbec5bf3 100644 --- a/pgdog/src/frontend/prepared_statements/global_cache.rs +++ b/pgdog/src/frontend/prepared_statements/global_cache.rs @@ -713,7 +713,6 @@ mod test { cache.close_unused(100); assert_eq!(cache.len(), 100); - // The sweep returns table memory to the allocator. assert!( cache.capacity() < spike_capacity / 8, "capacity {} must shrink well below spike {}", @@ -759,7 +758,6 @@ mod test { cache.close_unused(20_000); - // Nothing was removed; the table must not shrink. assert_eq!(cache.len(), 10_000); assert!(cache.capacity() >= capacity / 2); } From f5b8986f4bdbdb4835072c213ff4fead4f570ea7 Mon Sep 17 00:00:00 2001 From: Igor Ohrimenko Date: Wed, 29 Jul 2026 16:30:16 +0300 Subject: [PATCH 04/10] test: hoist assert message values so coverage sees them --- .../frontend/prepared_statements/global_cache.rs | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/pgdog/src/frontend/prepared_statements/global_cache.rs b/pgdog/src/frontend/prepared_statements/global_cache.rs index 8bbec5bf3..67ec90719 100644 --- a/pgdog/src/frontend/prepared_statements/global_cache.rs +++ b/pgdog/src/frontend/prepared_statements/global_cache.rs @@ -689,11 +689,10 @@ mod test { // The table allocates capacity, not len; the accounting // must report that memory. let table_floor = spike_capacity * (std::mem::size_of::<(CacheKey, CachedStmt)>() + 1); + let usage = cache.memory_usage(); assert!( - cache.memory_usage() >= table_floor, - "memory_usage {} must include table capacity {}", - cache.memory_usage(), - table_floor + usage >= table_floor, + "memory_usage {usage} must include table capacity {table_floor}" ); } @@ -713,11 +712,10 @@ mod test { cache.close_unused(100); assert_eq!(cache.len(), 100); + let shrunk_capacity = cache.capacity(); assert!( - cache.capacity() < spike_capacity / 8, - "capacity {} must shrink well below spike {}", - cache.capacity(), - spike_capacity + shrunk_capacity < spike_capacity / 8, + "capacity {shrunk_capacity} must shrink well below spike {spike_capacity}" ); assert!(cache.memory_usage() < spike_memory / 8); From b8958d852063e58f48d851e34de28d997064bfe3 Mon Sep 17 00:00:00 2001 From: Igor Ohrimenko Date: Wed, 29 Jul 2026 16:42:54 +0300 Subject: [PATCH 05/10] review: document MemoryUsage contract, tighten shrink threshold, precise metric help The MemoryUsage trait mixes inline-size scalars with heap-counting containers; state explicitly that aggregated numbers are upper-bound approximations. Compare capacity against len * factor directly instead of integer division, and say which table the capacity gauge measures. --- pgdog/src/frontend/prepared_statements/global_cache.rs | 2 +- pgdog/src/stats/memory.rs | 7 +++++++ pgdog/src/stats/query_cache.rs | 3 ++- 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/pgdog/src/frontend/prepared_statements/global_cache.rs b/pgdog/src/frontend/prepared_statements/global_cache.rs index 67ec90719..b1ffcd6e4 100644 --- a/pgdog/src/frontend/prepared_statements/global_cache.rs +++ b/pgdog/src/frontend/prepared_statements/global_cache.rs @@ -352,7 +352,7 @@ impl GlobalCache { const MIN_CAPACITY: usize = 4096; if self.statements.capacity() > MIN_CAPACITY - && self.statements.capacity() / SHRINK_FACTOR > self.statements.len() + && self.statements.capacity() > self.statements.len() * SHRINK_FACTOR { self.statements.shrink_to(self.statements.len() * 2); self.names.shrink_to(self.names.len() * 2); diff --git a/pgdog/src/stats/memory.rs b/pgdog/src/stats/memory.rs index afd40a5ce..d028e821b 100644 --- a/pgdog/src/stats/memory.rs +++ b/pgdog/src/stats/memory.rs @@ -3,6 +3,13 @@ use lru::LruCache; use std::collections::{BTreeMap, HashMap, HashSet, VecDeque}; use std::hash::Hash; +/// Approximate number of bytes attributable to a value, for observability. +/// +/// Scalar impls report their inline size (so summing over a collection's +/// elements works), container impls report allocated capacity plus the sum +/// over live elements. As a result, the inline portion of live entries can +/// be counted both in the container's capacity term and in the element sum: +/// numbers are upper-bound approximations, not exact accounting. pub trait MemoryUsage { fn memory_usage(&self) -> usize; } diff --git a/pgdog/src/stats/query_cache.rs b/pgdog/src/stats/query_cache.rs index a9ca9459d..0f1e65a95 100644 --- a/pgdog/src/stats/query_cache.rs +++ b/pgdog/src/stats/query_cache.rs @@ -94,7 +94,8 @@ impl QueryCache { }), Metric::new(QueryCacheMetric { name: "prepared_statements_capacity".into(), - help: "Number of slots allocated by the prepared statements cache".into(), + help: "Number of slots allocated by the statements table of the prepared statements cache" + .into(), value: self.prepared_statements_capacity, gauge: true, }), From 26a3522396dcdd2946be7ed0dd19e56adcfdfb81 Mon Sep 17 00:00:00 2001 From: Igor Ohrimenko Date: Wed, 29 Jul 2026 16:47:07 +0300 Subject: [PATCH 06/10] docs: trim MemoryUsage contract to the essentials --- pgdog/src/stats/memory.rs | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/pgdog/src/stats/memory.rs b/pgdog/src/stats/memory.rs index d028e821b..184eda5d8 100644 --- a/pgdog/src/stats/memory.rs +++ b/pgdog/src/stats/memory.rs @@ -3,13 +3,10 @@ use lru::LruCache; use std::collections::{BTreeMap, HashMap, HashSet, VecDeque}; use std::hash::Hash; -/// Approximate number of bytes attributable to a value, for observability. +/// Approximate bytes attributable to a value, for metrics. /// -/// Scalar impls report their inline size (so summing over a collection's -/// elements works), container impls report allocated capacity plus the sum -/// over live elements. As a result, the inline portion of live entries can -/// be counted both in the container's capacity term and in the element sum: -/// numbers are upper-bound approximations, not exact accounting. +/// Scalars report their inline size, containers report allocated capacity +/// plus the sum over elements: treat results as an upper bound. pub trait MemoryUsage { fn memory_usage(&self) -> usize; } From cd8862befdbfe91b6d852b3561d9a641b62950ff Mon Sep 17 00:00:00 2001 From: Igor Ohrimenko Date: Wed, 29 Jul 2026 18:33:37 +0300 Subject: [PATCH 07/10] test: single-line asserts so coverage does not count panic-branch lines --- pgdog/src/frontend/prepared_statements/global_cache.rs | 4 ++-- pgdog/src/stats/memory.rs | 5 +++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/pgdog/src/frontend/prepared_statements/global_cache.rs b/pgdog/src/frontend/prepared_statements/global_cache.rs index b1ffcd6e4..173e0db2f 100644 --- a/pgdog/src/frontend/prepared_statements/global_cache.rs +++ b/pgdog/src/frontend/prepared_statements/global_cache.rs @@ -692,7 +692,7 @@ mod test { let usage = cache.memory_usage(); assert!( usage >= table_floor, - "memory_usage {usage} must include table capacity {table_floor}" + "memory_usage {usage} < table {table_floor}" ); } @@ -715,7 +715,7 @@ mod test { let shrunk_capacity = cache.capacity(); assert!( shrunk_capacity < spike_capacity / 8, - "capacity {shrunk_capacity} must shrink well below spike {spike_capacity}" + "capacity {shrunk_capacity} not shrunk" ); assert!(cache.memory_usage() < spike_memory / 8); diff --git a/pgdog/src/stats/memory.rs b/pgdog/src/stats/memory.rs index 184eda5d8..ee7c2010e 100644 --- a/pgdog/src/stats/memory.rs +++ b/pgdog/src/stats/memory.rs @@ -131,9 +131,10 @@ mod tests { // The allocation survives removals; capacity() may dip slightly // due to tombstones but stays the same order of magnitude. assert!(map.capacity() * 2 >= capacity); + let floor = map.capacity() * (std::mem::size_of::<(usize, usize)>() + 1); assert!( - map.memory_usage() >= map.capacity() * (std::mem::size_of::<(usize, usize)>() + 1), - "spare capacity left by removed entries must be counted" + map.memory_usage() >= floor, + "spare capacity must be counted" ); } From bfd152772a520c58b9a664081b1307b1863fca2e Mon Sep 17 00:00:00 2001 From: Igor Ohrimenko Date: Wed, 29 Jul 2026 18:35:35 +0300 Subject: [PATCH 08/10] test: assertions without panic-branch message lines rustfmt splits macro calls longer than its small-heuristics width, which puts assert messages back on their own never-executed lines; the values are in named locals right above, so the messages add nothing. --- pgdog/src/frontend/prepared_statements/global_cache.rs | 10 ++-------- pgdog/src/stats/memory.rs | 5 +---- 2 files changed, 3 insertions(+), 12 deletions(-) diff --git a/pgdog/src/frontend/prepared_statements/global_cache.rs b/pgdog/src/frontend/prepared_statements/global_cache.rs index 173e0db2f..2585ad29a 100644 --- a/pgdog/src/frontend/prepared_statements/global_cache.rs +++ b/pgdog/src/frontend/prepared_statements/global_cache.rs @@ -690,10 +690,7 @@ mod test { // must report that memory. let table_floor = spike_capacity * (std::mem::size_of::<(CacheKey, CachedStmt)>() + 1); let usage = cache.memory_usage(); - assert!( - usage >= table_floor, - "memory_usage {usage} < table {table_floor}" - ); + assert!(usage >= table_floor); } #[test] @@ -713,10 +710,7 @@ mod test { assert_eq!(cache.len(), 100); let shrunk_capacity = cache.capacity(); - assert!( - shrunk_capacity < spike_capacity / 8, - "capacity {shrunk_capacity} not shrunk" - ); + assert!(shrunk_capacity < spike_capacity / 8); assert!(cache.memory_usage() < spike_memory / 8); // Statements that survived the sweep are still usable. diff --git a/pgdog/src/stats/memory.rs b/pgdog/src/stats/memory.rs index ee7c2010e..3b6877f4d 100644 --- a/pgdog/src/stats/memory.rs +++ b/pgdog/src/stats/memory.rs @@ -132,10 +132,7 @@ mod tests { // due to tombstones but stays the same order of magnitude. assert!(map.capacity() * 2 >= capacity); let floor = map.capacity() * (std::mem::size_of::<(usize, usize)>() + 1); - assert!( - map.memory_usage() >= floor, - "spare capacity must be counted" - ); + assert!(map.memory_usage() >= floor); } #[test] From fa5f598872c864a2d14f1b1fed131065adf0ff0a Mon Sep 17 00:00:00 2001 From: Igor Ohrimenko Date: Thu, 30 Jul 2026 13:55:42 +0300 Subject: [PATCH 09/10] review: O(1) memory reporting for the global cache, shrink_to_fit Track heap bytes owned by cached statements incrementally on insert/remove/rewrite/describe instead of iterating millions of entries on every metrics scrape; report tables via capacity. Addresses reviewer concern about O(n) accounting on a hot path (and the unused set is no longer iterated either). Shrink with shrink_to_fit as suggested; the hysteresis guards stay. Drift is guarded by a test that recomputes the counter from scratch after every mutation path. --- .../prepared_statements/global_cache.rs | 135 +++++++++++++----- 1 file changed, 100 insertions(+), 35 deletions(-) diff --git a/pgdog/src/frontend/prepared_statements/global_cache.rs b/pgdog/src/frontend/prepared_statements/global_cache.rs index 2585ad29a..c98f99efb 100644 --- a/pgdog/src/frontend/prepared_statements/global_cache.rs +++ b/pgdog/src/frontend/prepared_statements/global_cache.rs @@ -27,13 +27,9 @@ pub struct Statement { impl MemoryUsage for Statement { #[inline] fn memory_usage(&self) -> usize { - self.parse.len() - + if let Some(ref row_description) = self.row_description { - row_description.memory_usage() - } else { - 0 - } - + self.cache_key.memory_usage() + // Same content accounting the cache aggregates in content_bytes, + // so SHOW PREPARED STATEMENTS agrees with the metric. + self.content_bytes() + self.cache_key.memory_usage() } } @@ -42,6 +38,18 @@ impl Statement { self.parse.query() } + /// Heap bytes owned by this statement; tracked incrementally + /// in the cache's `content_bytes`. + fn content_bytes(&self) -> usize { + self.parse.len() + + self.rewrite.as_ref().map(|p| p.len()).unwrap_or(0) + + self + .row_description + .as_ref() + .map(|r| r.memory_usage()) + .unwrap_or(0) + } + fn cache_key(&self) -> &CacheKey { &self.cache_key } @@ -113,16 +121,20 @@ pub struct GlobalCache { unused: HashSet, counter: usize, versions: usize, + /// Heap bytes owned by cached statements, maintained on + /// insert/remove so memory reporting stays O(1). + content_bytes: usize, } impl MemoryUsage for GlobalCache { #[inline] fn memory_usage(&self) -> usize { - self.statements.memory_usage() - + self.names.memory_usage() - + self.counter.memory_usage() - + self.versions.memory_usage() - + self.unused.memory_usage() + // O(1): tables report their allocation via capacity, entry + // contents are tracked incrementally as statements come and go. + self.statements.capacity() * (std::mem::size_of::<(CacheKey, CachedStmt)>() + 1) + + self.names.capacity() * (std::mem::size_of::<(String, Statement)>() + 1) + + self.unused.capacity() * (std::mem::size_of::() + 1) + + self.content_bytes } } @@ -172,14 +184,14 @@ impl GlobalCache { }, ); - self.names.insert( - name.clone(), - Statement { - parse, - cache_key, - ..Default::default() - }, - ); + let key = name.clone(); + let statement = Statement { + parse, + cache_key, + ..Default::default() + }; + self.content_bytes += key.capacity() + statement.content_bytes(); + self.names.insert(key, statement); (true, name) } @@ -208,15 +220,15 @@ impl GlobalCache { }, ); - self.names.insert( - name.clone(), - Statement { - parse, - version: self.versions, - cache_key: key, - ..Default::default() - }, - ); + let name_key = name.clone(); + let statement = Statement { + parse, + version: self.versions, + cache_key: key, + ..Default::default() + }; + self.content_bytes += name_key.capacity() + statement.content_bytes(); + self.names.insert(name_key, statement); name } @@ -224,7 +236,9 @@ impl GlobalCache { /// Rewrite prepared statement in the global cache. pub fn rewrite(&mut self, parse: &Parse) { if let Some(stmt) = self.names.get_mut(parse.name()) { + let old = stmt.rewrite.as_ref().map(|p| p.len()).unwrap_or(0); stmt.rewrite = Some(parse.clone()); + self.content_bytes = self.content_bytes.saturating_sub(old) + parse.len(); } } @@ -234,6 +248,7 @@ impl GlobalCache { if let Some(ref mut entry) = self.names.get_mut(name) && entry.row_description.is_none() { + self.content_bytes += row_description.memory_usage(); entry.row_description = Some(row_description); } } @@ -345,8 +360,7 @@ impl GlobalCache { /// Return table memory to the allocator after a spike of unique /// statements. Hysteresis (mostly-empty table, above a minimum size) - /// avoids rehashing on every sweep; shrinking to twice the current - /// len leaves headroom for new statements. + /// avoids rehashing on every sweep. fn maybe_shrink(&mut self) { const SHRINK_FACTOR: usize = 8; const MIN_CAPACITY: usize = 4096; @@ -354,16 +368,19 @@ impl GlobalCache { if self.statements.capacity() > MIN_CAPACITY && self.statements.capacity() > self.statements.len() * SHRINK_FACTOR { - self.statements.shrink_to(self.statements.len() * 2); - self.names.shrink_to(self.names.len() * 2); - self.unused.shrink_to(self.unused.len() * 2); + self.statements.shrink_to_fit(); + self.names.shrink_to_fit(); + self.unused.shrink_to_fit(); } } /// Remove statement from global cache. fn remove(&mut self, name: &str) { - if let Some(stmt) = self.names.remove(name) { + if let Some((key, stmt)) = self.names.remove_entry(name) { self.statements.remove(stmt.cache_key()); + self.content_bytes = self + .content_bytes + .saturating_sub(key.capacity() + stmt.content_bytes()); } } @@ -387,6 +404,16 @@ impl GlobalCache { pub fn statements(&self) -> &HashMap { &self.statements } + + /// Recompute content bytes from scratch; used by tests to verify + /// the incremental counter never drifts. + #[cfg(test)] + fn recomputed_content_bytes(&self) -> usize { + self.names + .iter() + .map(|(k, v)| k.capacity() + v.content_bytes()) + .sum() + } } #[cfg(test)] @@ -753,4 +780,42 @@ mod test { assert_eq!(cache.len(), 10_000); assert!(cache.capacity() >= capacity / 2); } + + #[test] + fn test_content_bytes_tracks_all_mutations() { + use crate::net::messages::Field; + + let mut cache = GlobalCache::default(); + for i in 0..500 { + let parse = Parse::named("s", format!("SELECT {}", i)); + cache.insert(&parse); + cache.insert(&parse); // duplicate must not grow the counter + } + for i in 0..50 { + let parse = Parse::named("v", format!("SELECT 'v{}'", i)); + cache.insert_anyway(&parse); + } + let rewrite = Parse::named("__pgdog_1", "SELECT 1, 2, 3"); + cache.rewrite(&rewrite); + cache.rewrite(&rewrite); // replacing a rewrite must not double-count + let rd = RowDescription::new(&[Field::text("name"), Field::bigint("id")]); + cache.insert_row_description("__pgdog_2", &rd); + cache.insert_row_description("__pgdog_2", &rd); // second call is a no-op + assert_eq!(cache.content_bytes, cache.recomputed_content_bytes()); + + for i in 1..=500 { + // inserted twice, so close twice + cache.close(&global_name(i)); + cache.close(&global_name(i)); + } + for i in 501..=550 { + cache.close(&global_name(i)); + } + cache.close_unused(10); + assert_eq!(cache.len(), 10); + assert_eq!(cache.content_bytes, cache.recomputed_content_bytes()); + + cache.close_unused(0); + assert_eq!(cache.content_bytes, 0); + } } From 7b78c7ea1a12cf2fd0305afee9a73abafbd0a972 Mon Sep 17 00:00:00 2001 From: Igor Ohrimenko Date: Wed, 12 Aug 2026 13:48:03 +0300 Subject: [PATCH 10/10] Match insert_row_description's by-value signature in the test main takes the RowDescription by value now; the test still passed a reference. --- pgdog/src/frontend/prepared_statements/global_cache.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pgdog/src/frontend/prepared_statements/global_cache.rs b/pgdog/src/frontend/prepared_statements/global_cache.rs index c98f99efb..826186cf3 100644 --- a/pgdog/src/frontend/prepared_statements/global_cache.rs +++ b/pgdog/src/frontend/prepared_statements/global_cache.rs @@ -799,8 +799,8 @@ mod test { cache.rewrite(&rewrite); cache.rewrite(&rewrite); // replacing a rewrite must not double-count let rd = RowDescription::new(&[Field::text("name"), Field::bigint("id")]); - cache.insert_row_description("__pgdog_2", &rd); - cache.insert_row_description("__pgdog_2", &rd); // second call is a no-op + cache.insert_row_description("__pgdog_2", rd.clone()); + cache.insert_row_description("__pgdog_2", rd); // second call is a no-op assert_eq!(cache.content_bytes, cache.recomputed_content_bytes()); for i in 1..=500 {