perf(core.utils): track FileStream cache size incrementally

Avoid the per-write directory walk in `_enforce_size_cap` by maintaining
a running byte total. The tracker is seeded once from `_compute_total_size`
at open time, updated on `__setitem__` (net delta from the old entry's
size), `__delitem__` (subtract the unlinked file's size), and `clear`
(re-derive from the post-clear state). When a write doesn't push the
running total above `max_size_bytes`, eviction stays a no-op -- writes
become O(1) instead of O(n) in the cache size.

Cross-process drift (other writers/deleters working on the same root)
self-corrects: any time `_enforce_size_cap` actually runs its scan, it
reseeds `_tracked_size_bytes` from the observed disk total, so
overestimates trigger one extra scan and then settle.

Skipped entirely when `max_size_bytes is None` (no cap, no need for a
tracker). Mutations are guarded by `_size_lock` so multi-threaded
writers in the same process don't interleave the read-modify-write on
the int.

Also switch `_iter_entry_paths` from `Path.iterdir` to `os.scandir`:
the dirent type cache lets `is_dir`/`is_file` answer without a
separate `stat` syscall on filesystems that report it (ext4, NTFS,
...). Behaviourally equivalent (the cache layout never creates symlinks,
and we pass `follow_symlinks=False` to match `pathlib`'s previous
no-symlink-confusion default for our paths).

Tests cover three guarantees: (1) writes that stay under the cap don't
call `_enforce_size_cap`, (2) writes that cross the cap do call it,
and (3) external deletion behind the cache's back is reconciled by the
next eviction pass.

Author: cpcloud

cuda_core/cuda/core/utils/_program_cache/_file_stream.py

@@ -17,6 +17,7 @@
 import hashlib
 import os
 import tempfile
+import threading
 import time
 from pathlib import Path
 from typing import Iterable
@@ -401,6 +402,19 @@ def __init__(
         # crashed writers. Age-based so concurrent in-flight writes from
         # other processes are preserved.
         self._sweep_stale_tmp_files()
+        # Incremental size tracker. Without it every ``__setitem__`` would
+        # walk ``entries/`` + ``tmp/`` to compute the total -- O(n) per
+        # write. With it: writes update the tracker by the net delta in O(1)
+        # and only walk on eviction (which already needs the scan to sort
+        # entries by atime). The tracker is seeded by one full scan at open
+        # time and refreshed on every eviction pass; cross-process drift
+        # (other writers/deleters) self-corrects the next time eviction
+        # fires. The lock guards mutations so multi-threaded writers in
+        # the same process don't interleave the read-modify-write on the
+        # int. Skipped entirely when ``max_size_bytes is None`` -- without
+        # a cap the tracker is dead weight.
+        self._size_lock = threading.Lock()
+        self._tracked_size_bytes = self._compute_total_size() if max_size_bytes is not None else 0
 
     # -- key-to-path helpers -------------------------------------------------
 
@@ -445,6 +459,18 @@ def __setitem__(self, key: object, value: bytes | bytearray | memoryview | Objec
         # FileNotFoundError even though we could trivially recover.
         self._tmp.mkdir(parents=True, exist_ok=True)
 
+        # Stat the existing entry (if any) BEFORE the replace so we can
+        # update the tracker by the net delta. A racing writer that lands
+        # an ``os.replace`` between this stat and our own makes ``old_size``
+        # slightly off; the next ``_enforce_size_cap`` reconciles by
+        # re-scanning. Skipped when ``max_size_bytes is None`` (no tracker).
+        old_size = 0
+        if self._max_size_bytes is not None:
+            try:
+                old_size = target.stat().st_size
+            except FileNotFoundError:
+                old_size = 0
+
         fd, tmp_name = tempfile.mkstemp(prefix="entry-", dir=self._tmp)
         tmp_path = Path(tmp_name)
         try:
@@ -465,14 +491,39 @@ def __setitem__(self, key: object, value: bytes | bytearray | memoryview | Objec
             with contextlib.suppress(FileNotFoundError):
                 tmp_path.unlink()
             raise
-        self._enforce_size_cap()
+
+        if self._max_size_bytes is None:
+            return
+
+        # O(1) tracker update. Only run the scan-heavy ``_enforce_size_cap``
+        # when this write actually pushes the running total above the cap.
+        new_size = len(data)
+        with self._size_lock:
+            self._tracked_size_bytes += new_size - old_size
+            over_cap = self._tracked_size_bytes > self._max_size_bytes
+        if over_cap:
+            self._enforce_size_cap()
 
     def __delitem__(self, key: object) -> None:
         path = self._path_for_key(key)
+        # Stat before unlink so we can decrement the tracker by the actual
+        # on-disk size. Best-effort: if the file vanishes between stat and
+        # unlink (concurrent eviction), we treat the delete as a miss --
+        # matching the behaviour callers expect (KeyError) and leaving the
+        # tracker untouched (the racing eviction already accounted for it).
+        size = 0
+        if self._max_size_bytes is not None:
+            try:
+                size = path.stat().st_size
+            except FileNotFoundError:
+                raise KeyError(key) from None
         try:
             _unlink_with_sharing_retry(path)
         except FileNotFoundError:
             raise KeyError(key) from None
+        if self._max_size_bytes is not None:
+            with self._size_lock:
+                self._tracked_size_bytes -= size
 
     def __len__(self) -> int:
         """Return the number of files currently in ``entries/``.
@@ -513,18 +564,71 @@ def clear(self) -> None:
                 if sub.is_dir():
                     with contextlib.suppress(OSError):
                         sub.rmdir()
+        # The directory is now (almost) empty -- but a concurrent writer may
+        # have landed a fresh entry between the snapshot and the unlink, and
+        # young temp files were intentionally preserved. Re-derive the
+        # tracker from the post-clear state instead of zeroing blindly.
+        if self._max_size_bytes is not None:
+            actual = self._compute_total_size()
+            with self._size_lock:
+                self._tracked_size_bytes = actual
 
     # -- internals -----------------------------------------------------------
 
     def _iter_entry_paths(self) -> Iterable[Path]:
-        if not self._entries.exists():
+        # ``os.scandir`` returns ``DirEntry`` objects whose ``is_dir`` /
+        # ``is_file`` methods consult the cached dirent type from the
+        # ``readdir`` result on filesystems that report it (ext4, NTFS, ...),
+        # avoiding a per-entry ``stat`` syscall. ``Path.iterdir`` also wraps
+        # ``scandir`` but discards the cached type, forcing a separate
+        # ``stat`` for every ``Path.is_dir`` / ``Path.is_file`` -- a real
+        # cost on caches with thousands of entries. ``follow_symlinks=False``
+        # matches the prior behaviour: ``pathlib`` defaults to following,
+        # but our cache layout never creates symlinks, so the choice is
+        # only visible if a user manually points one inside ``entries/``;
+        # treating it as the file/dir it points to is the surprising
+        # outcome, so we don't.
+        try:
+            outer = os.scandir(self._entries)
+        except FileNotFoundError:
             return
-        for sub in self._entries.iterdir():
-            if not sub.is_dir():
+        with outer:
+            for sub in outer:
+                if not sub.is_dir(follow_symlinks=False):
+                    continue
+                try:
+                    inner = os.scandir(sub.path)
+                except FileNotFoundError:
+                    continue
+                with inner:
+                    for entry in inner:
+                        if entry.is_file(follow_symlinks=False):
+                            yield Path(entry.path)
+
+    def _compute_total_size(self) -> int:
+        """Walk ``entries/`` + ``tmp/`` and return the on-disk byte total.
+
+        Used to seed the tracker at open time and to refresh it after every
+        eviction pass. Best-effort: files that vanish under us during the
+        walk (concurrent eviction by this or another process) are skipped.
+        Tracked total may briefly differ from this scan's result under
+        cross-process contention; the next eviction will reconcile.
+        """
+        total = 0
+        for path in self._iter_entry_paths():
+            try:
+                total += path.stat().st_size
+            except FileNotFoundError:
                 continue
-            for entry in sub.iterdir():
-                if entry.is_file():
-                    yield entry
+        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
+        return total
 
     def _sweep_stale_tmp_files(self) -> None:
         """Remove temp files left behind by crashed writers.
@@ -587,11 +691,18 @@ def _enforce_size_cap(self) -> None:
                 except FileNotFoundError:
                     continue
         if total <= self._max_size_bytes:
+            # Re-seed the tracker from the scan: catches drift from
+            # cross-process writers/deleters that the per-write delta
+            # accounting wouldn't have observed. Reaching here means the
+            # tracker was over-cap but the disk truth is under-cap, so
+            # this assignment is the cheapest reconciliation point we get.
+            with self._size_lock:
+                self._tracked_size_bytes = total
             return
         entries.sort(key=lambda e: e[0])  # oldest atime first
         for _atime, size, path, st_before in entries:
             if total <= self._max_size_bytes:
-                return
+                break
             # _prune_if_stat_unchanged refuses if a writer replaced the file
             # between snapshot and now, so eviction can't silently delete a
             # freshly-committed entry from another process.
@@ -620,3 +731,8 @@ def _enforce_size_cap(self) -> None:
             except PermissionError as exc:
                 if not _is_windows_sharing_violation(exc):
                     raise
+        # Reconcile: after the eviction pass, ``total`` reflects what we
+        # believe the disk now holds. Re-seed the tracker so the next write
+        # accumulates from a fresh baseline.
+        with self._size_lock:
+            self._tracked_size_bytes = total

cuda_core/tests/test_program_cache.py

@@ -1909,6 +1909,80 @@ def test_filestream_cache_unbounded_by_default(tmp_path):
         assert len(cache) == 20
 
 
+def test_filestream_cache_writes_skip_scan_when_under_cap(tmp_path, monkeypatch):
+    """An incremental size tracker keeps writes O(1) while under the cap.
+    Counter-test against a regression that re-introduces a per-write disk
+    walk: if every write triggers the eviction scan, the running total has
+    to be recomputed for entries we never modified, and a cache with
+    thousands of files gets quadratic. Writes that stay under
+    ``max_size_bytes`` must not call ``_enforce_size_cap`` at all."""
+    from cuda.core.utils import FileStreamProgramCache
+
+    with FileStreamProgramCache(tmp_path / "fc", max_size_bytes=1_000_000) as cache:
+        # Seed one entry so the tracker has a non-zero starting point.
+        cache[b"k0"] = b"x" * 100
+
+        calls: list[int] = []
+        original = cache._enforce_size_cap
+        monkeypatch.setattr(
+            cache,
+            "_enforce_size_cap",
+            lambda: (calls.append(1), original())[1],
+        )
+
+        # 100 small writes well under the 1MB cap -- none should trigger
+        # the scan-heavy enforcement path.
+        for i in range(1, 101):
+            cache[f"k{i}".encode()] = b"y" * 100
+        assert calls == [], f"_enforce_size_cap called {len(calls)} times under cap"
+
+
+def test_filestream_cache_writes_trigger_scan_when_over_cap(tmp_path, monkeypatch):
+    """Counterpart to ``test_filestream_cache_writes_skip_scan_when_under_cap``:
+    once the tracker reports we crossed the cap, the eviction scan must
+    fire so the cache actually enforces its bound."""
+    from cuda.core.utils import FileStreamProgramCache
+
+    with FileStreamProgramCache(tmp_path / "fc", max_size_bytes=300) as cache:
+        calls: list[int] = []
+        original = cache._enforce_size_cap
+        monkeypatch.setattr(
+            cache,
+            "_enforce_size_cap",
+            lambda: (calls.append(1), original())[1],
+        )
+
+        # 200 + 200 > 300 -> second write must trigger enforcement.
+        cache[b"a"] = b"a" * 200
+        cache[b"b"] = b"b" * 200
+        assert len(calls) >= 1
+
+
+def test_filestream_cache_tracker_reconciles_after_external_drift(tmp_path):
+    """If another writer (or this cache reopened) deletes entries on disk
+    out from under us, the running tracker overestimates the on-disk total.
+    The next eviction pass walks the directory and re-seeds the tracker
+    from reality, so the overestimate self-corrects rather than growing
+    without bound and triggering bogus evictions every write thereafter."""
+    from cuda.core.utils import FileStreamProgramCache
+
+    with FileStreamProgramCache(tmp_path / "fc", max_size_bytes=1000) as cache:
+        cache[b"a"] = b"A" * 400
+        cache[b"b"] = b"B" * 400
+        assert cache._tracked_size_bytes == 800
+        # Simulate an external deleter (another process, manual rm,
+        # filesystem-level pruning script): unlink an entry behind the
+        # cache's back so the tracker is now stale-high.
+        path_a = cache._path_for_key(b"a")
+        path_a.unlink()
+        # Tracker still believes 800 bytes are on disk; trigger an
+        # eviction by writing one more entry that pushes the tracker
+        # above the cap. The scan should observe ~400 bytes (only 'b'
+        # remains plus the new write) and reset the tracker to that.
+        cache[b"c"] = b"C" * 700  # tracker becomes 800 + 700 = 1500 > cap
+        assert cache._tracked_size_bytes <= 1100  # actual on-disk is 'b' + 'c' or just 'c'
+
+
 def test_make_program_cache_key_changes_with_key_schema_version(monkeypatch):
     """Bumping ``_KEY_SCHEMA_VERSION`` produces a different cache key for
     the same logical inputs. That's what makes a schema bump invalidate