fix(events): drain JSONL mutations before propagating cancellation (#5439)

* fix(events): drain JSONL mutations before propagating cancellation

Signed-off-by: Undermoon1412 <Undermoon1412@users.noreply.github.com>

* fix(events): cover cross-thread cancellation and name drain tasks

Signed-off-by: Undermoon1412 <Undermoon1412@users.noreply.github.com>

---------

Signed-off-by: Undermoon1412 <Undermoon1412@users.noreply.github.com>
Co-authored-by: Undermoon1412 <Undermoon1412@users.noreply.github.com>
This commit is contained in:
Undermoon1412 2026-09-15 07:28:24 +08:00 committed by GitHub
parent 2536f24f81
commit 6c697f3067
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 321 additions and 8 deletions

View File

@ -351,6 +351,8 @@ DeerFlow still uses `Forwarded` / `X-Forwarded-*` headers to recover the browser
> [!IMPORTANT]
> The Gateway still owns active run tasks in process, so production defaults to a single Gateway worker (`GATEWAY_WORKERS=1`). Multi-worker deployments require Postgres, the Redis stream bridge (`stream_bridge.type: redis`), `run_ownership.heartbeat_enabled: true`, and `run_events.backend: db`; process-local memory/JSONL event stores cannot enforce singleton delivery receipts across workers. The bridge shares SSE delivery and bounded `Last-Event-ID` replay across workers. When a valid reconnect cursor has been trimmed, or a subscriber that already established an empty-stream wait falls behind before its first delivery, Memory and Redis emit a machine-readable SSE `gap` event instead of silently returning a partial replay; the Web UI reloads durable thread/event state and resumes from the retained tail. Lease reconciliation marks runs from dead workers as errors, persists their delivery receipts, publishes the terminal stream marker, schedules retained-stream cleanup, and updates the affected thread status. SSE, `/wait`, and internal stream consumers use `stream_bridge.heartbeat_interval_seconds` (default `15`) for idle liveness checks; changing it requires a Gateway restart. Malformed Redis reconnect IDs live-tail new events instead of replaying the retained buffer, and the rolling retained-buffer TTL (`stream_ttl_seconds`) remains a cleanup safety net rather than a run timeout. IM channel state and other process-local services still need their own multi-worker coordination.
>
> In single-process JSONL deployments, cancelling an admitted event-store mutation waits for its background file I/O, rollback, and bookkeeping to settle before releasing the thread write lock. This prevents an older cancelled write from recreating deleted records or rolling back a later successful write. Cancellation can therefore wait on slow storage; it does not stop an in-flight filesystem operation. Callers still waiting to acquire the lock can cancel without starting a mutation. A batch spanning multiple threads drains its current thread group before propagating cancellation; subsequent thread groups do not start.
>
> After a run publishes its terminal stream marker, its process-local `RunRecord` remains available for the existing five-minute grace period before cleanup; durable run history remains available through `RunStore`, while the stream bridge retains its delivery tail on its separate cleanup schedule.
>
> Run cancellation may land on any Gateway worker. A non-owning worker now persists the interrupt or rollback request for the live owner, which observes it during lease renewal and performs the normal cancellation flow; load-balancer routing alone no longer produces a 409. The first accepted action wins even if a retry lands on the owner, and accepted cancellation competes atomically with owner completion. Dead owners still follow lease takeover and orphan recovery. Cancellation latency is therefore bounded by the lease heartbeat interval.

View File

@ -289,3 +289,19 @@ rejects caller-supplied `__conversation_reader` values in both context carriers,
installs only the host value, and releases it during terminal cleanup. The
callback is not checkpoint state and must never be recovered from an earlier
run or serialized into run kwargs.
## JSONL mutation cancellation
`JsonlRunEventStore._run_mutation` acquires the per-thread lock before admitting
an operation, then drains the shielded operation through filesystem I/O, rollback,
and sequence/lock bookkeeping before releasing the lock or re-raising caller
cancellation. Repeated cancellation must not detach an active disk worker; a failed
mutation remains the cause of the propagated cancellation. There is deliberately
no drain timeout that would release ownership while a worker can still modify files.
A queued caller can cancel before admission, and unrelated threads remain independent.
Drain tasks are named `jsonl-mutation:{thread_id}` for asyncio task dumps. Multi-thread
`put_batch` drains its current group on cancellation and never starts later groups;
the admitted group keeps its records on success or completes rollback on failure.
This is a store-local guarantee, not a change to RunJournal cancellation policy or
JSONL's single-process deployment constraint. Regression coverage is in
`tests/test_jsonl_event_store_cancellation.py`.

View File

@ -30,6 +30,7 @@ import asyncio
import json
import logging
import re
from collections.abc import Callable, Coroutine
from datetime import UTC, datetime
from pathlib import Path
from typing import Any
@ -54,6 +55,35 @@ class JsonlRunEventStore(RunEventStore):
def _get_write_lock(self, thread_id: str) -> asyncio.Lock:
return self._write_locks.setdefault(thread_id, asyncio.Lock())
async def _run_mutation[T](self, thread_id: str, operation: Callable[[], Coroutine[Any, Any, T]]) -> T:
"""Drain an admitted mutation before propagating caller cancellation.
Cancelling ``to_thread`` only stops its awaiter, not the filesystem
worker. Keep the thread lock through I/O, rollback and bookkeeping,
even if the caller is cancelled repeatedly. Queued callers can still
cancel before acquiring the lock, without starting a mutation.
"""
async with self._get_write_lock(thread_id):
task = asyncio.create_task(operation(), name=f"jsonl-mutation:{thread_id}")
cancellation: asyncio.CancelledError | None = None
while not task.done():
try:
await asyncio.shield(task)
except asyncio.CancelledError as exc:
if cancellation is None:
cancellation = exc
except Exception:
# Retrieve the failure below, after preserving any earlier
# cancellation. The operation has already finished rollback.
break
if cancellation is not None:
try:
task.result()
except Exception as exc:
raise cancellation from exc
raise cancellation
return task.result()
@staticmethod
def _validate_id(value: str, label: str) -> str:
"""Validate that an ID is safe for use in filesystem paths."""
@ -145,7 +175,7 @@ class JsonlRunEventStore(RunEventStore):
path.unlink()
async def put(self, *, thread_id, run_id, event_type, category, content="", metadata=None, created_at=None):
async with self._get_write_lock(thread_id):
async def mutate():
await self._ensure_seq_loaded(thread_id)
seq = self._next_seq(thread_id)
record = {
@ -161,6 +191,8 @@ class JsonlRunEventStore(RunEventStore):
await asyncio.to_thread(self._write_record, record)
return record
return await self._run_mutation(thread_id, mutate)
async def put_batch(self, events):
"""Persist a batch of events under a per-thread write lock.
@ -171,8 +203,9 @@ class JsonlRunEventStore(RunEventStore):
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.
back earlier thread groups. Cancellation drains the current thread group
before propagating, without starting subsequent groups. This rollback
does not make a multi-file batch crash-atomic.
"""
if not events:
return []
@ -199,7 +232,7 @@ class JsonlRunEventStore(RunEventStore):
metadata=None,
created_at=None,
):
async with self._get_write_lock(thread_id):
async def mutate():
existing = await asyncio.to_thread(self._read_run_events, thread_id, run_id)
for event in existing:
if event.get("event_type") == event_type:
@ -218,8 +251,10 @@ class JsonlRunEventStore(RunEventStore):
await asyncio.to_thread(self._write_record, record)
return record, True
return await self._run_mutation(thread_id, mutate)
async def _write_batch_async(self, thread_id: str, batch: list[dict[str, Any]]) -> list[dict[str, Any]]:
async with self._get_write_lock(thread_id):
async def mutate():
await self._ensure_seq_loaded(thread_id)
records: list[dict[str, Any]] = []
for ev in batch:
@ -242,6 +277,8 @@ class JsonlRunEventStore(RunEventStore):
await asyncio.to_thread(self._append_record_groups, run_batches)
return records
return await self._run_mutation(thread_id, mutate)
def _append_records(self, path: Path, records: list[dict[str, Any]]) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
lines = "".join(json.dumps(r, default=str, ensure_ascii=False) + "\n" for r in records)
@ -376,21 +413,25 @@ class JsonlRunEventStore(RunEventStore):
return found
async def delete_by_thread(self, thread_id):
async with self._get_write_lock(thread_id):
async def mutate():
all_events = await asyncio.to_thread(self._read_thread_events, thread_id)
count = len(all_events)
await asyncio.to_thread(self._delete_thread_files, thread_id)
self._seq_counters.pop(thread_id, None)
# Pop the lock inside the held scope to minimise the window where a new caller
# Pop the lock inside the held mutation to minimise the window where a new caller
# could obtain a fresh lock while a waiting coroutine still holds the old one.
# Note: coroutines that already acquired a reference to this lock before the
# delete will still proceed after we release — this is an accepted narrow race.
self._write_locks.pop(thread_id, None)
return count
return await self._run_mutation(thread_id, mutate)
async def delete_by_run(self, thread_id, run_id):
async with self._get_write_lock(thread_id):
async def mutate():
events = await asyncio.to_thread(self._read_run_events, thread_id, run_id)
count = len(events)
await asyncio.to_thread(self._delete_run_file, thread_id, run_id)
return count
return await self._run_mutation(thread_id, mutate)

View File

@ -0,0 +1,254 @@
"""Disk mutations retain thread ownership after caller cancellation (#5438)."""
from __future__ import annotations
import asyncio
import threading
import pytest
from deerflow.runtime.events.store.jsonl import JsonlRunEventStore
def _event(run_id="r1", content="message"):
return {"thread_id": "t1", "run_id": run_id, "event_type": "message", "category": "message", "content": content}
class _PausedIO:
"""Pause real filesystem work at a known point, without timing-based races."""
def __init__(self, operation):
self.operation = operation
self.loop = asyncio.get_running_loop()
self.entered = asyncio.Event()
self.finished = asyncio.Event()
self.release = threading.Event()
def __call__(self, *args):
self.loop.call_soon_threadsafe(self.entered.set)
try:
if not self.release.wait(10):
raise TimeoutError("test did not release paused filesystem operation")
return self.operation(*args)
finally:
self.loop.call_soon_threadsafe(self.finished.set)
async def _checkpoint():
# Let cancellation delivery and already-ready lock waiters run. I/O completion
# is controlled by Events, not by elapsed time or a fixed sleep budget.
await asyncio.sleep(0)
await asyncio.sleep(0)
@pytest.mark.anyio
@pytest.mark.parametrize("method", ["put", "put_if_absent", "put_batch"])
@pytest.mark.parametrize("cancellations", [1, 3])
async def test_cancelled_write_cannot_recreate_deleted_records(tmp_path, monkeypatch, method, cancellations):
store = JsonlRunEventStore(tmp_path)
await store.put(**_event(content="baseline"))
io_method = "_append_record_groups" if method == "put_batch" else "_write_record"
paused = _PausedIO(getattr(store, io_method))
monkeypatch.setattr(store, io_method, paused)
operation = getattr(store, method)
pending = asyncio.create_task(operation([_event("r2")]) if method == "put_batch" else operation(**_event("r2")))
deletion = None
try:
await asyncio.wait_for(paused.entered.wait(), 5)
for _ in range(cancellations):
pending.cancel()
await _checkpoint()
deletion = asyncio.create_task(store.delete_by_thread("t1"))
await _checkpoint()
assert not pending.done(), "cancellation returned while disk mutation still owned the thread"
assert not deletion.done(), "deletion overtook the cancelled disk mutation"
paused.release.set()
with pytest.raises(asyncio.CancelledError):
await pending
assert await deletion == 2
assert await store.count_messages("t1") == 0
assert "t1" not in store._seq_counters
assert "t1" not in store._write_locks
finally:
paused.release.set()
await asyncio.gather(pending, *([deletion] if deletion is not None else []), return_exceptions=True)
await asyncio.wait_for(paused.finished.wait(), 5)
@pytest.mark.anyio
@pytest.mark.parametrize("cancel", [False, True])
async def test_old_batch_rollback_cannot_erase_later_acknowledged_write(tmp_path, monkeypatch, cancel):
store = JsonlRunEventStore(tmp_path)
await store.put(**_event(content="baseline"))
append = store._append_records
def fail_second_file(path, records):
raise OSError("injected append failure")
paused = _PausedIO(fail_second_file)
def append_with_failure(path, records):
if path.stem == "r2":
return paused(path, records)
return append(path, records)
monkeypatch.setattr(store, "_append_records", append_with_failure)
batch = asyncio.create_task(store.put_batch([_event(content="batch-a"), _event("r2", "batch-b")]))
writer = None
try:
await asyncio.wait_for(paused.entered.wait(), 5)
if cancel:
batch.cancel()
await _checkpoint()
batch.cancel()
await _checkpoint()
writer = asyncio.create_task(store.put(**_event(content="acknowledged-later")))
await _checkpoint()
assert not writer.done(), "a new writer entered before rollback settled"
paused.release.set()
with pytest.raises(asyncio.CancelledError if cancel else OSError) as caught:
await batch
if cancel:
assert isinstance(caught.value.__cause__, OSError)
saved = await writer
assert saved["content"] == "acknowledged-later"
assert [row["content"] for row in await store.list_messages("t1")] == ["baseline", "acknowledged-later"]
assert await store.list_events("t1", "r2") == []
finally:
paused.release.set()
await asyncio.gather(batch, *([writer] if writer is not None else []), return_exceptions=True)
await asyncio.wait_for(paused.finished.wait(), 5)
@pytest.mark.anyio
@pytest.mark.parametrize("method,io_method", [("delete_by_thread", "_delete_thread_files"), ("delete_by_run", "_delete_run_file")])
async def test_cancelled_delete_cannot_erase_later_write(tmp_path, monkeypatch, method, io_method):
store = JsonlRunEventStore(tmp_path)
await store.put(**_event(content="baseline"))
paused = _PausedIO(getattr(store, io_method))
monkeypatch.setattr(store, io_method, paused)
pending = asyncio.create_task(store.delete_by_thread("t1") if method == "delete_by_thread" else store.delete_by_run("t1", "r1"))
writer = None
try:
await asyncio.wait_for(paused.entered.wait(), 5)
pending.cancel()
await _checkpoint()
pending.cancel()
await _checkpoint()
writer = asyncio.create_task(store.put(**_event(content="later")))
await _checkpoint()
assert not pending.done()
assert not writer.done()
paused.release.set()
with pytest.raises(asyncio.CancelledError):
await pending
saved = await writer
assert saved["seq"] == (1 if method == "delete_by_thread" else 2)
assert [row["content"] for row in await store.list_messages("t1")] == ["later"]
finally:
paused.release.set()
await asyncio.gather(pending, *([writer] if writer is not None else []), return_exceptions=True)
await asyncio.wait_for(paused.finished.wait(), 5)
@pytest.mark.anyio
async def test_cancellation_while_waiting_for_lock_never_starts_write(tmp_path):
store = JsonlRunEventStore(tmp_path)
async with store._get_write_lock("t1"):
pending = asyncio.create_task(store.put(**_event()))
await _checkpoint()
pending.cancel()
with pytest.raises(asyncio.CancelledError):
await pending
assert await store.count_messages("t1") == 0
assert not list(tmp_path.rglob("*.jsonl"))
@pytest.mark.anyio
async def test_cancelled_idempotent_write_is_visible_to_retry_and_other_threads_progress(tmp_path, monkeypatch):
store = JsonlRunEventStore(tmp_path)
write = store._write_record
paused = _PausedIO(write)
def pause_first_thread(record):
return paused(record) if record["thread_id"] == "t1" else write(record)
monkeypatch.setattr(store, "_write_record", pause_first_thread)
pending = asyncio.create_task(store.put_if_absent(**_event()))
retry = None
try:
await asyncio.wait_for(paused.entered.wait(), 5)
pending.cancel()
await _checkpoint()
other = await asyncio.wait_for(store.put(**{**_event(), "thread_id": "t2"}), 5)
assert other["seq"] == 1
retry = asyncio.create_task(store.put_if_absent(**_event()))
await _checkpoint()
assert not retry.done()
paused.release.set()
with pytest.raises(asyncio.CancelledError):
await pending
record, inserted = await retry
assert not inserted
assert record["seq"] == 1
assert await store.count_messages("t1") == 1
finally:
paused.release.set()
await asyncio.gather(pending, *([retry] if retry is not None else []), return_exceptions=True)
await asyncio.wait_for(paused.finished.wait(), 5)
@pytest.mark.anyio
@pytest.mark.parametrize("fail_first_thread", [False, True])
async def test_cancelled_multithread_batch_drains_current_group_without_starting_next(tmp_path, monkeypatch, fail_first_thread):
store = JsonlRunEventStore(tmp_path)
await store.put(**_event(content="baseline"))
append = store._append_records
def finish_first_thread(path, records):
if fail_first_thread:
raise OSError("injected first-thread append failure")
return append(path, records)
paused = _PausedIO(finish_first_thread)
def append_with_pause(path, records):
if path.stem == "r2":
return paused(path, records)
return append(path, records)
monkeypatch.setattr(store, "_append_records", append_with_pause)
pending = asyncio.create_task(
store.put_batch(
[
_event(content="first-thread-a"),
_event("r2", "first-thread-b"),
{**_event(content="second-thread"), "thread_id": "t2"},
]
)
)
try:
await asyncio.wait_for(paused.entered.wait(), 5)
pending.cancel()
await _checkpoint()
pending.cancel()
await _checkpoint()
assert not pending.done(), "the current thread group must finish before cancellation propagates"
assert "t2" not in store._seq_counters, "the next thread group must not start"
paused.release.set()
with pytest.raises(asyncio.CancelledError) as caught:
await pending
if fail_first_thread:
assert isinstance(caught.value.__cause__, OSError)
assert [row["content"] for row in await store.list_messages("t1")] == ["baseline"]
assert await store.list_events("t1", "r2") == []
else:
assert caught.value.__cause__ is None
assert [row["content"] for row in await store.list_messages("t1")] == ["baseline", "first-thread-a", "first-thread-b"]
assert await store.list_messages("t2") == []
assert not store._run_file("t2", "r1").exists()
assert "t2" not in store._seq_counters
finally:
paused.release.set()
await asyncio.gather(pending, return_exceptions=True)
await asyncio.wait_for(paused.finished.wait(), 5)