diff --git a/pgdog/src/frontend/prepared_statements/global_cache.rs b/pgdog/src/frontend/prepared_statements/global_cache.rs index f61d388a7..df09d52ba 100644 --- a/pgdog/src/frontend/prepared_statements/global_cache.rs +++ b/pgdog/src/frontend/prepared_statements/global_cache.rs @@ -28,14 +28,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() - + self.evict_on_close.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() + self.evict_on_close.memory_usage() } } @@ -44,6 +39,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.clone() } @@ -115,16 +122,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.len() * std::mem::size_of::() + // 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 } } @@ -166,14 +177,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) } @@ -202,15 +213,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 } @@ -218,7 +229,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(); } } @@ -229,6 +242,7 @@ impl GlobalCache { && entry.row_description.is_none() { entry.row_description = Some(row_description.clone()); + self.content_bytes += row_description.memory_usage(); } } @@ -237,8 +251,12 @@ impl GlobalCache { self.statements.clear(); self.names.clear(); self.unused.clear(); + self.statements.shrink_to_fit(); + self.names.shrink_to_fit(); + self.unused.shrink_to_fit(); self.counter = 0; self.versions = 0; + self.content_bytes = 0; } /// Get the query string stored in the global cache @@ -289,6 +307,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 @@ -326,13 +350,34 @@ impl GlobalCache { self.remove(&global_name(*counter)); } + self.maybe_shrink(); + remove.len() } + /// 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. + fn maybe_shrink(&mut self) { + const SHRINK_FACTOR: usize = 8; + const MIN_CAPACITY: usize = 4096; + + if self.statements.capacity() > MIN_CAPACITY + && self.statements.capacity() > self.statements.len() * SHRINK_FACTOR + { + 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()); } } @@ -356,6 +401,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)] @@ -612,4 +667,120 @@ 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); + + // 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!(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); + + let shrunk_capacity = cache.capacity(); + assert!(shrunk_capacity < spike_capacity / 8); + 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); + + 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); + } } diff --git a/pgdog/src/stats/memory.rs b/pgdog/src/stats/memory.rs index d10645d07..3b6877f4d 100644 --- a/pgdog/src/stats/memory.rs +++ b/pgdog/src/stats/memory.rs @@ -3,6 +3,10 @@ use lru::LruCache; use std::collections::{BTreeMap, HashMap, HashSet, VecDeque}; use std::hash::Hash; +/// Approximate bytes attributable to a value, for metrics. +/// +/// 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; } @@ -54,12 +58,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 +81,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 +112,38 @@ 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); + let floor = map.capacity() * (std::mem::size_of::<(usize, usize)>() + 1); + assert!(map.memory_usage() >= floor); + } + + #[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..0f1e65a95 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,13 @@ impl QueryCache { value: self.prepared_statements, gauge: true, }), + Metric::new(QueryCacheMetric { + name: "prepared_statements_capacity".into(), + help: "Number of slots allocated by the statements table of 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 +182,7 @@ mod tests { }, len: 5, prepared_statements: 6, + prepared_statements_capacity: 8, prepared_statements_memory: 7, }; @@ -189,6 +199,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(), ] );