fix(cache): publish complete disk cache entries atomically

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
 

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 = [