mooncake-wheel/mooncake/structured_object_store.py

@@ -1263,7 +1263,6 @@ def _is_duplicate_buffer_registration(status: Any) -> bool:
 
 
 def _cleanup_keys(store: BundleStore, keys: Sequence[str], strict: bool) -> None:
-    errors = []
     pending_keys = list(dict.fromkeys(keys))
     batch_remove = getattr(store, "batch_remove", None)
     if callable(batch_remove) and pending_keys:
@@ -1283,7 +1282,6 @@ def _cleanup_keys(store: BundleStore, keys: Sequence[str], strict: bool) -> None
             ]
             if not failed_results:
                 return
-            errors.extend(failed_results)
             pending_keys = [key for key, _status in failed_results]
         except Exception:
             if strict:
@@ -1301,9 +1299,7 @@ def _cleanup_keys(store: BundleStore, keys: Sequence[str], strict: bool) -> None
             continue
         if status not in (None, 0, MISSING_OBJECT_ERROR):
             retry_errors.append((key, status))
-    if retry_errors:
-        errors = retry_errors
-    if errors and strict:
+    if retry_errors and strict:
         raise RuntimeError(
-            f"failed to remove {len(errors)} Mooncake keys: {errors[:3]}"
+            f"failed to remove {len(retry_errors)} Mooncake keys: {retry_errors[:3]}"
         )

mooncake-wheel/tests/test_structured_object_store.py

@@ -236,6 +236,25 @@ class ForceTrackingStore(InMemoryStore):
         return super().batch_remove(keys, force)
 
 
+class TransientBatchRemoveStore(InMemoryStore):
+    """batch_remove reports a transient failure for one key without deleting it.
+
+    The inherited per-key remove() retry then succeeds and deletes the key, so
+    cleanup should end with no outstanding error.
+    """
+
+    def batch_remove(self, keys: list[str], force: bool = False) -> list[int]:
+        self.batch_remove_calls += 1
+        results: list[int] = []
+        for index, key in enumerate(keys):
+            if index == 0:
+                results.append(-1)  # transient failure, key left in place
+            else:
+                self.remove(key, force)
+                results.append(0)
+        return results
+
+
 class StrictRegisterStore(InMemoryStore):
     def register_buffer(self, buffer_ptr: int, size: int) -> int:
         if buffer_ptr in self.registered:
@@ -580,6 +599,16 @@ def test_bundle_remove_uses_force_batch_remove_when_available() -> None:
     assert store.objects == {}
 
 
+def test_bundle_remove_recovers_after_transient_batch_failure() -> None:
+    store, transfer = make_transfer(TransientBatchRemoveStore())
+
+    ref = transfer.put_bundle(b"meta", {"a": b"x", "b": b"y", "c": b"z"})
+    transfer.remove_bundle(ref)
+
+    assert store.objects == {}
+    assert store.batch_remove_calls == 1
+
+
 def test_bundle_concurrent_put_and_read_spec_full_read() -> None:
     store, transfer = make_transfer(GetOnlyStore())
     payload = bytes(range(128))