From 1c6eed066b57e5bf387b7ebcad9e447515dece80 Mon Sep 17 00:00:00 2001
From: Ayodeji Akinola <akinayodeji4all@gmail.com>
Date: Fri, 25 Oct 2024 09:12:43 +0000
Subject: feat/remove-keys-from-tiny-ufo

Includes-commit: ac601366bd02f97c132e0c1e762e17a9c962b9a6
Includes-commit: f8f456ae72db0f78aa5823f5e4121322ecbb9816
Replicated-from: https://github.com/cloudflare/pingora/pull/442
Co-authored-by: Matthew Gumport <mbg@cloudflare.com>
---
 .bleep             |  2 +-
 tinyufo/src/lib.rs | 59 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/.bleep b/.bleep
index 51c295c..6be531e 100644
--- a/.bleep
+++ b/.bleep
@@ -1 +1 @@
-72b8880751c39d3cfa0998f0a7d4f5fda04ffcfb
\ No newline at end of file
+93a8806fd5a1e6df065f20bc40e9594fde0a21db
\ No newline at end of file
diff --git a/tinyufo/src/lib.rs b/tinyufo/src/lib.rs
index eb4cbaf..5587e1e 100644
--- a/tinyufo/src/lib.rs
+++ b/tinyufo/src/lib.rs
@@ -424,6 +424,34 @@ impl<K: Hash, T: Clone + Send + Sync + 'static> TinyUfo<K, T> {
         self.queues.admit(key, data, weight, false, &self.buckets)
     }
 
+    /// Remove the given key from the cache if it exists
+    ///
+    /// Returns Some(T) if the key was found and removed, None otherwise
+    pub fn remove(&self, key: &K) -> Option<T> {
+        let key = self.random_status.hash_one(key);
+
+        // Get data and update weights
+        let result = self.buckets.get_map(&key, |bucket| {
+            let data = bucket.data.clone();
+            let weight = bucket.weight;
+
+            // Update weight based on queue location
+            if bucket.queue.is_main() {
+                self.queues.main_weight.fetch_sub(weight as usize, SeqCst);
+            } else {
+                self.queues.small_weight.fetch_sub(weight as usize, SeqCst);
+            }
+
+            data
+        });
+
+        // If we found and processed the item, remove it from buckets
+        if result.is_some() {
+            self.buckets.remove(&key);
+        }
+
+        result
+    }
     /// Always put the key value to the [TinyUfo]
     ///
     /// Return a list of [KV] of key and `T` that are evicted
@@ -728,4 +756,35 @@ mod tests {
             assert_eq!(cache.peek_queue(k), Some(SMALL));
         }
     }
+    #[test]
+    fn test_remove() {
+        let mut cache = TinyUfo::new(5, 5);
+        cache.random_status = RandomState::with_seeds(2, 3, 4, 5);
+
+        cache.put(1, 1, 1);
+        cache.put(2, 2, 2);
+        cache.put(3, 3, 2);
+
+        assert_eq!(cache.remove(&1), Some(1));
+        assert_eq!(cache.remove(&3), Some(3));
+        assert_eq!(cache.get(&1), None);
+        assert_eq!(cache.get(&3), None);
+
+        // Verify empty keys get evicted when cache fills up
+        // Fill cache to trigger eviction
+        cache.put(5, 5, 2);
+        cache.put(6, 6, 2);
+        cache.put(7, 7, 2);
+
+        // The removed items (1, 3) should be naturally evicted now
+        // and new items should be in cache
+        assert_eq!(cache.get(&1), None);
+        assert_eq!(cache.get(&3), None);
+        assert!(cache.get(&5).is_some() || cache.get(&6).is_some() || cache.get(&7).is_some());
+
+        // Test weights after eviction cycles
+        let total_weight =
+            cache.queues.small_weight.load(SeqCst) + cache.queues.main_weight.load(SeqCst);
+        assert!(total_weight <= 5); // Should not exceed limit
+    }
 }
-- 
cgit v1.2.3