autocomplete: bound the in-memory caches

1 files changed, 24 insertions, 6 deletions

diff --git a/src/main.rs b/src/main.rs
index 8d11d76..56f81cb 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -25,6 +25,10 @@ use database::DbPool;
 /// How long a search cache entry stays fresh (10 seconds).
 const AC_SEARCH_CACHE_TTL: Duration = Duration::from_secs(10);
 
+/// Hard caps so autocomplete traffic cannot grow memory without bound.
+const AC_CACHE_MAX_ENTRIES: usize = 2_000;
+const AC_SEARCH_CACHE_MAX_ENTRIES: usize = 256;
+
 /// Per-user autocomplete stash: (UserId, filename) → SearchHit.
 type AcCache = tokio::sync::Mutex<HashMap<(serenity::UserId, String), copyparty::SearchHit>>;
 
@@ -70,7 +74,13 @@ impl Data {
         key: String,
         hit: copyparty::SearchHit,
     ) {
-        self.ac_cache.lock().await.insert((user_id, key), hit);
+        let mut cache = self.ac_cache.lock().await;
+        if cache.len() >= AC_CACHE_MAX_ENTRIES {
+            // Entries are only useful for the command invocation that
+            // triggered the autocomplete, so dropping all of them is safe.
+            cache.clear();
+        }
+        cache.insert((user_id, key), hit);
     }
 
     pub async fn ac_lookup(
@@ -88,9 +98,10 @@ impl Data {
     /// Look up search cache for an autocomplete query. Returns `None` if
     /// uncached or stale.
     pub async fn ac_search_cache_lookup(&self, query: &str) -> Option<Vec<copyparty::SearchHit>> {
-        let cache = self.ac_search_cache.lock().await;
+        let mut cache = self.ac_search_cache.lock().await;
         let (hits, cached_at) = cache.get(query)?;
         if cached_at.elapsed() > AC_SEARCH_CACHE_TTL {
+            cache.remove(query);
             None
         } else {
             Some(hits.clone())
@@ -99,10 +110,17 @@ impl Data {
 
     /// Store search results into the shared cache.
     pub async fn ac_search_cache_store(&self, query: String, hits: Vec<copyparty::SearchHit>) {
-        self.ac_search_cache
-            .lock()
-            .await
-            .insert(query, (hits, Instant::now()));
+        let mut cache = self.ac_search_cache.lock().await;
+        cache.retain(|_, (_, cached_at)| cached_at.elapsed() <= AC_SEARCH_CACHE_TTL);
+        if cache.len() >= AC_SEARCH_CACHE_MAX_ENTRIES
+            && let Some(oldest) = cache
+                .iter()
+                .min_by_key(|(_, (_, cached_at))| *cached_at)
+                .map(|(query, _)| query.clone())
+        {
+            cache.remove(&oldest);
+        }
+        cache.insert(query, (hits, Instant::now()));
     }
 }