diff --git a/backend/packages/harness/deerflow/runtime/events/store/jsonl.py b/backend/packages/harness/deerflow/runtime/events/store/jsonl.py index 105ce894d2..440d361bc2 100644 --- a/backend/packages/harness/deerflow/runtime/events/store/jsonl.py +++ b/backend/packages/harness/deerflow/runtime/events/store/jsonl.py @@ -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"] diff --git a/backend/tests/test_jsonl_event_store_async_io.py b/backend/tests/test_jsonl_event_store_async_io.py index db254113fa..677262fd85 100644 --- a/backend/tests/test_jsonl_event_store_async_io.py +++ b/backend/tests/test_jsonl_event_store_async_io.py @@ -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) # ---------------------------------------------------------------------------