From e2b0fa85eb4a99760886b5e11173753e5b6ddc76 Mon Sep 17 00:00:00 2001 From: yixi ning Date: Mon, 19 Oct 2026 03:49:39 +0000 Subject: [PATCH 1/2] fix(cache): publish complete disk cache entries atomically --- py/src/braintrust/prompt_cache/disk_cache.py | 29 +++++-- .../prompt_cache/test_disk_cache.py | 77 +++++++++++++++++++ 2 files changed, 99 insertions(+), 7 deletions(-) diff --git a/py/src/braintrust/prompt_cache/disk_cache.py b/py/src/braintrust/prompt_cache/disk_cache.py index ab20e0be4..eb85514b6 100644 --- a/py/src/braintrust/prompt_cache/disk_cache.py +++ b/py/src/braintrust/prompt_cache/disk_cache.py @@ -15,6 +15,7 @@ import json import logging import os +import tempfile from collections.abc import Callable from typing import Any, Generic, TypeVar @@ -23,6 +24,7 @@ log = logging.getLogger(__name__) +_TEMP_PREFIX = ".braintrust-cache-" class DiskCache(Generic[T]): @@ -116,14 +118,16 @@ 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) @@ -131,22 +135,33 @@ def set(self, key: str, value: T) -> None: os.makedirs(self._dir, exist_ok=True) 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")) + + os.replace(temporary_path, file_path) + 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 _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 7f49986ed..ee3f34d14 100644 --- a/py/src/braintrust/prompt_cache/test_disk_cache.py +++ b/py/src/braintrust/prompt_cache/test_disk_cache.py @@ -1,8 +1,10 @@ import os import shutil import tempfile +import threading import time import unittest +from concurrent.futures import ThreadPoolExecutor from typing import Any from braintrust import prompt @@ -24,6 +26,81 @@ 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_keys_with_invalid_paths(self): data = {"1": "2"} weird_keys = [ From 4c76ed680d6644f3a8aa58b6b8940a45e9eafca8 Mon Sep 17 00:00:00 2001 From: yixi ning Date: Mon, 19 Oct 2026 06:31:33 +0000 Subject: [PATCH 2/2] fix: retry Windows cache publication and reap expired staging files --- py/src/braintrust/prompt_cache/disk_cache.py | 31 ++++- .../prompt_cache/test_disk_cache.py | 110 ++++++++++++++++++ 2 files changed, 140 insertions(+), 1 deletion(-) diff --git a/py/src/braintrust/prompt_cache/disk_cache.py b/py/src/braintrust/prompt_cache/disk_cache.py index eb85514b6..16ea6a6f9 100644 --- a/py/src/braintrust/prompt_cache/disk_cache.py +++ b/py/src/braintrust/prompt_cache/disk_cache.py @@ -16,6 +16,7 @@ import logging import os import tempfile +import time from collections.abc import Callable from typing import Any, Generic, TypeVar @@ -25,6 +26,8 @@ log = logging.getLogger(__name__) _TEMP_PREFIX = ".braintrust-cache-" +_TEMP_MAX_AGE = 24 * 60 * 60 +_TEMP_CLEANUP_INTERVAL = 60 * 60 class DiskCache(Generic[T]): @@ -74,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.""" @@ -133,6 +137,7 @@ def set(self, key: str, value: T) -> None: # (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 tempfile.NamedTemporaryFile(dir=self._dir, prefix=_TEMP_PREFIX, delete=False) as temporary: @@ -142,7 +147,15 @@ def set(self, key: str, value: T) -> None: value = self._serializer(value) f.write(json.dumps(value).encode("utf-8")) - os.replace(temporary_path, file_path) + 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() @@ -157,6 +170,22 @@ def set(self, key: str, value: T) -> None: 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 diff --git a/py/src/braintrust/prompt_cache/test_disk_cache.py b/py/src/braintrust/prompt_cache/test_disk_cache.py index ee3f34d14..235d137ce 100644 --- a/py/src/braintrust/prompt_cache/test_disk_cache.py +++ b/py/src/braintrust/prompt_cache/test_disk_cache.py @@ -1,11 +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 @@ -101,6 +104,113 @@ def serialize(value): 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 = [