Try doing the invalidation in the background

Signed-off-by: Alex Snaps <[email protected]>
Kuadrant · Dec 11, 2024 · 113aef4 · 113aef4
1 parent 4f3512f
commit 113aef4
Showing 1 changed file with 17 additions and 5 deletions.
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(),
         }
     }
 
@@ -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);
+                        }
+                    }
                 }
             }
         }