fix(memory): ensure atomic write for history.jsonl
Use temp file + os.replace + fsync to prevent partial writes on crash. Add tests for atomic write behavior and tmp file cleanup on exception.
This commit is contained in:
+11
-2
@@ -4,6 +4,7 @@ from __future__ import annotations
|
|||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
import json
|
import json
|
||||||
|
import os
|
||||||
import re
|
import re
|
||||||
import weakref
|
import weakref
|
||||||
import tiktoken
|
import tiktoken
|
||||||
@@ -359,10 +360,18 @@ class MemoryStore:
|
|||||||
return None
|
return None
|
||||||
|
|
||||||
def _write_entries(self, entries: list[dict[str, Any]]) -> None:
|
def _write_entries(self, entries: list[dict[str, Any]]) -> None:
|
||||||
"""Overwrite history.jsonl with the given entries."""
|
"""Overwrite history.jsonl with the given entries (atomic write)."""
|
||||||
with open(self.history_file, "w", encoding="utf-8") as f:
|
tmp_path = self.history_file.with_suffix(self.history_file.suffix + ".tmp")
|
||||||
|
try:
|
||||||
|
with open(tmp_path, "w", encoding="utf-8") as f:
|
||||||
for entry in entries:
|
for entry in entries:
|
||||||
f.write(json.dumps(entry, ensure_ascii=False) + "\n")
|
f.write(json.dumps(entry, ensure_ascii=False) + "\n")
|
||||||
|
f.flush()
|
||||||
|
os.fsync(f.fileno())
|
||||||
|
os.replace(tmp_path, self.history_file)
|
||||||
|
except BaseException:
|
||||||
|
tmp_path.unlink(missing_ok=True)
|
||||||
|
raise
|
||||||
|
|
||||||
# -- dream cursor --------------------------------------------------------
|
# -- dream cursor --------------------------------------------------------
|
||||||
|
|
||||||
|
|||||||
@@ -141,6 +141,54 @@ class TestHistoryWithCursor:
|
|||||||
assert len(entries) == 2
|
assert len(entries) == 2
|
||||||
assert entries[0]["cursor"] in {4, 5}
|
assert entries[0]["cursor"] in {4, 5}
|
||||||
|
|
||||||
|
def test_write_entries_uses_atomic_write(self, tmp_path):
|
||||||
|
"""_write_entries uses temp file + os.replace for atomicity."""
|
||||||
|
store = MemoryStore(tmp_path)
|
||||||
|
store.append_history("event 1")
|
||||||
|
store.append_history("event 2")
|
||||||
|
store.append_history("event 3")
|
||||||
|
entries = store.read_unprocessed_history(since_cursor=0)
|
||||||
|
|
||||||
|
# Monitor temp file existence
|
||||||
|
tmp_path_obj = store.history_file.with_suffix(".jsonl.tmp")
|
||||||
|
assert not tmp_path_obj.exists() # Should not exist initially
|
||||||
|
|
||||||
|
# Call _write_entries
|
||||||
|
store._write_entries(entries)
|
||||||
|
|
||||||
|
# Temp file should be cleaned up
|
||||||
|
assert not tmp_path_obj.exists()
|
||||||
|
# Original file should exist
|
||||||
|
assert store.history_file.exists()
|
||||||
|
|
||||||
|
def test_write_entries_cleans_up_tmp_on_exception(self, tmp_path, monkeypatch):
|
||||||
|
"""Exception during _write_entries cleans up the temp file."""
|
||||||
|
store = MemoryStore(tmp_path)
|
||||||
|
store.append_history("event 1")
|
||||||
|
entries = store.read_unprocessed_history(since_cursor=0)
|
||||||
|
|
||||||
|
tmp_path_obj = store.history_file.with_suffix(".jsonl.tmp")
|
||||||
|
|
||||||
|
# Mock os.replace to raise an exception
|
||||||
|
original_replace = __import__('os').replace
|
||||||
|
|
||||||
|
def failing_replace(*args, **kwargs):
|
||||||
|
raise RuntimeError("Simulated failure")
|
||||||
|
|
||||||
|
monkeypatch.setattr('os.replace', failing_replace)
|
||||||
|
|
||||||
|
try:
|
||||||
|
store._write_entries(entries)
|
||||||
|
assert False, "Should have raised"
|
||||||
|
except RuntimeError:
|
||||||
|
pass
|
||||||
|
|
||||||
|
# Temp file should be cleaned up
|
||||||
|
assert not tmp_path_obj.exists()
|
||||||
|
|
||||||
|
# Original file should still exist (because replace failed)
|
||||||
|
assert store.history_file.exists()
|
||||||
|
|
||||||
|
|
||||||
class TestAppendHistoryHardCap:
|
class TestAppendHistoryHardCap:
|
||||||
"""append_history has a defensive cap that catches new callers who forgot
|
"""append_history has a defensive cap that catches new callers who forgot
|
||||||
|
|||||||
Reference in New Issue
Block a user