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); + } + } } } }