feat: Adds remove method to actively delete data from in-memory cache

pingora-memory-cache/src/lib.rs

@@ -122,6 +122,11 @@ impl<K: Hash, T: Clone + Send + Sync + 'static> MemoryCache<K, T> {
         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<Duration>) {
         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<i32, i32> = 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<i32, i32> = MemoryCache::new(10);

tinyufo/src/lib.rs

@@ -443,6 +443,22 @@ impl<K: Hash, T: Clone + Send + Sync + 'static> TinyUfo<K, T> {
         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<bool> {
         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);