From fa153561c6ebc4ae4a9d4f4d6170b999cc2e5b63 Mon Sep 17 00:00:00 2001
From: bjmeetsfo <bjmeetsfo@users.noreply.github.com>
Date: Sun, 6 Sep 2026 14:06:42 -0700
Subject: [PATCH] Release sharded cache handles from iterators

---
 src/runtime/multilayer_cache.rs | 29 +++++++++++++++++++++++++++++
 src/tests/mod.rs                | 21 +++++++++++++++++++++
 2 files changed, 50 insertions(+)

diff --git a/src/runtime/multilayer_cache.rs b/src/runtime/multilayer_cache.rs
index cd87b4c..28518e5 100644
--- a/src/runtime/multilayer_cache.rs
+++ b/src/runtime/multilayer_cache.rs
@@ -7121,6 +7121,35 @@ impl ShardedMultiLayerCache {
         for handle in handles {
             groups[self.shard_index_for_key(&handle.key)].push(handle);
         }
+        self.release_shard_groups(groups, released, started)
+    }
+
+    pub fn release_batch_iter<I>(&self, handles: I) -> usize
+    where
+        I: IntoIterator<Item = CachePinnedHandle>,
+    {
+        let started = Instant::now();
+        let mut groups = (0..self.shard_count())
+            .map(|_| Vec::<CachePinnedHandle>::new())
+            .collect::<Vec<_>>();
+        let mut released = 0usize;
+        for handle in handles {
+            groups[self.shard_index_for_key(&handle.key)].push(handle);
+            released = released.saturating_add(1);
+        }
+        self.release_shard_groups(groups, released, started)
+    }
+
+    fn release_shard_groups(
+        &self,
+        groups: Vec<Vec<CachePinnedHandle>>,
+        released: usize,
+        started: Instant,
+    ) -> usize {
+        if released == 0 {
+            self.sharded_stats.record_latency(started);
+            return 0;
+        }
         if released < Self::BATCH_FANOUT_THRESHOLD {
             for (index, group) in groups.into_iter().enumerate() {
                 if !group.is_empty() {
diff --git a/src/tests/mod.rs b/src/tests/mod.rs
index 8547945..55d364b 100644
--- a/src/tests/mod.rs
+++ b/src/tests/mod.rs
@@ -11534,6 +11534,27 @@ mod tests {
         assert_eq!(cache.stats().unpin_operations, 2);
     }
 
+    #[test]
+    fn sharded_release_batch_iter_drains_pins_across_shards() {
+        let dir = tempfile::tempdir().unwrap();
+        let cache = ShardedMultiLayerCache::new(
+            CacheOptions::new(1 << 20, 0, 0).with_ssd_paths(vec![dir.path().join("ssd")]),
+            4,
+        );
+        let first = CacheKey::string(1, "sharded-release-iter-a");
+        let second = CacheKey::string(2, "sharded-release-iter-b");
+        cache.put(first.clone(), b"alpha".to_vec()).unwrap();
+        cache.put(second.clone(), b"bravo".to_vec()).unwrap();
+
+        let handles = cache.acquire_batch(&[first, second]).unwrap();
+        assert_eq!(cache.stats().pinned_entries, 2);
+
+        let released = cache.release_batch_iter(handles.into_iter().flatten());
+        assert_eq!(released, 2);
+        assert_eq!(cache.stats().pinned_entries, 0);
+        assert_eq!(cache.stats().unpin_operations, 2);
+    }
+
     #[test]
     fn storage_config_buffer_manager_ops() {
         let mut mgr = BufferManager::with_config(4096, 0.75, 1024);