Merge pull request #399 from Kuadrant/qualified_deletes

Delete counters of qualified limits
Kuadrant · Jan 30, 2025 · 066b218 · 066b218
2 parents 9b9601f + 0acb833
commit 066b218
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 3 deletions.
diff --git 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
@@ -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(),
         }
     }
 
@@ -230,7 +233,21 @@ 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 {
+            let l = limit.clone();
+            if let Err(PredicateError::InvalidationClosuresDisabled) = self
+                .qualified_counters
+                .invalidate_entries_if(move |c, _| c.limit() == &l)
+            {
+                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 {