fix(memory): prepared statements cache retains gigabytes after statement spike

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<usize>,
     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::<usize>()
+        // 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::<usize>() + 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,23 +213,25 @@ 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
     }
 
     /// 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<CacheKey, CachedStmt> {
         &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<String> = 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);
+    }
 }