diff --git a/src/openultrasast/push/cache.py b/src/openultrasast/push/cache.py index ef8c306..9e321d6 100644 --- a/src/openultrasast/push/cache.py +++ b/src/openultrasast/push/cache.py @@ -2,6 +2,7 @@ from __future__ import annotations +import contextlib import fcntl import hashlib import json @@ -45,6 +46,19 @@ class ArtifactCache: self.last_reason = "" self.hits = 0 self.misses = 0 + # The last full scan of what the cache holds, keyed by the root directory's modification stamp. Every + # publication used to list and stat the whole cache to total its size, so a push that stores one + # answer per question paid O(entries) per answer: on a WordPress plugin, 12,115 entries in, each new + # answer cost ~50,000 filesystem calls and the evidence pass never finished a 7,200 s case. Any change + # another process makes -- an entry added, evicted or removed -- moves the stamp and forces a rescan, + # so the size limit stays exact under concurrent pushes; our own changes update the tally in place. + self._usage: tuple[int, list[tuple[float, Path, int]], int] | None = None + self._room: tuple[list[tuple[float, Path, int]], int] = ([], 0) + # Every lock stripe exists from the start, so creating one on first use never moves the root's stamp + # and invalidates the remembered scan. A stripe that cannot be created is left to `_lock`, as before. + for name in [".publication-lock", *(f".lock-{stripe:02d}" for stripe in range(64))]: + with contextlib.suppress(OSError): + os.close(os.open(self.root / name, os.O_CREAT | os.O_RDWR | os.O_NOFOLLOW | os.O_NONBLOCK, 0o600)) @contextmanager def _lock(self, name: str, budget: ExecutionBudget, *, shared: bool = False, wait: bool = True) -> Iterator[bool]: @@ -146,24 +160,22 @@ class ArtifactCache: _remove_owned_tree(root, budget.deadline_monotonic) def _make_room(self, incoming: int, key: str, budget: ExecutionBudget) -> bool: - entries = [] - used = 0 - for path in self.root.iterdir(): - if time.monotonic() >= budget.deadline_monotonic: + remembered = self._usage + self._usage = None # restored only by a publication that completes + if remembered is not None and remembered[0] == self.root.stat().st_mtime_ns: + entries, used = list(remembered[1]), remembered[2] + else: + scanned = self._scan(budget) + if scanned is None: return False - if path.is_symlink(): + entries, used = scanned + kept: list[tuple[float, Path, int]] = [] + fits = False + for mtime, path, size in sorted(entries): + if fits or used + incoming <= self.max_bytes: + fits = True + kept.append((mtime, path, size)) continue - if path.is_dir() and (_KEY.fullmatch(path.name) or path.name.startswith(".pending-")): - size = 0 - for child in path.iterdir(): - if time.monotonic() >= budget.deadline_monotonic: - return False - size += child.lstat().st_size - used += size - entries.append((path.stat().st_mtime, path, size)) - for _, path, size in sorted(entries): - if used + incoming <= self.max_bytes: - return True if path.name == key or path.name.startswith(".pending-"): self._remove(path, budget) used -= size @@ -172,8 +184,30 @@ class ArtifactCache: if locked: self._remove(path, budget) used -= size + else: + kept.append((mtime, path, size)) + self._room = (kept, used) return used + incoming <= self.max_bytes + def _scan(self, budget: ExecutionBudget) -> tuple[list[tuple[float, Path, int]], int] | None: + """Every entry and pending publication with its size, or ``None`` when the deadline passed.""" + entries = [] + used = 0 + for path in self.root.iterdir(): + if time.monotonic() >= budget.deadline_monotonic: + return None + if path.is_symlink(): + continue + if path.is_dir() and (_KEY.fullmatch(path.name) or path.name.startswith(".pending-")): + size = 0 + for child in path.iterdir(): + if time.monotonic() >= budget.deadline_monotonic: + return None + size += child.lstat().st_size + used += size + entries.append((path.stat().st_mtime, path, size)) + return entries, used + def publish(self, key: str, payload: bytes | Path, metadata: dict[str, Any], budget: ExecutionBudget, *, complete: bool = True) -> bool: if not complete or not _KEY.fullmatch(key): self.last_reason = "partial_or_invalid_cache_publication" @@ -224,6 +258,12 @@ class ArtifactCache: self._remove(existing, budget) os.replace(pending, existing) pending = None + kept, used = self._room + replaced = [entry for entry in kept if entry[1].name == key] + kept = [entry for entry in kept if entry[1].name != key] + used -= sum(entry[2] for entry in replaced) + kept.append((existing.stat().st_mtime, existing, size + len(manifest))) + self._usage = (self.root.stat().st_mtime_ns, kept, used + size + len(manifest)) return True except (OSError, ValueError, TimeoutError, _DeadlineExpired) as error: self.last_reason = type(error).__name__ diff --git a/tests/test_push_cache.py b/tests/test_push_cache.py index 2f33908..e1ef450 100644 --- a/tests/test_push_cache.py +++ b/tests/test_push_cache.py @@ -146,3 +146,39 @@ def test_abandoned_partial_files_count_toward_size_and_are_evicted(tmp_path): assert cache.publish(key, b"b" * 1000, {}, budget()) assert not abandoned.exists() assert sum(p.stat().st_size for p in cache.root.rglob("*") if p.is_file()) <= cache.max_bytes + + +def _real_usage(cache): + """The cache's size as a fresh scan sees it: the instrument the remembered tally must agree with.""" + return sum(child.lstat().st_size for entry in cache.root.iterdir() if entry.is_dir() for child in entry.iterdir()) + + +def test_publications_reuse_one_scan_and_keep_an_exact_tally(tmp_path, monkeypatch): + """A push stores one answer per question; rescanning the cache on each made the pass quadratic. + + On a WordPress plugin, 12,115 entries in, every new answer cost ~50,000 filesystem calls and the evidence + pass never finished its case. The tally must stay exact, not merely cheap. + """ + cache = ArtifactCache(tmp_path / "cache", max_bytes=10**7) + scans = [] + real_scan = cache._scan + monkeypatch.setattr(cache, "_scan", lambda b: scans.append(1) or real_scan(b)) + for index in range(60): + assert cache.publish(digest_value(f"answer-{index}"), b"rows" * (index + 1), {"kind": "query"}, budget()) + assert cache.publish(digest_value("answer-7"), b"replaced", {"kind": "query"}, budget()) # same key, new size + assert len(scans) == 1, f"{len(scans)} full scans for 61 publications" + assert cache._usage is not None and cache._usage[2] == _real_usage(cache) + + +def test_another_writer_forces_a_rescan_and_the_limit_holds(tmp_path): + """The tally is only trusted while nothing else touched the cache; the size limit is never approximate.""" + root = tmp_path / "cache" + first = ArtifactCache(root, max_bytes=6000) + second = ArtifactCache(root, max_bytes=6000) + for index in range(4): + assert first.publish(digest_value(f"first-{index}"), b"a" * 1000, {}, budget()) + for index in range(4): + assert second.publish(digest_value(f"second-{index}"), b"b" * 1000, {}, budget()) + assert first.publish(digest_value("first-late"), b"c" * 1000, {}, budget()) + assert _real_usage(first) <= 6000 + assert first._usage is not None and first._usage[2] == _real_usage(first)