377 lines
17 KiB
Python
377 lines
17 KiB
Python
"""Use/eviction races and byte budgets for the immutable native build cache."""
|
|
from __future__ import annotations
|
|
|
|
import gc
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
from app.simulation.native_codegen import cache_storage as storage
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
|
|
|
|
def key(number: int) -> str:
|
|
# Distinct first words make shard collisions explicit in the tests.
|
|
return f"{number:08x}" + "0" * 56
|
|
|
|
|
|
class NativeCacheStorageTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temporary = tempfile.TemporaryDirectory(prefix="native-cache-storage-")
|
|
self.addCleanup(self.temporary.cleanup)
|
|
self.folder = Path(self.temporary.name)
|
|
self.cache = self.folder / "cache"
|
|
|
|
def entry(self, number: int, size: int, age: int, kind: str = "models") -> Path:
|
|
path = self.cache / kind / key(number)
|
|
path.mkdir(parents=True)
|
|
(path / "artifact").write_bytes(b"x" * size)
|
|
os.utime(path, ns=(1_000_000_000 + age, 1_000_000_000 + age))
|
|
return path
|
|
|
|
def lease(self, number: int, **kwargs):
|
|
lease = storage.acquire_cache_lease(self.cache, "models", key(number), **kwargs)
|
|
if lease is not None:
|
|
self.addCleanup(lease.close)
|
|
return lease
|
|
|
|
def test_overlapping_shared_leases_and_nonblocking_exclusion(self):
|
|
first = self.lease(0)
|
|
second = self.lease(0, blocking=False)
|
|
self.assertIsNotNone(second)
|
|
self.assertIsNone(self.lease(0, exclusive=True, blocking=False))
|
|
first.close()
|
|
self.assertTrue(first.closed)
|
|
self.assertIsNone(self.lease(0, exclusive=True, blocking=False))
|
|
second.close()
|
|
writer = self.lease(0, exclusive=True, blocking=False)
|
|
self.assertIsNotNone(writer)
|
|
self.assertIsNone(self.lease(0, blocking=False))
|
|
writer.close()
|
|
writer.close() # Explicit close and cleanup are idempotent.
|
|
with self.assertRaises(RuntimeError):
|
|
with writer:
|
|
pass
|
|
|
|
def test_separate_process_sees_shared_and_exclusive_lock_contract(self):
|
|
self.lease(1)
|
|
script = r'''
|
|
import json, sys
|
|
from pathlib import Path
|
|
from app.simulation.native_codegen.cache_storage import acquire_cache_lease
|
|
cache, key = Path(sys.argv[1]), sys.argv[2]
|
|
shared = acquire_cache_lease(cache, 'models', key, blocking=False)
|
|
exclusive = acquire_cache_lease(cache, 'models', key, exclusive=True, blocking=False)
|
|
print(json.dumps({'shared': shared is not None, 'exclusive': exclusive is not None}))
|
|
if shared: shared.close()
|
|
if exclusive: exclusive.close()
|
|
'''
|
|
result = subprocess.run(
|
|
[sys.executable, "-c", script, str(self.cache), key(1)], cwd=ROOT,
|
|
capture_output=True, text=True, timeout=15, check=True,
|
|
)
|
|
self.assertEqual(json.loads(result.stdout), {"shared": True, "exclusive": False})
|
|
|
|
def test_blocking_process_proceeds_when_last_reader_releases(self):
|
|
lease = self.lease(2)
|
|
script = r'''
|
|
import sys
|
|
from pathlib import Path
|
|
from app.simulation.native_codegen.cache_storage import acquire_cache_lease
|
|
print('waiting', flush=True)
|
|
with acquire_cache_lease(Path(sys.argv[1]), 'models', sys.argv[2], exclusive=True):
|
|
print('acquired', flush=True)
|
|
'''
|
|
process = subprocess.Popen(
|
|
[sys.executable, "-c", script, str(self.cache), key(2)], cwd=ROOT,
|
|
stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True,
|
|
)
|
|
self.addCleanup(lambda: process.kill() if process.poll() is None else None)
|
|
self.assertEqual(process.stdout.readline().strip(), "waiting")
|
|
self.assertIsNone(process.poll())
|
|
lease.close()
|
|
stdout, stderr = process.communicate(timeout=15)
|
|
self.assertEqual(process.returncode, 0, stderr)
|
|
self.assertEqual(stdout.strip(), "acquired")
|
|
|
|
def test_finalizer_releases_abandoned_lease(self):
|
|
lease = storage.acquire_cache_lease(self.cache, "models", key(3))
|
|
self.assertIsNone(self.lease(3, exclusive=True, blocking=False))
|
|
del lease
|
|
gc.collect()
|
|
self.assertIsNotNone(self.lease(3, exclusive=True, blocking=False))
|
|
|
|
def test_lru_removes_oldest_until_separate_budgets_are_met(self):
|
|
oldest = self.entry(0, 60, 1)
|
|
middle = self.entry(1, 60, 2)
|
|
newest = self.entry(2, 60, 3)
|
|
object_old = self.entry(3, 40, 1, "objects")
|
|
object_new = self.entry(4, 40, 2, "objects")
|
|
report = storage.prune_cache(self.cache, model_limit_bytes=120, object_limit_bytes=40)
|
|
self.assertFalse(oldest.exists())
|
|
self.assertTrue(middle.exists() and newest.exists())
|
|
self.assertFalse(object_old.exists())
|
|
self.assertTrue(object_new.exists())
|
|
self.assertEqual(report["models"]["afterBytes"], 120)
|
|
self.assertEqual(report["objects"]["afterBytes"], 40)
|
|
self.assertEqual(report["models"]["removedEntries"], 1)
|
|
self.assertEqual(report["models"]["errors"], [])
|
|
|
|
def test_held_entry_survives_and_recent_touch_changes_lru_order(self):
|
|
first = self.entry(0, 60, 1)
|
|
second = self.entry(1, 60, 2)
|
|
third = self.entry(2, 60, 3)
|
|
held = self.lease(0)
|
|
with self.lease(1):
|
|
storage.touch_cache_entry(self.cache, "models", key(1))
|
|
report = storage.prune_cache(self.cache, model_limit_bytes=60)
|
|
self.assertTrue(first.exists())
|
|
self.assertFalse(second.exists() or third.exists())
|
|
self.assertEqual(report["models"]["skippedInUse"], 1)
|
|
held.close()
|
|
|
|
# A separate round verifies touch order without an active-entry override.
|
|
old = self.entry(3, 60, 3)
|
|
with self.lease(0):
|
|
storage.touch_cache_entry(self.cache, "models", key(0))
|
|
storage.prune_cache(self.cache, model_limit_bytes=60)
|
|
self.assertTrue(first.exists())
|
|
self.assertFalse(old.exists())
|
|
|
|
def test_in_use_overage_is_reported_and_later_sweep_recovers(self):
|
|
first = self.entry(0, 60, 1)
|
|
second = self.entry(1, 60, 2)
|
|
a, b = self.lease(0), self.lease(1)
|
|
report = storage.prune_cache(self.cache, model_limit_bytes=60)
|
|
self.assertTrue(first.exists() and second.exists())
|
|
self.assertEqual(report["models"]["overLimitBytes"], 60)
|
|
self.assertEqual(report["models"]["skippedInUse"], 2)
|
|
a.close()
|
|
b.close()
|
|
report = storage.prune_cache(self.cache, model_limit_bytes=60)
|
|
self.assertFalse(first.exists())
|
|
self.assertTrue(second.exists())
|
|
self.assertEqual(report["models"]["overLimitBytes"], 0)
|
|
|
|
def test_only_final_oversized_entry_is_retained(self):
|
|
old = self.entry(0, 200, 1)
|
|
new = self.entry(1, 200, 2)
|
|
report = storage.prune_cache(self.cache, model_limit_bytes=100)
|
|
self.assertFalse(old.exists())
|
|
self.assertTrue(new.exists())
|
|
self.assertEqual(report["models"]["oversizedEntries"], 1)
|
|
self.assertEqual(report["models"]["overLimitBytes"], 100)
|
|
|
|
def test_lock_shards_are_bounded_and_collisions_only_defer_eviction(self):
|
|
for number in range(512):
|
|
with storage.acquire_cache_lease(self.cache, "models", key(number)):
|
|
pass
|
|
self.assertEqual(len(list((self.cache / ".locks").iterdir())), storage.LOCK_SHARDS)
|
|
self.assertTrue(all(path.stat().st_size == 0 for path in (self.cache / ".locks").iterdir()))
|
|
colliding = self.entry(128, 60, 1)
|
|
other = self.entry(1, 60, 2)
|
|
self.lease(0)
|
|
report = storage.prune_cache(self.cache, model_limit_bytes=60)
|
|
self.assertTrue(colliding.exists())
|
|
self.assertFalse(other.exists())
|
|
self.assertEqual(report["models"]["skippedInUse"], 1)
|
|
|
|
def test_legacy_unknown_and_linked_entries_are_not_followed_or_removed(self):
|
|
old = self.entry(0, 60, 1)
|
|
self.entry(1, 60, 2)
|
|
legacy = self.cache / key(50)
|
|
legacy.mkdir()
|
|
(legacy / "model").write_text("old cache")
|
|
unknown = self.cache / "models" / "notes.txt"
|
|
unknown.write_text("user notes")
|
|
outside = self.folder / "outside"
|
|
outside.mkdir()
|
|
(outside / "important").write_text("keep")
|
|
direct = self.cache / "models" / key(3)
|
|
nested = self.cache / "models" / key(4)
|
|
nested.mkdir()
|
|
try:
|
|
direct.symlink_to(outside, target_is_directory=True)
|
|
(nested / "outside").symlink_to(outside, target_is_directory=True)
|
|
except OSError as exc:
|
|
self.skipTest(f"Symlinks unavailable: {exc}")
|
|
report = storage.prune_cache(self.cache, model_limit_bytes=60)
|
|
self.assertFalse(old.exists())
|
|
self.assertEqual((legacy / "model").read_text(), "old cache")
|
|
self.assertEqual(unknown.read_text(), "user notes")
|
|
self.assertEqual((outside / "important").read_text(), "keep")
|
|
self.assertTrue(direct.is_symlink() and (nested / "outside").is_symlink())
|
|
self.assertEqual(report["models"]["skippedUnmanaged"], 3)
|
|
|
|
def test_invalid_keys_and_symlink_lock_files_are_rejected(self):
|
|
for kind, digest in (("../models", key(0)), ("models", "../outside"),
|
|
("models", "A" * 64), ("models", key(0) + "\n")):
|
|
with self.subTest(kind=kind, key=digest):
|
|
with self.assertRaises(ValueError):
|
|
storage.acquire_cache_lease(self.cache, kind, digest)
|
|
lock_dir = self.cache / ".locks"
|
|
lock_dir.mkdir(parents=True)
|
|
target = self.folder / "outside-lock"
|
|
target.write_text("untouched")
|
|
try:
|
|
(lock_dir / "models-000.lock").symlink_to(target)
|
|
except OSError as exc:
|
|
self.skipTest(f"Symlinks unavailable: {exc}")
|
|
with self.assertRaises(RuntimeError):
|
|
self.lease(0)
|
|
self.assertEqual(target.read_text(), "untouched")
|
|
|
|
def test_prune_serializes_sweeps_without_blocking_readers(self):
|
|
old = self.entry(0, 200, 1)
|
|
new = self.entry(1, 200, 2)
|
|
import shutil
|
|
original = shutil.rmtree
|
|
concurrent = []
|
|
|
|
def removing(path, *args, **kwargs):
|
|
concurrent.append(storage.prune_cache(self.cache, model_limit_bytes=100))
|
|
# Reader leases do not contend with the global sweep lock.
|
|
with self.lease(2, blocking=False):
|
|
pass
|
|
return original(path, *args, **kwargs)
|
|
|
|
with patch.object(storage.shutil, "rmtree", side_effect=removing):
|
|
storage.prune_cache(self.cache, model_limit_bytes=100)
|
|
self.assertFalse(old.exists())
|
|
self.assertTrue(new.exists())
|
|
self.assertEqual(len(concurrent), 1)
|
|
self.assertTrue(concurrent[0]["models"]["skippedConcurrentSweep"])
|
|
|
|
def stage(self, number: int, size: int, namespace: str = "root") -> Path:
|
|
parent = self.cache if namespace == "root" else self.cache / namespace
|
|
parent.mkdir(parents=True, exist_ok=True)
|
|
stage = Path(tempfile.mkdtemp(prefix=f"building-{key(number)}-", dir=parent))
|
|
(stage / "unfinished.o").write_bytes(b"x" * size)
|
|
return stage
|
|
|
|
def test_orphan_stages_are_removed_in_all_three_managed_locations(self):
|
|
root = self.stage(10, 30)
|
|
model = self.stage(11, 40, "models")
|
|
obj = self.stage(12, 50, "objects")
|
|
legacy = self.cache / "building-old12345"
|
|
legacy.mkdir()
|
|
unknown = self.cache / "models" / f"building-{key(13)}-short"
|
|
unknown.mkdir()
|
|
# Similar names inside an unrelated entry are outside the sweep scope.
|
|
nested = self.cache / key(14) / f"building-{key(15)}-abcdefgh"
|
|
nested.mkdir(parents=True)
|
|
report = storage.prune_cache(self.cache)
|
|
self.assertFalse(root.exists() or model.exists() or obj.exists())
|
|
self.assertTrue(legacy.exists() and unknown.exists() and nested.exists())
|
|
self.assertEqual(report["orphanStages"]["orphanStagesRemoved"], 3)
|
|
self.assertEqual(report["orphanStages"]["orphanStagesBytes"], 120)
|
|
self.assertEqual(report["orphanStages"]["skippedUnmanaged"], 2)
|
|
self.assertEqual(report["orphanStages"]["errors"], [])
|
|
|
|
def test_active_stages_use_their_correct_namespace_and_survive(self):
|
|
model_lease = self.lease(10)
|
|
object_lease = storage.acquire_cache_lease(self.cache, "objects", key(11))
|
|
self.addCleanup(object_lease.close)
|
|
root_active = self.stage(10, 30)
|
|
model_active = self.stage(10, 40, "models")
|
|
object_active = self.stage(11, 50, "objects")
|
|
collision = self.stage(10 + storage.LOCK_SHARDS, 20)
|
|
# Holding models/key(10) does not protect objects/key(10), or vice versa.
|
|
object_orphan = self.stage(10, 60, "objects")
|
|
model_orphan = self.stage(11, 70)
|
|
report = storage.prune_cache(self.cache)["orphanStages"]
|
|
self.assertTrue(all(path.exists() for path in (
|
|
root_active, model_active, object_active, collision,
|
|
)))
|
|
self.assertFalse(object_orphan.exists() or model_orphan.exists())
|
|
self.assertEqual(report["skippedInUse"], 4)
|
|
self.assertEqual(report["orphanStagesRemoved"], 2)
|
|
self.assertEqual(report["orphanStagesBytes"], 130)
|
|
model_lease.close()
|
|
object_lease.close()
|
|
report = storage.prune_cache(self.cache)["orphanStages"]
|
|
self.assertEqual(report["orphanStagesRemoved"], 4)
|
|
self.assertEqual(report["orphanStagesBytes"], 140)
|
|
|
|
def test_killed_builder_releases_stage_protection_for_next_sweep(self):
|
|
script = r"""
|
|
import sys, tempfile, time
|
|
from pathlib import Path
|
|
from app.simulation.native_codegen.cache_storage import acquire_cache_lease
|
|
cache, key = Path(sys.argv[1]), sys.argv[2]
|
|
lease = acquire_cache_lease(cache, 'models', key)
|
|
stage = Path(tempfile.mkdtemp(prefix='building-' + key + '-', dir=cache))
|
|
(stage / 'unfinished.o').write_bytes(b'x' * 37)
|
|
print(stage.name, flush=True)
|
|
time.sleep(60)
|
|
"""
|
|
process = subprocess.Popen(
|
|
[sys.executable, "-c", script, str(self.cache), key(16)], cwd=ROOT,
|
|
stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True,
|
|
)
|
|
self.addCleanup(lambda: process.kill() if process.poll() is None else None)
|
|
stage = self.cache / process.stdout.readline().strip()
|
|
self.assertTrue(stage.is_dir())
|
|
report = storage.prune_cache(self.cache)["orphanStages"]
|
|
self.assertTrue(stage.exists())
|
|
self.assertEqual(report["skippedInUse"], 1)
|
|
process.kill()
|
|
process.communicate(timeout=15)
|
|
report = storage.prune_cache(self.cache)["orphanStages"]
|
|
self.assertFalse(stage.exists())
|
|
self.assertEqual(report["orphanStagesRemoved"], 1)
|
|
self.assertEqual(report["orphanStagesBytes"], 37)
|
|
|
|
def test_orphan_stage_symlinks_and_symlink_namespaces_are_not_followed(self):
|
|
self.cache.mkdir()
|
|
outside = self.folder / "outside-stage"
|
|
outside.mkdir()
|
|
artifact = outside / "important"
|
|
artifact.write_text("keep")
|
|
direct = self.cache / f"building-{key(20)}-abcdefgh"
|
|
nested = self.stage(21, 10)
|
|
object_namespace = self.cache / "objects"
|
|
try:
|
|
direct.symlink_to(outside, target_is_directory=True)
|
|
(nested / "link").symlink_to(artifact)
|
|
object_namespace.symlink_to(outside, target_is_directory=True)
|
|
except OSError as exc:
|
|
self.skipTest(f"Symlinks unavailable: {exc}")
|
|
# A matching name behind a namespace link is never inspected or removed.
|
|
outside_stage = outside / f"building-{key(22)}-abcdefgh"
|
|
outside_stage.mkdir()
|
|
report = storage.prune_cache(self.cache)["orphanStages"]
|
|
self.assertEqual(report["orphanStagesRemoved"], 0)
|
|
self.assertEqual(report["skippedUnmanaged"], 2)
|
|
self.assertTrue(direct.is_symlink() and nested.exists() and outside_stage.exists())
|
|
self.assertEqual(artifact.read_text(), "keep")
|
|
|
|
def test_orphan_cleanup_failure_is_reported_without_removing_other_data(self):
|
|
stage = self.stage(24, 80)
|
|
with patch.object(storage.shutil, "rmtree", side_effect=PermissionError("busy stage")):
|
|
report = storage.prune_cache(self.cache)["orphanStages"]
|
|
self.assertTrue(stage.exists())
|
|
self.assertEqual(report["orphanStagesRemoved"], 0)
|
|
self.assertEqual(report["orphanStagesBytes"], 0)
|
|
self.assertEqual(len(report["errors"]), 1)
|
|
self.assertIn("busy stage", report["errors"][0])
|
|
|
|
def test_filesystem_cleanup_errors_do_not_fail_simulation(self):
|
|
self.entry(0, 60, 1)
|
|
self.entry(1, 60, 2)
|
|
with patch.object(storage.shutil, "rmtree", side_effect=PermissionError("busy")):
|
|
report = storage.prune_cache(self.cache, model_limit_bytes=60)
|
|
self.assertEqual(report["models"]["overLimitBytes"], 60)
|
|
self.assertTrue(report["models"]["errors"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|