Release sharded cache handles from iterators

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() {

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);