diff --git a/pingora-memory-cache/src/lib.rs b/pingora-memory-cache/src/lib.rs index 5e0254bd..bab7dbb3 100644 --- a/pingora-memory-cache/src/lib.rs +++ b/pingora-memory-cache/src/lib.rs @@ -122,6 +122,11 @@ impl MemoryCache { self.store.put(hashed_key, node, 1); } + pub fn remove(&self, key: &K) { + let hashed_key = self.hasher.hash_one(key); + self.store.remove(hashed_key); + } + pub(crate) fn force_put(&self, key: &K, value: T, ttl: Option) { if let Some(t) = ttl { if t.is_zero() { @@ -223,6 +228,19 @@ mod tests { assert_eq!(hit, CacheStatus::Hit); } + #[test] + fn test_remove() { + let cache: MemoryCache = MemoryCache::new(10); + cache.put(&1, 2, None); + let (res, hit) = cache.get(&1); + assert_eq!(res.unwrap(), 2); + assert_eq!(hit, CacheStatus::Hit); + cache.remove(&1); + let (res, hit) = cache.get(&1); + assert_eq!(res, None); + assert_eq!(hit, CacheStatus::Miss); + } + #[test] fn test_multi_get() { let cache: MemoryCache = MemoryCache::new(10); diff --git a/tinyufo/src/lib.rs b/tinyufo/src/lib.rs index 015afe0e..ac101820 100644 --- a/tinyufo/src/lib.rs +++ b/tinyufo/src/lib.rs @@ -443,6 +443,22 @@ impl TinyUfo { self.queues.admit(key, data, weight, true, &self.buckets) } + pub fn remove(&self, key: K) { + let key = self.random_status.hash_one(key); + self.buckets.get_map(&key, |p| { + if p.queue.is_main() { + self.queues.main_weight.fetch_sub(p.weight as usize, SeqCst); + } else { + self.queues + .small_weight + .fetch_sub(p.weight as usize, SeqCst); + } + }); + + //estimator should not be updated due to collision possibility + self.buckets.remove(&key); + } + #[cfg(test)] fn peek_queue(&self, key: K) -> Option { let key = self.random_status.hash_one(&key); @@ -553,6 +569,28 @@ mod tests { assert_eq!(cache.peek_queue(4), Some(SMALL)); } + #[test] + fn test_remove_key() { + let cache = TinyUfo::new(5, 5); + + cache.put(1, 1, 1); + cache.put(2, 2, 2); + cache.put(3, 3, 2); + // cache full now + + assert_eq!(cache.peek_queue(1), Some(SMALL)); + assert_eq!(cache.peek_queue(2), Some(SMALL)); + assert_eq!(cache.peek_queue(3), Some(SMALL)); + + cache.remove(1); + cache.put(4, 4, 1); + + assert_eq!(cache.peek_queue(1), None); + assert_eq!(cache.peek_queue(2), Some(SMALL)); + assert_eq!(cache.peek_queue(3), Some(SMALL)); + assert_eq!(cache.peek_queue(4), Some(SMALL)); + } + #[test] fn test_evict_entry_denied() { let cache = TinyUfo::new(5, 5);