1 files changed, 52 insertions, 3 deletions

diff --git a/pingora-cache/src/memory.rs b/pingora-cache/src/memory.rs
index 6a66577..afe57ae 100644
--- a/pingora-cache/src/memory.rs
+++ b/pingora-cache/src/memory.rs
@@ -307,12 +307,12 @@ impl Storage for MemCache {
     }
 
     async fn purge(&'static self, key: &CompactCacheKey, _trace: &SpanHandle) -> Result<bool> {
-        // TODO: purge partial
-
         // This usually purges the primary key because, without a lookup, variance key is usually
         // empty
         let hash = key.combined();
-        Ok(self.cached.write().remove(&hash).is_some())
+        let temp_removed = self.temp.write().remove(&hash).is_some();
+        let cache_removed = self.cached.write().remove(&hash).is_some();
+        Ok(temp_removed || cache_removed)
     }
 
     async fn update_meta(
@@ -507,4 +507,53 @@ mod test {
         let data = hit_handler2.read_body().await.unwrap();
         assert!(data.is_none());
     }
+
+    #[tokio::test]
+    async fn test_purge_partial() {
+        static MEM_CACHE: Lazy<MemCache> = Lazy::new(MemCache::new);
+        let cache = &MEM_CACHE;
+
+        let key = CacheKey::new("", "a", "1").to_compact();
+        let hash = key.combined();
+        let meta = (
+            "meta_key".as_bytes().to_vec(),
+            "meta_value".as_bytes().to_vec(),
+        );
+
+        let temp_obj = TempObject::new(meta);
+        cache.temp.write().insert(hash.clone(), temp_obj);
+
+        assert!(cache.temp.read().contains_key(&hash));
+
+        let result = cache.purge(&key, &Span::inactive().handle()).await;
+        assert!(result.is_ok());
+
+        assert!(!cache.temp.read().contains_key(&hash));
+    }
+
+    #[tokio::test]
+    async fn test_purge_complete() {
+        static MEM_CACHE: Lazy<MemCache> = Lazy::new(MemCache::new);
+        let cache = &MEM_CACHE;
+
+        let key = CacheKey::new("", "a", "1").to_compact();
+        let hash = key.combined();
+        let meta = (
+            "meta_key".as_bytes().to_vec(),
+            "meta_value".as_bytes().to_vec(),
+        );
+        let body = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 0];
+        let cache_obj = CacheObject {
+            meta,
+            body: Arc::new(body),
+        };
+        cache.cached.write().insert(hash.clone(), cache_obj);
+
+        assert!(cache.cached.read().contains_key(&hash));
+
+        let result = cache.purge(&key, &Span::inactive().handle()).await;
+        assert!(result.is_ok());
+
+        assert!(!cache.cached.read().contains_key(&hash));
+    }
 }