fixup! feat(core.utils): validate name_expressions; sweep stale tmp files

Author: cpcloud

cuda_core/cuda/core/utils/_program_cache.py

@@ -597,10 +597,15 @@ def make_program_cache_key(
     # (_program.pyx:759). Returning a cached ObjectCode whose mapping-key
     # type differs from what the caller's later ``get_kernel`` passes would
     # silently miss -- so treat ``"foo"`` and ``b"foo"`` as distinct here.
+    # Reject anything other than str/bytes/bytearray up front; Program.compile
+    # would fail at compile time anyway, and persisting a key for an invalid
+    # input is just a foot-gun.
     def _tag_name(n):
         if isinstance(n, (bytes, bytearray)):
             return b"b:" + bytes(n)
-        return b"s:" + str(n).encode("utf-8")
+        if isinstance(n, str):
+            return b"s:" + n.encode("utf-8")
+        raise TypeError(f"name_expressions elements must be str, bytes, or bytearray; got {type(n).__name__}")
 
     names = tuple(sorted(_tag_name(n) for n in name_expressions))
 
@@ -988,6 +993,11 @@ def _enforce_size_cap(self) -> None:
 _ENTRIES_SUBDIR = "entries"
 _TMP_SUBDIR = "tmp"
 _SCHEMA_FILE = "SCHEMA_VERSION"
+# Temp files older than this are assumed to belong to a crashed writer and
+# are eligible for cleanup. Picked large enough that no real ``os.replace``
+# write should still be in flight (writes are bounded by mkstemp + write +
+# fsync + replace, all fast on healthy disks).
+_TMP_STALE_AGE_SECONDS = 3600
 
 
 _SHARING_VIOLATION_WINERRORS = (32, 33)  # ERROR_SHARING_VIOLATION, ERROR_LOCK_VIOLATION
@@ -1114,6 +1124,10 @@ def __init__(
                     with contextlib.suppress(FileNotFoundError):
                         entry.unlink()
                 self._schema_path.write_text(expected)
+        # Opportunistic startup sweep of orphaned temp files left by any
+        # crashed writers. Age-based so concurrent in-flight writes from
+        # other processes are preserved.
+        self._sweep_stale_tmp_files()
 
     # -- key-to-path helpers -------------------------------------------------
 
@@ -1219,6 +1233,15 @@ def clear(self) -> None:
         for path in list(self._iter_entry_paths()):
             with contextlib.suppress(FileNotFoundError):
                 path.unlink()
+        # The user explicitly asked to wipe this cache, so also drop every
+        # temp file we can see (whether stale or in flight from this process).
+        # Other processes' in-flight writes will still complete to ``entries``
+        # via ``os.replace``, but their staging files are intentionally gone.
+        if self._tmp.exists():
+            for tmp in list(self._tmp.iterdir()):
+                if tmp.is_file():
+                    with contextlib.suppress(FileNotFoundError):
+                        tmp.unlink()
         # Remove empty subdirs (best-effort; concurrent writers may re-create).
         if self._entries.exists():
             for sub in sorted(self._entries.iterdir(), reverse=True):
@@ -1238,18 +1261,51 @@ def _iter_entry_paths(self) -> Iterable[Path]:
                 if entry.is_file():
                     yield entry
 
+    def _sweep_stale_tmp_files(self) -> None:
+        """Remove temp files left behind by crashed writers.
+
+        Age threshold is conservative (``_TMP_STALE_AGE_SECONDS``) so an
+        in-flight write from another process is not interrupted. Best
+        effort: a missing file or a permission failure is ignored.
+        """
+        if not self._tmp.exists():
+            return
+        cutoff = time.time() - _TMP_STALE_AGE_SECONDS
+        for tmp in self._tmp.iterdir():
+            if not tmp.is_file():
+                continue
+            try:
+                if tmp.stat().st_mtime < cutoff:
+                    tmp.unlink()
+            except (FileNotFoundError, PermissionError):
+                continue
+
     def _enforce_size_cap(self) -> None:
         if self._max_size_bytes is None:
             return
+        # Sweep stale temp files first so a long-dead writer's leftovers
+        # don't drag the apparent size up and force needless eviction.
+        self._sweep_stale_tmp_files()
         entries = []
         total = 0
+        # Count both committed entries AND surviving temp files: temp files
+        # occupy disk too, even if they're young. Without this the soft cap
+        # silently undercounts in-flight writes.
         for path in self._iter_entry_paths():
             try:
                 st = path.stat()
             except FileNotFoundError:
                 continue
             entries.append((st.st_mtime, st.st_size, path))
             total += st.st_size
+        if self._tmp.exists():
+            for tmp in self._tmp.iterdir():
+                if not tmp.is_file():
+                    continue
+                try:
+                    total += tmp.stat().st_size
+                except FileNotFoundError:
+                    continue
         if total <= self._max_size_bytes:
             return
         entries.sort()  # oldest mtime first

cuda_core/tests/test_program_cache.py

@@ -232,6 +232,14 @@ def test_make_program_cache_key_name_expressions_order_insensitive():
     assert _make_key(name_expressions=("f", "g")) == _make_key(name_expressions=("g", "f"))
 
 
+@pytest.mark.parametrize("bad", [123, 1.5, object(), None])
+def test_make_program_cache_key_rejects_invalid_name_expressions_element(bad):
+    """Program.compile only forwards str/bytes name_expressions to NVRTC;
+    persisting a key for an invalid input is just a foot-gun. Reject up front."""
+    with pytest.raises(TypeError, match="name_expressions"):
+        _make_key(name_expressions=("ok", bad))
+
+
 def test_make_program_cache_key_name_expressions_str_bytes_distinct():
     """``Program.compile`` records the *original* Python object as the key in
     ``ObjectCode.symbol_mapping``. Returning a cached ObjectCode whose
@@ -1376,6 +1384,75 @@ def test_filestream_cache_rejects_negative_size_cap(tmp_path):
         FileStreamProgramCache(tmp_path / "fc", max_size_bytes=-1)
 
 
+def test_filestream_cache_sweeps_stale_tmp_files_on_open(tmp_path):
+    """A crashed writer can leave files in ``tmp/``; the next ``open`` must
+    sweep ones older than the staleness threshold so disk usage doesn't
+    grow without bound."""
+    import os as _os
+
+    from cuda.core.utils import FileStreamProgramCache, _program_cache
+
+    root = tmp_path / "fc"
+    # Create the cache directory layout, then plant two temp files: one
+    # young (must be preserved as it could be an in-flight write) and one
+    # ancient (must be swept).
+    with FileStreamProgramCache(root):
+        pass
+    young = root / "tmp" / "entry-young"
+    young.write_bytes(b"in-flight")
+    ancient = root / "tmp" / "entry-ancient"
+    ancient.write_bytes(b"crashed-writer-leftover")
+    ancient_mtime = time.time() - _program_cache._TMP_STALE_AGE_SECONDS - 60
+    _os.utime(ancient, (ancient_mtime, ancient_mtime))
+
+    with FileStreamProgramCache(root):
+        # Reopen triggers _sweep_stale_tmp_files.
+        assert young.exists(), "young temp file must not be swept"
+        assert not ancient.exists(), "ancient temp file should have been swept"
+
+
+def test_filestream_cache_clear_drops_all_tmp_files(tmp_path):
+    """clear() is an explicit user wipe, so it removes every temp file too
+    -- including young ones (other processes' in-flight writes will still
+    complete to entries/, just without their staging file)."""
+    from cuda.core.utils import FileStreamProgramCache
+
+    root = tmp_path / "fc"
+    with FileStreamProgramCache(root) as cache:
+        cache[b"k"] = _fake_object_code(b"v")
+    young_tmp = root / "tmp" / "entry-young"
+    young_tmp.write_bytes(b"in-flight")
+
+    with FileStreamProgramCache(root) as cache:
+        cache.clear()
+    assert not young_tmp.exists()
+
+
+def test_filestream_cache_size_cap_counts_tmp_files(tmp_path):
+    """Surviving temp files occupy disk too; the soft cap must include
+    them, otherwise an attacker (or a flurry of crashed writers) could
+    inflate disk usage well past max_size_bytes."""
+    from cuda.core.utils import FileStreamProgramCache
+
+    cap = 4000
+    root = tmp_path / "fc"
+    with FileStreamProgramCache(root, max_size_bytes=cap) as cache:
+        cache[b"a"] = _fake_object_code(b"A" * 1500, name="a")
+        time.sleep(0.02)
+        cache[b"b"] = _fake_object_code(b"B" * 1500, name="b")
+    # Plant a young temp file that pushes total over the cap.
+    young_tmp = root / "tmp" / "entry-leftover"
+    young_tmp.write_bytes(b"X" * 2500)
+
+    with FileStreamProgramCache(root, max_size_bytes=cap) as cache:
+        # New write triggers _enforce_size_cap; 'a' must be evicted because
+        # the temp file's bytes count toward the cap now.
+        time.sleep(0.02)
+        cache[b"c"] = _fake_object_code(b"C" * 200, name="c")
+        assert b"a" not in cache
+        assert b"c" in cache
+
+
 def test_filestream_cache_handles_long_keys(tmp_path):
     """Arbitrary-length keys must not overflow per-component filename limits.
     The filename is a fixed-length hash; the original key is verified from