diff --git a/pgdog/src/frontend/prepared_statements/global_cache.rs b/pgdog/src/frontend/prepared_statements/global_cache.rs index bb0689f0..51241af7 100644 --- a/pgdog/src/frontend/prepared_statements/global_cache.rs +++ b/pgdog/src/frontend/prepared_statements/global_cache.rs @@ -32,15 +32,17 @@ pub struct GlobalCache { names: HashMap, unused: HashSet, counter: Counter, + content_bytes: usize, } impl MemoryUsage for GlobalCache { #[inline] fn memory_usage(&self) -> usize { - self.statements.memory_usage() - + self.names.memory_usage() + 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.counter.memory_usage() - + self.unused.capacity() * 1usize.memory_usage() + + self.content_bytes } } @@ -119,18 +121,32 @@ impl GlobalCache { /// Rewrite prepared statement in the global cache. pub(crate) fn rewrite(&mut self, parse: &Parse) { - if let Some(stmt) = self.names.get_mut(parse.name()) { + let delta = self.names.get_mut(parse.name()).map(|stmt| { + let before = stmt.content_bytes(); stmt.set_rewrite(parse); + (before, stmt.content_bytes()) + }); + + if let Some((before, after)) = delta { + self.content_bytes = self.content_bytes.saturating_sub(before) + after; } } /// Client sent a Describe for a prepared statement and received a RowDescription. /// We record the RowDescription for later use by the results decoder. pub fn insert_row_description(&mut self, name: &str, row_description: RowDescription) { - if let Some(entry) = self.names.get_mut(name) - && entry.row_description.is_none() - { - entry.row_description = Some(row_description); + let added = self + .names + .get_mut(name) + .filter(|entry| entry.row_description.is_none()) + .map(|entry| { + let added = row_description.memory_usage(); + entry.row_description = Some(row_description); + added + }); + + if let Some(added) = added { + self.content_bytes += added; } } /// Get the Parse message for a globally unique prepared statement @@ -186,6 +202,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. #[cfg(test)] pub(crate) fn is_empty(&self) -> bool { @@ -233,6 +255,8 @@ impl GlobalCache { // unused will hold the remaining elements that was not extracted above self.unused = unused; + self.maybe_shrink(); + removed } @@ -246,9 +270,36 @@ impl GlobalCache { &self.statements } + /// 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(); + } + } + + #[cfg(test)] + fn recomputed_content_bytes(&self) -> usize { + self.names + .iter() + .map(|(k, v)| k.capacity() + v.content_bytes()) + .sum() + } + /// Remove statement from global cache. fn remove(&mut self, name: &str) { if let Some(stmt) = self.names.remove(name) { + self.content_bytes = self + .content_bytes + .saturating_sub(name.len() + stmt.content_bytes()); self.statements.remove(stmt.cache_key()); } } @@ -279,7 +330,9 @@ impl GlobalCache { used: 1, }, ); - self.names.insert(name.to_owned(), statement); + let key = name.to_owned(); + self.content_bytes += key.capacity() + statement.content_bytes(); + self.names.insert(key, statement); } } @@ -612,4 +665,117 @@ 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); + + 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); + + 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); + + 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); + } + for i in 0..50 { + cache.insert_prepare( + Bytes::from(format!("SELECT 'v{}'", i)), + &RewritePlan::default(), + ); + } + let rewrite = Parse::named("__pgdog_1", "SELECT 1, 2, 3"); + cache.rewrite(&rewrite); + cache.rewrite(&rewrite); + let row_description = RowDescription::new(&[Field::text("name"), Field::bigint("id")]); + cache.insert_row_description("__pgdog_2", row_description.clone()); + cache.insert_row_description("__pgdog_2", row_description); + assert_eq!(cache.content_bytes, cache.recomputed_content_bytes()); + + for i in 1..=500 { + 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/frontend/prepared_statements/statement.rs b/pgdog/src/frontend/prepared_statements/statement.rs index 990c8432..82252a3d 100644 --- a/pgdog/src/frontend/prepared_statements/statement.rs +++ b/pgdog/src/frontend/prepared_statements/statement.rs @@ -45,13 +45,7 @@ impl MemoryUsage for StatementType { impl MemoryUsage for Statement { #[inline] fn memory_usage(&self) -> usize { - self.stmt.memory_usage() - + if let Some(row_description) = &self.row_description { - row_description.memory_usage() - } else { - 0 - } - + self.cache_key.memory_usage() + self.content_bytes() + self.cache_key.memory_usage() } } @@ -92,6 +86,15 @@ impl Statement { &self.cache_key } + pub(super) fn content_bytes(&self) -> usize { + self.stmt.memory_usage() + + self + .row_description + .as_ref() + .map(|row_description| row_description.memory_usage()) + .unwrap_or_default() + } + pub(super) fn set_rewrite(&mut self, parse: &Parse) { if let StatementType::Parse { ref mut rewrite, .. diff --git a/pgdog/src/stats/memory.rs b/pgdog/src/stats/memory.rs index d10645d0..985de2dc 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,11 @@ 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::() + self.capacity() * (std::mem::size_of::() + 1) + + self.iter().map(|v| v.memory_usage()).sum::() } } @@ -101,3 +111,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 472b6093..0f1e65a9 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(), ] );