From 74880d5c85de65822477b3de9cb1f704871a1d95 Mon Sep 17 00:00:00 2001
From: caio_nogueira_27 <caio.nogueira@losch.lu>
Date: Mon, 15 Apr 2024 16:47:47 +0100
Subject: [PATCH 1/2] feat: Adds remove method to actively delete data from
 in-memory cache

---
 pingora-memory-cache/src/lib.rs | 18 ++++++++++++++++
 tinyufo/src/lib.rs              | 37 +++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+)

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<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);
diff --git a/tinyufo/src/lib.rs b/tinyufo/src/lib.rs
index 015afe0e..b402b574 100644
--- a/tinyufo/src/lib.rs
+++ b/tinyufo/src/lib.rs
@@ -443,6 +443,20 @@ 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);
@@ -552,6 +566,29 @@ mod tests {
         assert_eq!(cache.peek_queue(3), Some(SMALL));
         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() {

From a0753374289d71bd3a1fa427bc13a46ad409021e Mon Sep 17 00:00:00 2001
From: caio_nogueira_27 <caio.nogueira@losch.lu>
Date: Sat, 4 May 2024 13:08:20 +0100
Subject: [PATCH 2/2] fix: fmt

---
 tinyufo/src/lib.rs | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/tinyufo/src/lib.rs b/tinyufo/src/lib.rs
index b402b574..ac101820 100644
--- a/tinyufo/src/lib.rs
+++ b/tinyufo/src/lib.rs
@@ -449,10 +449,12 @@ impl<K: Hash, T: Clone + Send + Sync + 'static> TinyUfo<K, T> {
             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);
+                self.queues
+                    .small_weight
+                    .fetch_sub(p.weight as usize, SeqCst);
             }
         });
-        
+
         //estimator should not be updated due to collision possibility
         self.buckets.remove(&key);
     }
@@ -566,7 +568,7 @@ mod tests {
         assert_eq!(cache.peek_queue(3), Some(SMALL));
         assert_eq!(cache.peek_queue(4), Some(SMALL));
     }
-    
+
     #[test]
     fn test_remove_key() {
         let cache = TinyUfo::new(5, 5);
@@ -587,7 +589,6 @@ mod tests {
         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]