fixup! test(core.utils): exercise real _enforce_size_cap via generator-cleanup race injection

Author: cpcloud

cuda_core/tests/test_program_cache.py

@@ -1550,84 +1550,57 @@ def test_filestream_cache_clear_does_not_break_concurrent_writer(tmp_path):
 
 
 def test_filestream_cache_size_cap_does_not_unlink_replaced_file(tmp_path):
-    """``_enforce_size_cap``'s eviction loop must compare the snapshot stat
+    """The PRODUCTION ``_enforce_size_cap`` must compare the snapshot stat
     to the current stat before unlinking; if the file was replaced under
-    us (a concurrent writer's os.replace), the unlink is skipped.
-
-    Race injection: subclass the cache and override ``_iter_entry_paths`` to
-    yield the entry paths AND, on the second invocation (the eviction
-    loop's re-stat), have the wrapper invalidate a target by os.replace-ing
-    it before the eviction loop calls ``path.stat()``.
+    us (a concurrent writer's ``os.replace``), the unlink is skipped.
+
+    Race injection without reimplementing the method: subclass the cache
+    and override only ``_iter_entry_paths`` so that the cleanup code
+    *after* the generator's last yield runs an ``os.replace`` on path_a.
+    Python's for-loop calls ``next()`` until ``StopIteration``; the
+    generator code after its last yield runs at that ``StopIteration``,
+    which is exactly between ``_enforce_size_cap``'s scan loop and its
+    eviction loop. Eviction's per-entry re-stat then sees a different
+    stat for path_a and the production code's stat-guard must skip it.
     """
     import os as _os
 
     from cuda.core.utils import FileStreamProgramCache
 
-    cap = 4000
+    # Cap fits two entries (each ~2123 bytes on disk for a 2000-byte
+    # payload, including pickle overhead) but not three.
+    cap = 5000
     root = tmp_path / "fc"
     with FileStreamProgramCache(root, max_size_bytes=cap) as cache:
         cache[b"a"] = _fake_object_code(b"A" * 2000, name="a")
         time.sleep(0.02)
         cache[b"b"] = _fake_object_code(b"B" * 2000, name="b")
         path_a = cache._path_for_key(b"a")
+    assert path_a.exists(), "cap too small -- 'a' was evicted before the test ran"
 
-    # Reopen and simulate the race: replace path_a's contents on disk
-    # AFTER _enforce_size_cap has already snapshotted its stat. We do that
-    # by subclassing the cache and intercepting the `for ... in entries`
-    # eviction loop -- mutating the file between the snapshot list build
-    # and the per-entry unlink check.
     class _RaceCache(FileStreamProgramCache):
-        def _enforce_size_cap(self):
-            # Mirror the parent's scan + race injection + unlink path.
-            if self._max_size_bytes is None:
-                return
-            self._sweep_stale_tmp_files()
-            entries = []
-            total = 0
-            for path in self._iter_entry_paths():
-                try:
-                    st = path.stat()
-                except FileNotFoundError:
-                    continue
-                entries.append((st.st_mtime, st.st_size, path, st))
-                total += st.st_size
-            if total <= self._max_size_bytes:
-                return
-            entries.sort(key=lambda e: e[0])
-            # Race injection: replace path_a's bytes via os.replace BEFORE
-            # the eviction loop re-stats it.
-            tmp = path_a.parent / "_inflight"
-            tmp.write_bytes(b"\x80\x05fresh-by-other-writer-" * 32)
-            _os.replace(tmp, path_a)
-            # Now run the parent's eviction loop using the SAME snapshot.
-            for _mtime, size, path, st_before in entries:
-                if total <= self._max_size_bytes:
-                    return
-                try:
-                    stat_now = path.stat()
-                except FileNotFoundError:
-                    total -= size
-                    continue
-                if (stat_now.st_ino, stat_now.st_size, stat_now.st_mtime_ns) != (
-                    st_before.st_ino,
-                    st_before.st_size,
-                    st_before.st_mtime_ns,
-                ):
-                    total += stat_now.st_size - size
-                    continue
-                import contextlib
-
-                with contextlib.suppress(FileNotFoundError):
-                    path.unlink()
-                    total -= size
+        race_armed = True
+
+        def _iter_entry_paths(self):
+            yield from super()._iter_entry_paths()
+            # Generator cleanup runs at StopIteration, between
+            # _enforce_size_cap's scan and its eviction loop. Fire the race
+            # here exactly once.
+            if _RaceCache.race_armed and path_a.exists():
+                _RaceCache.race_armed = False
+                tmp = path_a.parent / "_inflight"
+                tmp.write_bytes(b"\x80\x05fresh-by-other-writer-" * 32)
+                _os.replace(tmp, path_a)
 
     with _RaceCache(root, max_size_bytes=cap) as cache:
-        # Trigger eviction by adding 'c'; eviction will scan, then
-        # _RaceCache replaces path_a, then eviction re-stats and skips it.
+        # Trigger eviction by adding 'c'; eviction's scan exhausts our
+        # racing generator, the cleanup fires, then the eviction loop's
+        # re-stat sees the new stat and the production stat-guard MUST
+        # refuse to unlink path_a.
         time.sleep(0.02)
         cache[b"c"] = _fake_object_code(b"C" * 2000, name="c")
 
-    # The race-injected fresh file must survive (NOT unlinked by eviction).
+    # The race-injected fresh file must survive: production stat-guard worked.
     assert path_a.exists(), "stat guard failed -- evicted a concurrently-replaced file"
     assert path_a.read_bytes().startswith(b"\x80\x05fresh-by-other-writer-")