diff --git a/py/src/braintrust/prompt_cache/disk_cache.py b/py/src/braintrust/prompt_cache/disk_cache.py index ab20e0be..16ea6a6f 100644 --- a/py/src/braintrust/prompt_cache/disk_cache.py +++ b/py/src/braintrust/prompt_cache/disk_cache.py @@ -15,6 +15,8 @@ import json import logging import os +import tempfile +import time from collections.abc import Callable from typing import Any, Generic, TypeVar @@ -23,6 +25,9 @@ log = logging.getLogger(__name__) +_TEMP_PREFIX = ".braintrust-cache-" +_TEMP_MAX_AGE = 24 * 60 * 60 +_TEMP_CLEANUP_INTERVAL = 60 * 60 class DiskCache(Generic[T]): @@ -72,6 +77,7 @@ def __init__( self._deserializer = deserializer self._log_warnings = log_warnings self._mkdirs = mkdirs + self._last_temp_cleanup: float | None = None def _get_entry_path(self, key: str) -> str: """Gets the file path for a cache entry.""" @@ -116,37 +122,75 @@ def set(self, key: str, value: T) -> None: """ Stores a value in the cache. If the cache is at its maximum size, the least recently used entries will be evicted. + Publish the entry only after serialization and compression complete, so concurrent + readers see a complete old or new value. Failed writes leave the old entry intact. Args: key: The key to store the value under. value: The value to store in the cache. - Raises: - RuntimeError: If there is an error writing to the disk cache. + Write errors are logged and ignored. """ + temporary_path: str | None = None try: # mkdirs exists only to make it easy to simulate cross-platform write errors # (permissions, etc wouldn't work on github actions on windows) if self._mkdirs: os.makedirs(self._dir, exist_ok=True) + self._cleanup_temporary_files() file_path = self._get_entry_path(key) - with gzip.open(file_path, "wb") as f: - if self._serializer is not None: - value = self._serializer(value) - f.write(json.dumps(value).encode("utf-8")) + with tempfile.NamedTemporaryFile(dir=self._dir, prefix=_TEMP_PREFIX, delete=False) as temporary: + temporary_path = temporary.name + with gzip.GzipFile(fileobj=temporary, mode="wb") as f: + if self._serializer is not None: + value = self._serializer(value) + f.write(json.dumps(value).encode("utf-8")) + + for attempt in range(5): + try: + os.replace(temporary_path, file_path) + break + except PermissionError as e: + # Windows readers can briefly deny replacement of an open file. + if getattr(e, "winerror", None) not in (5, 32, 33) or attempt == 4: + raise + time.sleep(0.01 * 2**attempt) + temporary_path = None self._evict_if_full() except Exception as e: # Swallow any cache write errors. Don't crash the app. if self._log_warnings: log.warning(f"Failed to write to disk cache: {e}") + finally: + if temporary_path is not None: + try: + os.unlink(temporary_path) + except OSError: + pass + + def _cleanup_temporary_files(self): + now = time.time() + if self._last_temp_cleanup is not None and now - self._last_temp_cleanup < _TEMP_CLEANUP_INTERVAL: + return + self._last_temp_cleanup = now + # A killed writer cannot run its finally block. Allow a day for active writes, + # then reap abandoned staging files, including for unbounded caches. + for name in os.listdir(self._dir): + if name.startswith(_TEMP_PREFIX): + path = os.path.join(self._dir, name) + try: + if os.path.getmtime(path) < now - _TEMP_MAX_AGE: + os.unlink(path) + except OSError: + pass def _evict_if_full(self): if self._max_size is None or self._max_size <= 0: return None - paths = [os.path.join(self._dir, f) for f in os.listdir(self._dir)] + paths = [os.path.join(self._dir, f) for f in os.listdir(self._dir) if not f.startswith(_TEMP_PREFIX)] if not paths or len(paths) <= self._max_size: return diff --git a/py/src/braintrust/prompt_cache/test_disk_cache.py b/py/src/braintrust/prompt_cache/test_disk_cache.py index 7f49986e..235d137c 100644 --- a/py/src/braintrust/prompt_cache/test_disk_cache.py +++ b/py/src/braintrust/prompt_cache/test_disk_cache.py @@ -1,9 +1,14 @@ import os import shutil +import subprocess +import sys import tempfile +import threading import time import unittest +from concurrent.futures import ThreadPoolExecutor from typing import Any +from unittest.mock import patch from braintrust import prompt from braintrust.prompt_cache import disk_cache @@ -24,6 +29,188 @@ def tearDown(self): except Exception: pass + def test_failed_json_encoding_preserves_previous_entry(self): + self.cache.set("key", {"version": 1}) + + self.cache.set("key", {"unsupported": {1, 2}}) + + self.assertEqual(self.cache.get("key"), {"version": 1}) + self.assertEqual(os.listdir(self.cache_dir), [os.path.basename(self.cache._get_entry_path("key"))]) + + def test_failed_serializer_preserves_previous_entry(self): + def serialize(value): + if value["version"] == 2: + raise ValueError("Cannot serialize this version") + return value + + cache = disk_cache.DiskCache[dict](self.cache_dir, serializer=serialize, log_warnings=False) + cache.set("key", {"version": 1}) + + cache.set("key", {"version": 2}) + + self.assertEqual(cache.get("key"), {"version": 1}) + self.assertEqual(os.listdir(self.cache_dir), [os.path.basename(cache._get_entry_path("key"))]) + + def test_read_during_replacement_returns_complete_previous_entry(self): + started = threading.Event() + release = threading.Event() + + def serialize(value): + if value == "new": + started.set() + assert release.wait(timeout=5) + return value + + cache = disk_cache.DiskCache[str](self.cache_dir, serializer=serialize, log_warnings=False) + cache.set("key", "old") + with ThreadPoolExecutor(max_workers=1) as executor: + write = executor.submit(cache.set, "key", "new") + try: + self.assertTrue(started.wait(timeout=5)) + self.assertEqual(cache.get("key"), "old") + finally: + release.set() + write.result(timeout=5) + + self.assertEqual(cache.get("key"), "new") + + def test_eviction_does_not_remove_an_in_progress_write(self): + started = threading.Event() + release = threading.Event() + + def serialize(value): + if value == "pending": + started.set() + assert release.wait(timeout=5) + return value + + cache = disk_cache.DiskCache[str](self.cache_dir, max_size=1, serializer=serialize, log_warnings=False) + cache.set("existing", "old") + with ThreadPoolExecutor(max_workers=1) as executor: + write = executor.submit(cache.set, "pending", "pending") + try: + self.assertTrue(started.wait(timeout=5)) + cache.set("existing", "updated") + self.assertEqual(cache.get("existing"), "updated") + # Avoid equal-mtime ties on filesystems with a coarse clock. + old_time = time.time() - 60 + os.utime(cache._get_entry_path("existing"), (old_time, old_time)) + finally: + release.set() + write.result(timeout=5) + + self.assertEqual(cache.get("pending"), "pending") + with self.assertRaises(KeyError): + cache.get("existing") + self.assertEqual(len(os.listdir(self.cache_dir)), 1) + + def test_retries_windows_sharing_violations(self): + self.cache.set("key", {"version": 1}) + replace = os.replace + sharing_violation = PermissionError("The file is in use") + sharing_violation.winerror = 32 + calls = 0 + + def replace_after_reader_closes(source, destination): + nonlocal calls + calls += 1 + if calls == 1: + raise sharing_violation + return replace(source, destination) + + with patch.object(disk_cache.os, "replace", side_effect=replace_after_reader_closes): + self.cache.set("key", {"version": 2}) + + self.assertEqual(self.cache.get("key"), {"version": 2}) + self.assertEqual(calls, 2) + self.assertEqual(len(os.listdir(self.cache_dir)), 1) + + def test_permanent_replace_error_preserves_previous_entry(self): + self.cache.set("key", {"version": 1}) + with patch.object(disk_cache.os, "replace", side_effect=PermissionError("Permission denied")) as replace: + self.cache.set("key", {"version": 2}) + + self.assertEqual(self.cache.get("key"), {"version": 1}) + self.assertEqual(replace.call_count, 1) + self.assertEqual(len(os.listdir(self.cache_dir)), 1) + + @unittest.skipUnless(os.name == "nt", "Windows file sharing behavior") + def test_replacement_overlaps_an_open_windows_reader(self): + started = threading.Event() + release = threading.Event() + + def deserialize(value): + if value == "old": + started.set() + assert release.wait(timeout=5) + return value + + cache = disk_cache.DiskCache[str](self.cache_dir, deserializer=deserialize, log_warnings=False) + cache.set("key", "old") + replace = os.replace + + def replace_and_release_reader(source, destination): + try: + return replace(source, destination) + except PermissionError: + release.set() + raise + + with ThreadPoolExecutor(max_workers=1) as executor: + read = executor.submit(cache.get, "key") + try: + self.assertTrue(started.wait(timeout=5)) + with patch.object(disk_cache.os, "replace", side_effect=replace_and_release_reader): + cache.set("key", "new") + finally: + release.set() + self.assertEqual(read.result(timeout=5), "old") + + self.assertEqual(cache.get("key"), "new") + + def test_persistent_windows_sharing_violation_is_bounded(self): + self.cache.set("key", {"version": 1}) + sharing_violation = PermissionError("The file is in use") + sharing_violation.winerror = 32 + with patch.object(disk_cache.os, "replace", side_effect=sharing_violation) as replace: + self.cache.set("key", {"version": 2}) + + self.assertEqual(self.cache.get("key"), {"version": 1}) + self.assertEqual(replace.call_count, 5) + self.assertEqual(len(os.listdir(self.cache_dir)), 1) + + def test_reaps_old_staging_files_left_by_an_interrupted_process(self): + # Exit in the serializer, after the temporary gzip file has been opened. + result = subprocess.run( + [ + sys.executable, + "-c", + "import os, sys; from braintrust.prompt_cache.disk_cache import DiskCache; " + "DiskCache(sys.argv[1], serializer=lambda value: os._exit(0)).set('key', 'value')", + self.cache_dir, + ], + check=True, + capture_output=True, + timeout=10, + ) + self.assertEqual(result.returncode, 0) + staging_files = os.listdir(self.cache_dir) + self.assertEqual(len(staging_files), 1) + orphan = os.path.join(self.cache_dir, staging_files[0]) + old_time = time.time() - 2 * 24 * 60 * 60 + os.utime(orphan, (old_time, old_time)) + with tempfile.NamedTemporaryFile(dir=self.cache_dir, prefix=".braintrust-cache-", delete=False) as active: + active_path = active.name + + # Cleanup also applies to unbounded caches and leaves recent writers alone. + cache = disk_cache.DiskCache[str](self.cache_dir) + cache.set("new", "value") + + self.assertEqual(cache.get("new"), "value") + self.assertFalse(os.path.exists(orphan)) + self.assertTrue(os.path.exists(active_path)) + self.assertEqual(len(os.listdir(self.cache_dir)), 2) + def test_keys_with_invalid_paths(self): data = {"1": "2"} weird_keys = [