Fix dynamic NIXL store secondary-lookup byte overcount race

cursoragent · cursoragent · commit c49599c49c41 · 2026-04-13T20:22:25.000Z
diff --git a/lmcache/v1/distributed/l2_adapters/nixl_store_dynamic_l2_adapter.py b/lmcache/v1/distributed/l2_adapters/nixl_store_dynamic_l2_adapter.py
@@ -571,9 +571,12 @@ async def _execute_store_in_the_loop(
                     pin_count=1,
                 )
                 with self._lock:
-                    self._memory_objects[key] = store_obj
-                    self._total_bytes += store_obj.size
-                    store_obj.decrease_pin_count()
+                    # Re-check in case lookup/recovery populated this key
+                    # while we awaited dynamic_store_file.
+                    if key not in self._memory_objects:
+                        self._memory_objects[key] = store_obj
+                        self._total_bytes += store_obj.size
+                        store_obj.decrease_pin_count()
                 stored_keys.append(key)
 
         except Exception:
diff --git a/tests/v1/distributed/test_nixl_store_dynamic_l2_adapter.py b/tests/v1/distributed/test_nixl_store_dynamic_l2_adapter.py
@@ -7,6 +7,7 @@
 """
 
 # Standard
+import asyncio
 import os
 import select
 import shutil
@@ -561,6 +562,47 @@ def test_get_usage_decreases_after_delete(self, adapter):
 
         assert usage_after < usage_before
 
+    def test_store_secondary_lookup_race_does_not_overcount_usage(
+        self, adapter_with_persist
+    ):
+        """Store+recover race should not leave stale bytes after delete."""
+        adpt, buf, _, _, _ = adapter_with_persist
+        key = create_object_key(1001)
+        obj = create_memory_obj(buf, page_index=0)
+        lookup_tasks: list[int] = []
+
+        async def _dynamic_store_with_lookup(mem_indices, file_path, page_size):
+            file_size = len(mem_indices) * page_size
+            with open(file_path, "wb") as f:
+                f.truncate(file_size)
+            lookup_tasks.append(adpt.submit_lookup_and_lock_task([key]))
+            await asyncio.sleep(0)
+
+        try:
+            adpt.nixl_agent.dynamic_store_file = _dynamic_store_with_lookup
+            adpt.submit_store_task([key], [obj])
+            wait_for_event_fd(adpt.get_store_event_fd())
+            wait_for_event_fd(adpt.get_lookup_and_lock_event_fd())
+            adpt.pop_completed_store_tasks()
+
+            bitmap = adpt.query_lookup_and_lock_result(lookup_tasks[0])
+            assert bitmap is not None
+            assert bitmap.test(0)
+
+            adpt.submit_unlock([key])
+            sync_task = adpt.submit_lookup_and_lock_task([])
+            wait_for_event_fd(adpt.get_lookup_and_lock_event_fd())
+            adpt.query_lookup_and_lock_result(sync_task)
+
+            usage_before_delete, _ = adpt.get_usage()
+            adpt.delete([key])
+            usage_after_delete, _ = adpt.get_usage()
+
+            assert usage_before_delete > 0.0
+            assert usage_after_delete == 0.0
+        finally:
+            adpt.close()
+
     def test_store_rejected_when_capacity_exceeded(self):
         """Store should stop when max capacity is reached."""
         tmp_dir = tempfile.mkdtemp(prefix="nixl_dyn_cap_test_")
