From 4f3512fb9e651ba3c2448d57da1e3578fdd04f31 Mon Sep 17 00:00:00 2001
From: Alex Snaps <alex@wcgw.dev>
Date: Mon, 9 Dec 2024 14:39:33 -0500
Subject: [PATCH 1/2] Delete counters of qualified limits

Signed-off-by: Alex Snaps <alex@wcgw.dev>
---
 limitador/src/lib.rs               | 28 ++++++++++++++++++++++++++++
 limitador/src/storage/in_memory.rs | 10 +++++++++-
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/limitador/src/lib.rs b/limitador/src/lib.rs
index 6902028c..44977fb8 100644
--- a/limitador/src/lib.rs
+++ b/limitador/src/lib.rs
@@ -697,6 +697,7 @@ fn classify_limits_by_namespace(
 mod test {
     use crate::limit::{Context, Expression, Limit};
     use crate::RateLimiter;
+    use std::collections::HashMap;
 
     #[test]
     fn properly_updates_existing_limits() {
@@ -729,4 +730,31 @@ mod test {
             .unwrap();
         assert_eq!(r.counters.first().unwrap().max_value(), 50);
     }
+
+    #[test]
+    fn deletes_qualified_counters() {
+        let rl = RateLimiter::new(100);
+        let namespace = "foo";
+
+        let l = Limit::new(
+            namespace,
+            42,
+            100,
+            vec![],
+            vec![Expression::parse("x").unwrap()],
+        );
+        let ctx = Context::from(HashMap::from([("x".to_string(), "a".to_string())]));
+        rl.add_limit(l.clone());
+        let r = rl
+            .check_rate_limited_and_update(&namespace.into(), &ctx, 1, true)
+            .unwrap();
+        assert_eq!(r.counters.first().unwrap().remaining(), Some(41));
+        rl.delete_limit(&l).unwrap();
+
+        rl.add_limit(l.clone());
+        let r = rl
+            .check_rate_limited_and_update(&namespace.into(), &ctx, 1, true)
+            .unwrap();
+        assert_eq!(r.counters.first().unwrap().remaining(), Some(41));
+    }
 }
diff --git a/limitador/src/storage/in_memory.rs b/limitador/src/storage/in_memory.rs
index 28dc2119..c41fbc18 100644
--- a/limitador/src/storage/in_memory.rs
+++ b/limitador/src/storage/in_memory.rs
@@ -230,7 +230,15 @@ impl InMemoryStorage {
     }
 
     fn delete_counters_of_limit(&self, limit: &Limit) {
-        self.simple_limits.write().unwrap().remove(limit);
+        if limit.variables().is_empty() {
+            self.simple_limits.write().unwrap().remove(limit);
+        } else {
+            for (c, _) in self.qualified_counters.iter() {
+                if c.limit() == limit {
+                    self.qualified_counters.invalidate(&c);
+                }
+            }
+        }
     }
 
     fn counter_is_within_limits(counter: &Counter, current_val: Option<&u64>, delta: u64) -> bool {

From 113aef40866d8d13e39ac5141f45f356b40b44bc Mon Sep 17 00:00:00 2001
From: Alex Snaps <alex@wcgw.dev>
Date: Wed, 11 Dec 2024 07:26:31 -0500
Subject: [PATCH 2/2] Try doing the invalidation in the background

Signed-off-by: Alex Snaps <alex@wcgw.dev>
---
 limitador/src/storage/in_memory.rs | 22 +++++++++++++++++-----
 1 file changed, 17 insertions(+), 5 deletions(-)

diff --git a/limitador/src/storage/in_memory.rs b/limitador/src/storage/in_memory.rs
index c41fbc18..96211d4c 100644
--- a/limitador/src/storage/in_memory.rs
+++ b/limitador/src/storage/in_memory.rs
@@ -2,7 +2,8 @@ use crate::counter::Counter;
 use crate::limit::{Context, Limit, Namespace};
 use crate::storage::atomic_expiring_value::AtomicExpiringValue;
 use crate::storage::{Authorization, CounterStorage, StorageErr};
-use moka::sync::Cache;
+use moka::sync::{Cache, CacheBuilder};
+use moka::PredicateError;
 use std::collections::btree_map::Entry;
 use std::collections::{BTreeMap, HashMap, HashSet};
 use std::ops::Deref;
@@ -198,7 +199,9 @@ impl InMemoryStorage {
     pub fn new(cache_size: u64) -> Self {
         Self {
             simple_limits: RwLock::new(BTreeMap::new()),
-            qualified_counters: Cache::new(cache_size),
+            qualified_counters: CacheBuilder::new(cache_size)
+                .support_invalidation_closures()
+                .build(),
         }
     }
 
@@ -233,9 +236,18 @@ impl InMemoryStorage {
         if limit.variables().is_empty() {
             self.simple_limits.write().unwrap().remove(limit);
         } else {
-            for (c, _) in self.qualified_counters.iter() {
-                if c.limit() == limit {
-                    self.qualified_counters.invalidate(&c);
+            let l = limit.clone();
+            match self
+                .qualified_counters
+                .invalidate_entries_if(move |c, _| c.limit() == &l)
+            {
+                Ok(_) => {}
+                Err(PredicateError::InvalidationClosuresDisabled) => {
+                    for (c, _) in self.qualified_counters.iter() {
+                        if c.limit() == limit {
+                            self.qualified_counters.invalidate(&c);
+                        }
+                    }
                 }
             }
         }