fix(runtime): preserve per-run files in JSONL batch writes (#4938)

* fix(runtime): preserve per-run JSONL batch files

* fix(events): clarify JSONL rollback failures
This commit is contained in:
Roc 2026-08-23 15:00:37 +08:00 committed by GitHub
parent e5bf3ccf45
commit c88c2975c9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 155 additions and 32 deletions

View File

@ -157,13 +157,17 @@ class JsonlRunEventStore(RunEventStore):
return record
async def put_batch(self, events):
"""Persist a batch of events atomically per-thread.
"""Persist a batch of events under a per-thread write lock.
All seq numbers for the batch are reserved under a single per-thread
write lock and every record is appended in one file write so a
mid-batch failure cannot leave a partial set of records on disk that
a retry would then duplicate. Callers (e.g. worker.py's flush-retry
path) may safely re-buffer the entire batch on failure.
write lock. Records are grouped by run_id and appended to their own
run files while that lock is held. If a write fails and rollback
succeeds, already-appended groups for the current thread are restored
so callers (e.g. worker.py's flush-retry path) may safely re-buffer
that thread's batch. When a batch contains multiple thread IDs, thread
groups are processed sequentially, so a later failure does not roll
back earlier thread groups. This rollback does not make a multi-file
batch crash-atomic.
"""
if not events:
return []
@ -226,10 +230,11 @@ class JsonlRunEventStore(RunEventStore):
"created_at": ev.get("created_at") or datetime.now(UTC).isoformat(),
}
records.append(record)
path = self._run_file(thread_id, batch[0]["run_id"])
# Single append/write per thread. If this raises, no records were
# persisted; the caller's re-buffer reproduces no duplicates.
await asyncio.to_thread(self._append_records, path, records)
records_by_run: dict[str, list[dict[str, Any]]] = {}
for record in records:
records_by_run.setdefault(record["run_id"], []).append(record)
run_batches = [(self._run_file(thread_id, run_id), run_records) for run_id, run_records in records_by_run.items()]
await asyncio.to_thread(self._append_record_groups, run_batches)
return records
def _append_records(self, path: Path, records: list[dict[str, Any]]) -> None:
@ -238,6 +243,30 @@ class JsonlRunEventStore(RunEventStore):
with open(path, "a", encoding="utf-8") as f:
f.write(lines)
def _append_record_groups(self, groups: list[tuple[Path, list[dict[str, Any]]]]) -> None:
"""Append run groups and restore their original sizes if one fails."""
original_sizes: dict[Path, int | None] = {}
try:
for path, records in groups:
original_sizes[path] = path.stat().st_size if path.exists() else None
self._append_records(path, records)
except Exception:
for path, original_size in original_sizes.items():
try:
if original_size is None:
if path.exists():
path.unlink()
else:
with open(path, "r+b") as f:
f.truncate(original_size)
except OSError:
logger.error(
"Failed to roll back JSONL batch append for %s; retrying the batch may create duplicate records",
path,
exc_info=True,
)
raise
async def list_messages(self, thread_id, *, limit=50, before_seq=None, after_seq=None, user_id: str | None | _AutoSentinel = AUTO):
all_events = await asyncio.to_thread(self._read_thread_events, thread_id)
messages = [e for e in all_events if e.get("category") == "message"]

View File

@ -10,6 +10,7 @@ Verifies:
from __future__ import annotations
import asyncio
import logging
import tempfile
from pathlib import Path
@ -88,7 +89,7 @@ async def test_concurrent_puts_different_threads_independent_seqs():
# ---------------------------------------------------------------------------
# put_batch: delegates to put() and preserves order
# put_batch: assigns monotonic seqs and preserves per-run files
# ---------------------------------------------------------------------------
@ -103,6 +104,24 @@ async def test_put_batch_seqs_are_monotonic():
assert len(set(seqs)) == 5
@pytest.mark.anyio
async def test_put_batch_writes_mixed_run_ids_to_their_run_files():
with tempfile.TemporaryDirectory() as tmp:
store = _make_store(Path(tmp))
events = [
{"thread_id": "t1", "run_id": "r1", "event_type": "human_message", "category": "message", "content": "r1-first"},
{"thread_id": "t1", "run_id": "r2", "event_type": "human_message", "category": "message", "content": "r2-only"},
{"thread_id": "t1", "run_id": "r1", "event_type": "ai_message", "category": "message", "content": "r1-last"},
]
records = await store.put_batch(events)
r1_messages = await store.list_messages_by_run("t1", "r1")
r2_messages = await store.list_messages_by_run("t1", "r2")
assert [(record["run_id"], record["seq"]) for record in records] == [("r1", 1), ("r2", 2), ("r1", 3)]
assert [message["content"] for message in r1_messages] == ["r1-first", "r1-last"]
assert [message["content"] for message in r2_messages] == ["r2-only"]
# ---------------------------------------------------------------------------
# _ensure_seq_loaded: recovers max_seq from disk after fresh store init
# ---------------------------------------------------------------------------
@ -148,23 +167,15 @@ async def test_put_offloads_write_via_to_thread():
# ---------------------------------------------------------------------------
# put_batch atomicity: a failed append must not leave partial records so a
# caller re-buffering the batch on retry does not produce duplicates.
# put_batch failure rollback: a failed append must not leave partial records
# so a caller re-buffering the batch on retry does not produce duplicates.
# Regression for deer-flow PR #4082 (review feedback from willem-bd).
# ---------------------------------------------------------------------------
@pytest.mark.anyio
async def test_put_batch_failure_rolls_back_no_partial_records(monkeypatch):
"""If the disk write inside ``put_batch`` raises after partial output,
no records should remain on disk because the seq counter is reserved
under the write lock but the seqs were NOT written. A subsequent retry
therefore reproduces no duplicates.
Concretely: the implementation uses a single ``open().write()`` so on
failure the file is either empty or has the prior batch's records —
never a partial slice of the new batch.
"""
"""A failed append is rolled back before the re-buffered batch is retried."""
import json
from deerflow.runtime.events.store import jsonl as jsonl_mod
@ -194,26 +205,109 @@ async def test_put_batch_failure_rolls_back_no_partial_records(monkeypatch):
}
for i in range(4)
]
# First attempt — fails mid-batch; expect raise; the file may have
# partial lines but the in-memory seq counter has been advanced
# (because seq reservation happened under the lock).
# First attempt — fails after partial output; expect raise. The
# in-memory seq counter is advanced because reservation happens under
# the lock, but the partial file contents must be rolled back.
with pytest.raises(OSError):
await store.put_batch(events)
# Now retry with the real append (no failure): only the unreserved
# records will be written — but our implementation appends the whole
# batch again, so what we really verify here is that after a failure
# the seq counter is monotonic and consistent with the recovered
# disk state (no half-batch leftover gets accidentally re-numbered).
# Retry the full batch with the real append, matching worker.py's
# re-buffer path. Verify persisted contents below, not only return
# values, so a duplicate or partial disk write cannot pass unnoticed.
monkeypatch.setattr(jsonl_mod.JsonlRunEventStore, "_append_records", real_append)
# Retry the full batch — the re-buffer pattern from worker.py.
records = await store.put_batch(events)
persisted_events = await store.list_events("t1", "r1")
# The batch succeeded on retry, every event ended up exactly once in the
# file (no duplicates), and seqs are still strictly monotonic.
# The batch succeeded on retry, every event ended up exactly once on disk,
# and seqs are still strictly monotonic.
assert len(records) == 4, f"Expected 4 records, got {len(records)}"
seqs = [r["seq"] for r in records]
assert seqs == sorted(seqs) and len(set(seqs)) == 4, f"seqs not unique monotonic: {seqs}"
assert len(persisted_events) == 4
assert [event["content"] for event in persisted_events] == [f"event-{i}" for i in range(4)]
assert [event["seq"] for event in persisted_events] == seqs
@pytest.mark.anyio
async def test_mixed_run_batch_failure_restores_all_run_files(monkeypatch):
"""A failed mixed-run append restores prior bytes in every touched file."""
from deerflow.runtime.events.store import jsonl as jsonl_mod
real_append = jsonl_mod.JsonlRunEventStore._append_records
append_calls = 0
def failing_append(self, path, records):
nonlocal append_calls
append_calls += 1
if append_calls == 2:
path.parent.mkdir(parents=True, exist_ok=True)
with open(path, "a", encoding="utf-8") as f:
f.write("partial\n")
raise OSError("simulated second-run write failure")
real_append(self, path, records)
monkeypatch.setattr(jsonl_mod.JsonlRunEventStore, "_append_records", failing_append)
with tempfile.TemporaryDirectory() as tmp:
store = _make_store(Path(tmp))
await store.put(thread_id="t1", run_id="r1", event_type="trace", category="trace", content="existing")
r1_path = store._run_file("t1", "r1")
r2_path = store._run_file("t1", "r2")
original_r1 = r1_path.read_bytes()
events = [
{"thread_id": "t1", "run_id": "r1", "event_type": "trace", "category": "trace", "content": "new-r1"},
{"thread_id": "t1", "run_id": "r2", "event_type": "trace", "category": "trace", "content": "new-r2"},
]
with pytest.raises(OSError, match="second-run"):
await store.put_batch(events)
assert r1_path.read_bytes() == original_r1
assert not r2_path.exists()
assert [event["content"] for event in await store.list_events("t1", "r1")] == ["existing"]
assert await store.list_events("t1", "r2") == []
@pytest.mark.anyio
async def test_mixed_run_batch_logs_error_when_rollback_fails(monkeypatch, caplog):
"""A rollback failure must make possible retry duplicates visible to operators."""
from deerflow.runtime.events.store import jsonl as jsonl_mod
real_append = jsonl_mod.JsonlRunEventStore._append_records
real_unlink = Path.unlink
append_calls = 0
def failing_append(self, path, records):
nonlocal append_calls
append_calls += 1
if append_calls == 2:
real_append(self, path, records)
raise OSError("simulated second-run write failure")
real_append(self, path, records)
monkeypatch.setattr(jsonl_mod.JsonlRunEventStore, "_append_records", failing_append)
with tempfile.TemporaryDirectory() as tmp:
store = _make_store(Path(tmp))
r2_path = store._run_file("t1", "r2")
def failing_unlink(path, missing_ok=False):
if path == r2_path:
raise OSError("simulated rollback failure")
return real_unlink(path, missing_ok=missing_ok)
monkeypatch.setattr(Path, "unlink", failing_unlink)
events = [
{"thread_id": "t1", "run_id": "r1", "event_type": "trace", "category": "trace"},
{"thread_id": "t1", "run_id": "r2", "event_type": "trace", "category": "trace"},
]
with caplog.at_level(logging.ERROR, logger=jsonl_mod.__name__), pytest.raises(OSError, match="second-run"):
await store.put_batch(events)
assert r2_path.exists()
rollback_errors = [record for record in caplog.records if record.levelno >= logging.ERROR]
assert any("duplicate records" in record.getMessage() for record in rollback_errors)
# ---------------------------------------------------------------------------