fix: serialize state before truncating in the direct-write fallback

This commit is contained in:
Matt Van Horn
2026-07-05 04:04:54 -07:00
parent 96e88a3555
commit b00d4785ee
2 changed files with 26 additions and 5 deletions
+9 -5
View File
@@ -126,6 +126,7 @@ class AtomicJsonStore:
self._direct_write(payload) self._direct_write(payload)
def _atomic_write(self, payload: dict[str, Any]) -> None: def _atomic_write(self, payload: dict[str, Any]) -> None:
text = self._serialize(payload)
parent = os.path.dirname(self.path) or "." parent = os.path.dirname(self.path) or "."
fd, tmp_path = tempfile.mkstemp( fd, tmp_path = tempfile.mkstemp(
prefix=f".{os.path.basename(self.path)}.", prefix=f".{os.path.basename(self.path)}.",
@@ -135,7 +136,7 @@ class AtomicJsonStore:
) )
try: try:
with os.fdopen(fd, "w", encoding="utf-8") as f: with os.fdopen(fd, "w", encoding="utf-8") as f:
self._write_payload(payload, f) f.write(text)
f.flush() f.flush()
self._best_effort_fsync(f.fileno()) self._best_effort_fsync(f.fileno())
os.replace(tmp_path, self.path) os.replace(tmp_path, self.path)
@@ -148,6 +149,10 @@ class AtomicJsonStore:
raise raise
def _direct_write(self, payload: dict[str, Any]) -> None: def _direct_write(self, payload: dict[str, Any]) -> None:
# Serialize before truncating so a serialization failure never destroys
# the existing state file (the atomic path gets this for free via its
# temp file).
text = self._serialize(payload)
# Create with 0o600 so the fallback keeps the owner-only permissions the # Create with 0o600 so the fallback keeps the owner-only permissions the
# atomic path gets from mkstemp; state files can contain URLs and # atomic path gets from mkstemp; state files can contain URLs and
# per-download option overrides that must not leak on shared mounts. # per-download option overrides that must not leak on shared mounts.
@@ -161,7 +166,7 @@ class AtomicJsonStore:
os.fchmod(f.fileno(), 0o600) os.fchmod(f.fileno(), 0o600)
except OSError: except OSError:
pass pass
self._write_payload(payload, f) f.write(text)
f.flush() f.flush()
self._best_effort_fsync(f.fileno()) self._best_effort_fsync(f.fileno())
@@ -179,9 +184,8 @@ class AtomicJsonStore:
raise raise
@staticmethod @staticmethod
def _write_payload(payload: dict[str, Any], f: Any) -> None: def _serialize(payload: dict[str, Any]) -> str:
json.dump(payload, f, ensure_ascii=False, separators=(",", ":")) return json.dumps(payload, ensure_ascii=False, separators=(",", ":")) + "\n"
f.write("\n")
def _warn_direct_write_fallback(self, exc: OSError) -> None: def _warn_direct_write_fallback(self, exc: OSError) -> None:
if self._direct_write_fallback_warned: if self._direct_write_fallback_warned:
+17
View File
@@ -134,6 +134,23 @@ class StateStoreTests(unittest.TestCase):
# Existing state is untouched. # Existing state is untouched.
self.assertEqual(store.load()["items"], [{"key": "good"}]) self.assertEqual(store.load()["items"], [{"key": "good"}])
def test_serialization_failure_preserves_existing_state(self):
with tempfile.TemporaryDirectory() as tmp:
path = os.path.join(tmp, "queue.json")
store = AtomicJsonStore(path, kind="persistent_queue:queue")
store.save({"items": [{"key": "good"}]})
# Even on the fallback path, a non-serializable payload must raise
# before the existing good state file is touched.
with patch(
"state_store.tempfile.mkstemp",
side_effect=PermissionError(1, "Operation not permitted"),
):
with self.assertRaises(TypeError):
store.save({"items": object()})
self.assertEqual(store.load()["items"], [{"key": "good"}])
def test_invalid_file_is_quarantined(self): def test_invalid_file_is_quarantined(self):
with tempfile.TemporaryDirectory() as tmp: with tempfile.TemporaryDirectory() as tmp:
path = os.path.join(tmp, "queue.json") path = os.path.join(tmp, "queue.json")