mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
fix(runtime): own RunJournal writes and quarantine ambiguous outcomes (#6846)
* test(runtime): pin RunJournal write-ownership and cancellation safety * fix(runtime): own started journal writes and quarantine unknown outcomes * docs(runtime): document RunJournal write-ownership contract * fix(runtime): address RunJournal review feedback * fix(runtime): log successful settlement of abandoned journal writes * fix(runtime): serialise explicit RunJournal flush and close ownership * test(runtime): assert concurrent RunJournal close and flush success paths * fix(runtime): make RunJournal close abort-aware and terminal-monotonic --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
spud
GitHub
parent
3195ed1f09
commit
e7a2a1c331
@ -124,6 +124,8 @@ prompts, and derived hashes do not enter the event. Ordinary task-tool subagents
|
||||
recorder proxy, never the journal itself; durable batch runs have no parent
|
||||
journal and emit no such event. Recorder failures are fail-open.
|
||||
|
||||
**RunJournal write ownership** (`runtime/journal.py`): a batch handed to `store.put_batch()` is owned by a journal task until its outcome settles, so a cancelled `flush()`/`close()` caller detaches the waiter and re-raises `CancelledError` only after the owned write is drained — never abandon or replay the in-flight batch. A generic `put_batch` exception is an ambiguous outcome: it is quarantined, blocks every successor write, and makes `flush()`/`close()` report non-success instead of replaying the batch. Only `RunEventStoreNotCommittedError` means the store proved none of the batch committed and may be deliberately retried; stores must raise it instead of a bare exception for that case. Explicit `flush()` and `close(flush=True)` share one per-journal operation lock, so a successor batch can never overtake an unsettled predecessor, every waiter observes the same settlement, and a normal close detaches only after the whole buffer has settled. `close(flush=False)` is a separate abort path that never queues on that lock: it revokes write ownership before its first await and makes any in-flight flush report non-success rather than a late success. A cancelled write waits a bounded time for the store call to end; `Task.cancelled()`/a clean return does not prove an uncooperative backend's I/O actually stopped, so this is a bound on waiting, not a termination guarantee. An inner store task that outlives that bound stays supervised until it settles, so its late result is retrieved and logged rather than dropped. Two adjacent limits are known and tracked separately, not fixed here: the bounded wait covers the store write, not a `progress_reporter` that ignores cancellation (a pre-existing hazard on the force-close path); and background-write warnings log the exception type and message, whose text is backend-controlled and may echo driver-level parameters, so log consumers must not treat it as payload-free.
|
||||
|
||||
**JSONL record boundaries** (`runtime/events/store/jsonl.py`): thread reads,
|
||||
run reads, and sequence recovery split on physical newlines. Do not use
|
||||
`str.splitlines()`: U+0085/U+2028/U+2029 inside valid JSON strings must remain
|
||||
|
||||
@ -1,4 +1,4 @@
|
||||
from deerflow.runtime.events.store.base import RunEventStore
|
||||
from deerflow.runtime.events.store.base import RunEventStore, RunEventStoreNotCommittedError
|
||||
from deerflow.runtime.events.store.memory import MemoryRunEventStore
|
||||
|
||||
|
||||
@ -23,4 +23,4 @@ def make_run_event_store(config=None) -> RunEventStore:
|
||||
raise ValueError(f"Unknown run_events backend: {config.backend!r}")
|
||||
|
||||
|
||||
__all__ = ["MemoryRunEventStore", "RunEventStore", "make_run_event_store"]
|
||||
__all__ = ["MemoryRunEventStore", "RunEventStore", "RunEventStoreNotCommittedError", "make_run_event_store"]
|
||||
|
||||
@ -24,6 +24,19 @@ class IncompleteMessageRunLookupError(RuntimeError):
|
||||
"""Raised when a store cannot prove that a targeted lookup is complete."""
|
||||
|
||||
|
||||
class RunEventStoreNotCommittedError(RuntimeError):
|
||||
"""A batch write definitely did not commit any of its events.
|
||||
|
||||
``RunEventStore.put_batch()`` implementations may raise this only when they
|
||||
can prove that none of the batch was persisted (for example, the write was
|
||||
rejected before reaching storage). ``RunJournal`` may then deliberately
|
||||
retry the batch. Any other exception from ``put_batch`` is treated as an
|
||||
ambiguous/UNKNOWN outcome: the batch is quarantined and is never replayed,
|
||||
because a store that committed and then lost its acknowledgement would
|
||||
otherwise produce duplicate events.
|
||||
"""
|
||||
|
||||
|
||||
def normalize_message_ids(message_ids: set[str]) -> set[str]:
|
||||
"""Return the non-empty string IDs that can participate in a lookup."""
|
||||
return {message_id for message_id in message_ids if isinstance(message_id, str) and message_id}
|
||||
|
||||
@ -27,7 +27,7 @@ from copy import deepcopy
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime
|
||||
from functools import partial
|
||||
from typing import TYPE_CHECKING, Any, cast
|
||||
from typing import Any, cast
|
||||
from uuid import UUID
|
||||
|
||||
from langchain_core.callbacks import BaseCallbackHandler
|
||||
@ -47,15 +47,18 @@ from deerflow.runtime.events.catalog import (
|
||||
RUN_ERROR_EVENT,
|
||||
RUN_START_EVENT,
|
||||
)
|
||||
from deerflow.runtime.events.store.base import RunEventStore, RunEventStoreNotCommittedError
|
||||
from deerflow.utils.llm_text import strip_leading_think_blocks
|
||||
from deerflow.utils.messages import message_to_text, restore_original_human_message
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from deerflow.runtime.events.store.base import RunEventStore
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_LEGACY_SUMMARY_MESSAGE_NAME = "summary"
|
||||
# A cancelled owned write may not stop a store that ignores cancellation.
|
||||
# The wrapper waits this long for the inner call to end before returning, so
|
||||
# lease-loss close cannot hang forever on an uncooperative backend. This is a
|
||||
# bound on waiting, not proof the underlying I/O stopped.
|
||||
_WRITE_SETTLE_TIMEOUT_SECONDS = 5.0
|
||||
_PERSISTED_HIDDEN_HUMAN_INPUT_RESPONSE_SOURCES = frozenset({"ask_clarification", "sandbox_network"})
|
||||
_MAX_RUN_SKILL_SNAPSHOTS = 64
|
||||
|
||||
@ -334,8 +337,27 @@ class RunJournal(BaseCallbackHandler):
|
||||
# Write buffer
|
||||
self._buffer: list[dict] = []
|
||||
self._pending_llm_response: _PendingLlmResponse | None = None
|
||||
self._pending_flush_tasks: set[asyncio.Task[None]] = set()
|
||||
self._explicit_flush_in_progress = False
|
||||
# Journal-owned in-flight writes: task -> original batch. A live drain
|
||||
# classifies and re-raises; the done callback only handles a write that
|
||||
# settles with no drain waiting (so its exception is never unretrieved).
|
||||
self._owned_writes: dict[asyncio.Task[None], list[dict]] = {}
|
||||
self._draining_tasks: set[asyncio.Task[None]] = set()
|
||||
# Serialises explicit persistence operations (flush and
|
||||
# close(flush=True)) for this journal instance. close(flush=False) never
|
||||
# queues on it.
|
||||
self._operation_lock = asyncio.Lock()
|
||||
# Set by close(flush=False) before its first await. It revokes write
|
||||
# ownership so no flush/callback may start or revive a write.
|
||||
self._abandoned = False
|
||||
# Inner store tasks that outlived a bounded writer teardown (a backend
|
||||
# that ignored cancellation). Held strongly until they settle so their
|
||||
# late result is retrieved and logged rather than dropped. This is not a
|
||||
# termination guarantee: it only keeps settlement responsibility.
|
||||
self._abandoned_store_tasks: dict[asyncio.Task[None], int] = {}
|
||||
# Batches whose outcome is UNKNOWN (a generic store failure) must never
|
||||
# be replayed: the store may have committed before losing its ACK.
|
||||
# Keeping them here also fences every successor write for this run.
|
||||
self._quarantined_batches: list[tuple[list[dict], BaseException]] = []
|
||||
self._pending_progress_task: asyncio.Task[None] | None = None
|
||||
self._pending_progress_delayed = False
|
||||
self._progress_dirty = False
|
||||
@ -1060,45 +1082,214 @@ class RunJournal(BaseCallbackHandler):
|
||||
self._commit_pending_llm_response()
|
||||
if not self._buffer:
|
||||
return
|
||||
# Skip if a flush is already in flight — avoids concurrent writes
|
||||
# to the same SQLite file from multiple fire-and-forget tasks.
|
||||
if self._pending_flush_tasks or self._explicit_flush_in_progress:
|
||||
return
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
# No event loop — keep events in buffer for later async flush.
|
||||
# Skip while another owned write is in flight, a normal persistence
|
||||
# operation holds the operation lock, a batch is quarantined, or the
|
||||
# journal has been force-closed. This keeps the batch in the buffer
|
||||
# instead of starting an overlapping write.
|
||||
if self._owned_writes or self._operation_lock.locked() or self._quarantined_batches or self._abandoned:
|
||||
return
|
||||
batch = self._buffer.copy()
|
||||
self._buffer.clear()
|
||||
task = loop.create_task(self._flush_async(batch))
|
||||
self._pending_flush_tasks.add(task)
|
||||
task.add_done_callback(self._on_flush_done)
|
||||
|
||||
async def _flush_async(self, batch: list[dict]) -> None:
|
||||
try:
|
||||
store = self._store
|
||||
if store is None:
|
||||
return
|
||||
await store.put_batch(batch)
|
||||
self._feed_generation += 1
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Failed to flush %d events for run %s — returning to buffer",
|
||||
len(batch),
|
||||
self.run_id,
|
||||
exc_info=True,
|
||||
)
|
||||
# Return failed events to buffer for retry on next flush
|
||||
task = self._start_write(batch)
|
||||
if task is None:
|
||||
self._buffer = batch + self._buffer
|
||||
|
||||
def _on_flush_done(self, task: asyncio.Task) -> None:
|
||||
self._pending_flush_tasks.discard(task)
|
||||
if task.cancelled():
|
||||
def _start_write(self, batch: list[dict]) -> asyncio.Task[None] | None:
|
||||
"""Create a journal-owned write task, or None when no loop is running.
|
||||
|
||||
The caller must have already removed ``batch`` from ``_buffer``.
|
||||
"""
|
||||
if self._abandoned:
|
||||
return None
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
return None
|
||||
task: asyncio.Task[None] = loop.create_task(self._write_batch(batch))
|
||||
self._owned_writes[task] = batch
|
||||
task.add_done_callback(self._reap_settled_write)
|
||||
return task
|
||||
|
||||
async def _write_batch(self, batch: list[dict]) -> None:
|
||||
"""Run one owned ``put_batch`` isolated from caller cancellation.
|
||||
|
||||
A dedicated inner task performs the store call; the outer owned task
|
||||
only bookkeeps it. Cancelling a waiter never cancels the store task, and
|
||||
the store result is classified by the drain or the reap callback.
|
||||
"""
|
||||
store = self._store
|
||||
if store is None:
|
||||
return
|
||||
exc = task.exception()
|
||||
if exc:
|
||||
logger.warning("Journal flush task failed: %s", exc)
|
||||
store_task: asyncio.Task[None] = asyncio.ensure_future(store.put_batch(batch))
|
||||
try:
|
||||
await asyncio.shield(store_task)
|
||||
except asyncio.CancelledError:
|
||||
# The waiter (or an explicit lease-loss close) abandoned durability.
|
||||
# Ask the inner store task to stop, but bound the wait: Task.cancelled()
|
||||
# only reflects the Python task, not whether the underlying I/O stopped,
|
||||
# and a backend may swallow cancellation. We never claim the I/O was
|
||||
# force-terminated; when the store outlives the bound we hand it to an
|
||||
# explicit supervisor so its late result is still observed.
|
||||
store_task.cancel()
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
asyncio.shield(store_task),
|
||||
timeout=_WRITE_SETTLE_TIMEOUT_SECONDS,
|
||||
)
|
||||
except (TimeoutError, asyncio.CancelledError):
|
||||
pass
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Journal store write for run %s failed while the writer was being torn down (%s: %s)",
|
||||
self.run_id,
|
||||
type(exc).__name__,
|
||||
exc,
|
||||
)
|
||||
if not store_task.done():
|
||||
self._supervise_abandoned_store_task(store_task, len(batch))
|
||||
raise
|
||||
|
||||
def _supervise_abandoned_store_task(self, store_task: asyncio.Task, batch_size: int) -> None:
|
||||
"""Hold and settle an inner store task that outlived writer teardown.
|
||||
|
||||
The task is kept strongly referenced until it settles, so it is neither
|
||||
garbage-collected nor silently dropped, and its eventual result is
|
||||
retrieved and logged. Bounded teardown is not a claim that the
|
||||
underlying I/O has stopped.
|
||||
"""
|
||||
self._abandoned_store_tasks[store_task] = batch_size
|
||||
store_task.add_done_callback(self._on_abandoned_store_task_settled)
|
||||
|
||||
def _on_abandoned_store_task_settled(self, store_task: asyncio.Task) -> None:
|
||||
batch_size = self._abandoned_store_tasks.pop(store_task, 0)
|
||||
if store_task.cancelled():
|
||||
return
|
||||
exc = store_task.exception()
|
||||
if exc is not None:
|
||||
logger.warning(
|
||||
"Journal store write for run %s failed after the writer was abandoned (%d event(s), %s: %s)",
|
||||
self.run_id,
|
||||
batch_size,
|
||||
type(exc).__name__,
|
||||
exc,
|
||||
)
|
||||
else:
|
||||
# Observability only: a late successful return is reported but does
|
||||
# NOT clear quarantine, replay events, release successor writes, or
|
||||
# make close(flush=False) report durability. "completed successfully"
|
||||
# describes put_batch() returning, not an independent commit proof.
|
||||
logger.warning(
|
||||
"Journal store write for run %s completed successfully after writer abandonment (%d event(s)); quarantine policy remains unchanged",
|
||||
self.run_id,
|
||||
batch_size,
|
||||
)
|
||||
|
||||
async def _drain_pending_writes(self) -> None:
|
||||
"""Settle every journal-owned write before any successor may start.
|
||||
|
||||
Each owned task is awaited through ``asyncio.shield`` so a cancelled
|
||||
caller detaches the waiter, never the durable write. Classification uses
|
||||
the ``(task, batch)`` snapshot taken here, so a concurrent force close
|
||||
clearing the registry can never turn a real outcome into a silent
|
||||
success.
|
||||
"""
|
||||
if not self._owned_writes:
|
||||
return
|
||||
pending = list(self._owned_writes.items())
|
||||
tasks = [task for task, _batch in pending]
|
||||
self._draining_tasks.update(tasks)
|
||||
try:
|
||||
result = await asyncio.gather(
|
||||
*(asyncio.shield(task) for task in tasks),
|
||||
return_exceptions=True,
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
self._draining_tasks.difference_update(tasks)
|
||||
raise
|
||||
failures: list[BaseException] = []
|
||||
first_not_committed: BaseException | None = None
|
||||
for (task, batch), outcome in zip(pending, result):
|
||||
self._draining_tasks.discard(task)
|
||||
self._owned_writes.pop(task, None)
|
||||
if self._abandoned:
|
||||
# Ownership was revoked by a force close: never re-buffer,
|
||||
# quarantine, or count a successor. Report non-success instead.
|
||||
failures.append(RuntimeError(f"RunJournal for run {self.run_id} was force-closed during flush"))
|
||||
continue
|
||||
if isinstance(outcome, asyncio.CancelledError):
|
||||
if task.done():
|
||||
self._quarantined_batches.append((batch, RuntimeError("journal write outcome is unknown after cancellation")))
|
||||
failures.append(outcome)
|
||||
else:
|
||||
# The caller's cancellation is still propagating; put the
|
||||
# task back so the done callback settles it later.
|
||||
self._owned_writes[task] = batch
|
||||
elif isinstance(outcome, RunEventStoreNotCommittedError):
|
||||
self._buffer = batch + self._buffer
|
||||
if first_not_committed is None:
|
||||
first_not_committed = outcome
|
||||
elif isinstance(outcome, BaseException):
|
||||
self._quarantined_batches.append((batch, outcome))
|
||||
failures.append(outcome)
|
||||
else:
|
||||
self._feed_generation += 1
|
||||
if failures:
|
||||
raise failures[0]
|
||||
if first_not_committed is not None:
|
||||
raise first_not_committed
|
||||
|
||||
def _reap_settled_write(self, task: asyncio.Task) -> None:
|
||||
"""Settle a write whose task finished with no drain waiting on it.
|
||||
|
||||
This runs on every completion but is a no-op for the normal drained
|
||||
path (the drain removes the task first). It exists so a write that
|
||||
settles after the caller detached does not leave an unretrieved task
|
||||
exception or an unquarantined result.
|
||||
"""
|
||||
if task in self._draining_tasks:
|
||||
return
|
||||
try:
|
||||
batch = self._owned_writes.pop(task, None)
|
||||
if batch is None:
|
||||
return
|
||||
if task.cancelled():
|
||||
self._quarantined_batches.append((batch, RuntimeError("journal write outcome is unknown after cancellation")))
|
||||
logger.warning(
|
||||
"Journal write for run %s was cancelled before its outcome was known (%d event(s)); quarantined and not retried",
|
||||
self.run_id,
|
||||
len(batch),
|
||||
)
|
||||
return
|
||||
exc = task.exception()
|
||||
if exc is None:
|
||||
self._feed_generation += 1
|
||||
elif isinstance(exc, RunEventStoreNotCommittedError):
|
||||
if self._abandoned:
|
||||
logger.warning(
|
||||
"Journal write for run %s came back not-committed after a force close (%d event(s)); not re-buffered",
|
||||
self.run_id,
|
||||
len(batch),
|
||||
)
|
||||
return
|
||||
self._buffer = batch + self._buffer
|
||||
logger.warning(
|
||||
"Journal write for run %s came back not-committed (%d event(s), %s: %s); re-buffered for a later flush",
|
||||
self.run_id,
|
||||
len(batch),
|
||||
type(exc).__name__,
|
||||
exc,
|
||||
)
|
||||
else:
|
||||
self._quarantined_batches.append((batch, exc))
|
||||
logger.warning(
|
||||
"Journal write for run %s has an unknown outcome (%d event(s), %s: %s); quarantined and not retried",
|
||||
self.run_id,
|
||||
len(batch),
|
||||
type(exc).__name__,
|
||||
exc,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Failed to reap a settled journal write for run %s", self.run_id)
|
||||
|
||||
def _identify_caller(self, tags: list[str] | None) -> str:
|
||||
_tags = tags or []
|
||||
@ -1352,54 +1543,104 @@ class RunJournal(BaseCallbackHandler):
|
||||
|
||||
async def flush(self) -> None:
|
||||
"""Force flush remaining buffer. Called in worker's finally block."""
|
||||
if self._abandoned:
|
||||
raise RuntimeError(f"RunJournal for run {self.run_id} was force-closed; flush cannot report success")
|
||||
if self._closed:
|
||||
return
|
||||
async with self._operation_lock:
|
||||
if self._abandoned:
|
||||
raise RuntimeError(f"RunJournal for run {self.run_id} was force-closed; flush cannot report success")
|
||||
if self._closed:
|
||||
return
|
||||
await self._flush_locked()
|
||||
|
||||
async def _flush_locked(self) -> None:
|
||||
"""Drain the buffer to quiescence while the operation lock is held.
|
||||
|
||||
Callers MUST hold ``self._operation_lock``. It deliberately does not
|
||||
re-acquire the lock, so ``close(flush=True)`` can run flush and detach as
|
||||
one indivisible protected operation. Returns only when every accepted
|
||||
event has settled and nothing is quarantined; raises otherwise.
|
||||
"""
|
||||
if self._abandoned:
|
||||
raise RuntimeError(f"RunJournal for run {self.run_id} was force-closed; flush cannot report success")
|
||||
# Events recorded from worker threads reach this loop through
|
||||
# call_soon_threadsafe (see record_middleware). Since Python 3.13 an
|
||||
# awaited run_in_executor()/to_thread() may complete without yielding
|
||||
# to the loop when the worker finished before its future was chained,
|
||||
# so such callbacks can still be queued when flush() starts. Yield once
|
||||
# so they land in the buffer instead of being dropped after detach.
|
||||
# awaited run_in_executor()/to_thread() may complete without yielding to
|
||||
# the loop when the worker finished before its future was chained, so
|
||||
# such callbacks can still be queued when the flush starts. Yield once
|
||||
# so they land in the buffer instead of being dropped after detach. This
|
||||
# lives here so close(flush=True), which calls this directly under the
|
||||
# operation lock, gets the same guarantee.
|
||||
await asyncio.sleep(0)
|
||||
if self._closed:
|
||||
return
|
||||
self._explicit_flush_in_progress = True
|
||||
if self._abandoned:
|
||||
raise RuntimeError(f"RunJournal for run {self.run_id} was force-closed; flush cannot report success")
|
||||
if self._quarantined_batches:
|
||||
_batch, error = self._quarantined_batches[0]
|
||||
raise RuntimeError(f"RunJournal flush refused: {len(_batch)} event(s) for run {self.run_id} have an unknown durable outcome") from error
|
||||
self._commit_pending_llm_response()
|
||||
try:
|
||||
self._commit_pending_llm_response()
|
||||
if self._pending_flush_tasks:
|
||||
await asyncio.gather(*tuple(self._pending_flush_tasks), return_exceptions=True)
|
||||
while self._pending_progress_task is not None:
|
||||
pending_progress_task = self._pending_progress_task
|
||||
if pending_progress_task.done():
|
||||
if self._pending_progress_task is pending_progress_task:
|
||||
self._pending_progress_task = None
|
||||
break
|
||||
if self._pending_progress_delayed:
|
||||
pending_progress_task.cancel()
|
||||
await asyncio.gather(pending_progress_task, return_exceptions=True)
|
||||
if self._pending_progress_task is pending_progress_task:
|
||||
self._pending_progress_task = None
|
||||
self._progress_dirty = False
|
||||
self._pending_progress_delayed = False
|
||||
break
|
||||
await self._drain_pending_writes()
|
||||
except asyncio.CancelledError:
|
||||
# A caller cancelled a flush that already started a durable write.
|
||||
# Let the journal-owned write settle before propagating the
|
||||
# cancellation so the batch can never be lost between the task and
|
||||
# the buffer. Do NOT call uncancel(): we are propagating the
|
||||
# cancellation, not suppressing it, and touching the task's cancel
|
||||
# count would corrupt the caller's cancellation accounting.
|
||||
try:
|
||||
await self._drain_pending_writes()
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Journal write failed while flushing after cancellation for run %s",
|
||||
self.run_id,
|
||||
exc_info=True,
|
||||
)
|
||||
raise
|
||||
while self._pending_progress_task is not None:
|
||||
pending_progress_task = self._pending_progress_task
|
||||
if pending_progress_task.done():
|
||||
if self._pending_progress_task is pending_progress_task:
|
||||
self._pending_progress_task = None
|
||||
break
|
||||
if self._pending_progress_delayed:
|
||||
pending_progress_task.cancel()
|
||||
await asyncio.gather(pending_progress_task, return_exceptions=True)
|
||||
if self._pending_progress_task is pending_progress_task:
|
||||
self._pending_progress_task = None
|
||||
self._progress_dirty = False
|
||||
self._pending_progress_delayed = False
|
||||
break
|
||||
await asyncio.gather(pending_progress_task, return_exceptions=True)
|
||||
if self._pending_progress_task is pending_progress_task:
|
||||
self._pending_progress_task = None
|
||||
|
||||
while self._buffer:
|
||||
batch = self._buffer[: self._flush_threshold]
|
||||
del self._buffer[: self._flush_threshold]
|
||||
try:
|
||||
store = self._store
|
||||
if store is None:
|
||||
return
|
||||
await store.put_batch(batch)
|
||||
self._feed_generation += 1
|
||||
except Exception:
|
||||
self._buffer = batch + self._buffer
|
||||
raise
|
||||
finally:
|
||||
self._explicit_flush_in_progress = False
|
||||
# Only batches already accepted before this flush are attempted; a
|
||||
# NOT_COMMITTED batch re-buffered by the drain waits for a later explicit
|
||||
# flush instead of being retried in a tight loop. A force close revokes
|
||||
# ownership, so no further batch may start.
|
||||
while self._buffer:
|
||||
if self._abandoned or self._quarantined_batches:
|
||||
break
|
||||
batch = self._buffer[: self._flush_threshold]
|
||||
del self._buffer[: self._flush_threshold]
|
||||
if self._store is None or self._abandoned:
|
||||
self._buffer = batch + self._buffer
|
||||
break
|
||||
task = self._start_write(batch)
|
||||
if task is None:
|
||||
# No running loop (or write ownership was revoked): keep the batch.
|
||||
self._buffer = batch + self._buffer
|
||||
break
|
||||
await self._drain_pending_writes()
|
||||
|
||||
if self._abandoned:
|
||||
raise RuntimeError(f"RunJournal for run {self.run_id} was force-closed during flush")
|
||||
if self._quarantined_batches:
|
||||
_batch, error = self._quarantined_batches[0]
|
||||
raise RuntimeError(f"RunJournal flush refused: {len(_batch)} event(s) for run {self.run_id} have an unknown durable outcome") from error
|
||||
|
||||
def _detach_runtime_dependencies(self) -> None:
|
||||
"""Drop every external or potentially cyclic run-scoped reference."""
|
||||
@ -1410,8 +1651,10 @@ class RunJournal(BaseCallbackHandler):
|
||||
self._progress_reporter = None
|
||||
self._buffer.clear()
|
||||
self._pending_llm_response = None
|
||||
self._pending_flush_tasks.clear()
|
||||
self._explicit_flush_in_progress = False
|
||||
self._owned_writes = {}
|
||||
self._draining_tasks = set()
|
||||
# Intentionally NOT cleared: abandoned inner store tasks stay strongly
|
||||
# referenced here until they settle and their result is logged.
|
||||
self._pending_progress_task = None
|
||||
self._pending_progress_delayed = False
|
||||
self._progress_dirty = False
|
||||
@ -1432,29 +1675,46 @@ class RunJournal(BaseCallbackHandler):
|
||||
|
||||
async def close(self, *, flush: bool = True) -> None:
|
||||
"""Release run-scoped references, optionally flushing buffered events."""
|
||||
if self._closed:
|
||||
return
|
||||
if flush:
|
||||
# A failed terminal write returns its batch to ``_buffer``. Keep the
|
||||
# store and all buffered state attached so a later close/flush can retry
|
||||
# instead of silently discarding the tail of the run event stream.
|
||||
await self.flush()
|
||||
self._detach_runtime_dependencies()
|
||||
# A normal close is one protected operation: it serialises with every
|
||||
# other flush / close(True) and detaches only after the whole buffer
|
||||
# has settled, so no concurrent waiter can observe a
|
||||
# settled-but-detached journal or start a successor batch that nobody
|
||||
# drains.
|
||||
if self._abandoned:
|
||||
raise RuntimeError(f"RunJournal for run {self.run_id} was force-closed; close cannot report success")
|
||||
async with self._operation_lock:
|
||||
# ABORT is checked before the closed flag: a force close sets
|
||||
# both, so a normal close that queued behind the lock must fail
|
||||
# rather than observe `_closed` and report a false success.
|
||||
if self._abandoned:
|
||||
raise RuntimeError(f"RunJournal for run {self.run_id} was force-closed; close cannot report success")
|
||||
if self._closed:
|
||||
return
|
||||
await self._flush_locked()
|
||||
self._detach_runtime_dependencies()
|
||||
return
|
||||
|
||||
# A worker that lost its lease must detach without starting another
|
||||
# durable write. Drop dependencies before cancelling already-scheduled
|
||||
# work so tasks that have not begun observe the detached state. The
|
||||
# final detach must survive a second cancellation while those tasks stop.
|
||||
# Force / lease-loss close. It must NOT queue behind a possibly-stuck
|
||||
# normal operation, so it does not take the operation lock. It marks
|
||||
# ABORT before its first await, then cancels and reaps every owned write
|
||||
# (bounded, with late-result supervision) before detaching.
|
||||
# Terminal-state monotonicity: a journal that already closed durably must
|
||||
# not be re-marked ABORT (which would flip flush() from a no-op to an
|
||||
# error). A force close still preempts an in-flight normal close, where
|
||||
# `_closed` is not yet set.
|
||||
if self._closed:
|
||||
return
|
||||
self._abandoned = True
|
||||
self._closed = True
|
||||
self._store = None
|
||||
self._progress_reporter = None
|
||||
try:
|
||||
pending_flush_tasks = tuple(self._pending_flush_tasks)
|
||||
for task in pending_flush_tasks:
|
||||
pending_owned = tuple(self._owned_writes)
|
||||
for task in pending_owned:
|
||||
task.cancel()
|
||||
if pending_flush_tasks:
|
||||
await asyncio.gather(*pending_flush_tasks, return_exceptions=True)
|
||||
if pending_owned:
|
||||
await asyncio.gather(*pending_owned, return_exceptions=True)
|
||||
pending_progress_task = self._pending_progress_task
|
||||
if pending_progress_task is not None:
|
||||
pending_progress_task.cancel()
|
||||
|
||||
@ -4,6 +4,7 @@ Uses MemoryRunEventStore as the backend for direct event inspection.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import threading
|
||||
import weakref
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
@ -15,6 +16,7 @@ import pytest
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langchain_core.outputs import ChatGeneration, LLMResult
|
||||
|
||||
from deerflow.runtime.events.store.base import RunEventStoreNotCommittedError
|
||||
from deerflow.runtime.events.store.memory import MemoryRunEventStore
|
||||
from deerflow.runtime.journal import RunJournal
|
||||
from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY
|
||||
@ -227,7 +229,6 @@ async def test_close_flushes_and_detaches_runtime_dependencies():
|
||||
assert journal._store is None
|
||||
assert journal._progress_reporter is None
|
||||
assert journal._buffer == []
|
||||
assert journal._pending_flush_tasks == set()
|
||||
del store, reporter
|
||||
await asyncio.sleep(0)
|
||||
assert store_ref() is None
|
||||
@ -262,14 +263,14 @@ async def test_close_preserves_buffer_and_dependencies_when_flush_fails():
|
||||
async def put_batch(self, events):
|
||||
self.put_batch_calls += 1
|
||||
if self.put_batch_calls == 1:
|
||||
raise RuntimeError("transient store failure")
|
||||
raise RunEventStoreNotCommittedError("transient store failure")
|
||||
return await super().put_batch(events)
|
||||
|
||||
store = FailOnceRunEventStore()
|
||||
journal = RunJournal("r-close-retry", "t-close-retry", store, flush_threshold=100)
|
||||
journal.record_middleware("test", name="test", hook="after", action="record", changes={})
|
||||
|
||||
with pytest.raises(RuntimeError, match="transient store failure"):
|
||||
with pytest.raises(RunEventStoreNotCommittedError, match="transient store failure"):
|
||||
await journal.close()
|
||||
|
||||
assert journal._closed is False
|
||||
@ -295,7 +296,7 @@ async def test_close_retries_pending_no_usage_response_without_duplication():
|
||||
async def put_batch(self, events):
|
||||
self.put_batch_calls += 1
|
||||
if self.put_batch_calls == 1:
|
||||
raise RuntimeError("transient store failure")
|
||||
raise RunEventStoreNotCommittedError("transient store failure")
|
||||
return await super().put_batch(events)
|
||||
|
||||
async def progress_reporter(snapshot):
|
||||
@ -320,7 +321,7 @@ async def test_close_retries_pending_no_usage_response_without_duplication():
|
||||
assert journal._pending_llm_response is not None
|
||||
assert journal.get_completion_data()["message_count"] == 0
|
||||
|
||||
with pytest.raises(RuntimeError, match="transient store failure"):
|
||||
with pytest.raises(RunEventStoreNotCommittedError, match="transient store failure"):
|
||||
await journal.close()
|
||||
|
||||
assert journal._closed is False
|
||||
@ -390,7 +391,7 @@ async def test_close_without_flush_detaches_when_cancellation_interrupts_pending
|
||||
await asyncio.Event().wait()
|
||||
|
||||
pending_flush = asyncio.create_task(stubborn_pending_flush())
|
||||
journal._pending_flush_tasks.add(pending_flush)
|
||||
journal._owned_writes[pending_flush] = [{"event_type": "test", "category": "trace", "content": ""}]
|
||||
close_task = asyncio.create_task(journal.close(flush=False))
|
||||
await asyncio.wait_for(first_cancellation_seen.wait(), timeout=1)
|
||||
|
||||
@ -402,7 +403,443 @@ async def test_close_without_flush_detaches_when_cancellation_interrupts_pending
|
||||
assert journal._closed is True
|
||||
assert journal._store is None
|
||||
assert journal._buffer == []
|
||||
assert journal._pending_flush_tasks == set()
|
||||
assert journal._owned_writes == {}
|
||||
|
||||
|
||||
class TestWriteOwnership:
|
||||
@pytest.mark.anyio
|
||||
async def test_close_without_flush_cancels_owned_store_task(self):
|
||||
"""A hung store write must be cancelled and reaped on lease-loss close."""
|
||||
started = asyncio.Event()
|
||||
store_task_ref: list[asyncio.Task] = []
|
||||
|
||||
class HungStore(MemoryRunEventStore):
|
||||
async def put_batch(self, events):
|
||||
started.set()
|
||||
store_task_ref.append(asyncio.current_task())
|
||||
await asyncio.Event().wait()
|
||||
|
||||
store = HungStore()
|
||||
journal = RunJournal("r-hung", "t-hung", store, flush_threshold=100)
|
||||
journal._put(event_type="A", category="trace", content="one")
|
||||
journal._flush_sync()
|
||||
|
||||
await started.wait()
|
||||
await journal.close(flush=False)
|
||||
|
||||
assert journal._owned_writes == {}
|
||||
assert store_task_ref and store_task_ref[0].cancelled(), "inner store task was orphaned"
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_close_without_flush_returns_when_store_ignores_cancellation(self, monkeypatch):
|
||||
"""A store that ignores cancellation must not hang lease-loss close.
|
||||
|
||||
``Task.cancelled()`` cannot prove the underlying I/O stopped, so the
|
||||
journal bounds its wait instead of claiming the write was terminated.
|
||||
"""
|
||||
import deerflow.runtime.journal as journal_module
|
||||
|
||||
monkeypatch.setattr(journal_module, "_WRITE_SETTLE_TIMEOUT_SECONDS", 0.05, raising=False)
|
||||
started = asyncio.Event()
|
||||
never = asyncio.Event()
|
||||
|
||||
class UncooperativeStore(MemoryRunEventStore):
|
||||
async def put_batch(self, events):
|
||||
started.set()
|
||||
try:
|
||||
await never.wait()
|
||||
except asyncio.CancelledError:
|
||||
# Swallow cancellation and keep the I/O "running".
|
||||
await never.wait()
|
||||
|
||||
journal = RunJournal("r-uncooperative", "t-uncooperative", UncooperativeStore(), flush_threshold=100)
|
||||
journal._put(event_type="A", category="trace", content="one")
|
||||
journal._flush_sync()
|
||||
|
||||
await started.wait()
|
||||
await asyncio.wait_for(journal.close(flush=False), timeout=2.0)
|
||||
assert journal._owned_writes == {}
|
||||
assert journal._closed is True
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_committed_batch_with_lost_ack_is_not_replayed(self):
|
||||
class CommitThenRaiseStore(MemoryRunEventStore):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.calls = 0
|
||||
|
||||
async def put_batch(self, events):
|
||||
self.calls += 1
|
||||
await super().put_batch(events)
|
||||
raise ConnectionError("write committed but ACK lost")
|
||||
|
||||
store = CommitThenRaiseStore()
|
||||
journal = RunJournal("r-lost-ack", "t-lost-ack", store, flush_threshold=100)
|
||||
journal._put(event_type="A", category="trace", content="one")
|
||||
|
||||
with pytest.raises(Exception):
|
||||
await journal.flush()
|
||||
# The runner retries on its next flush; an ambiguous failure must not
|
||||
# replay a batch the store already committed.
|
||||
for _ in range(2):
|
||||
with pytest.raises(Exception):
|
||||
await journal.flush()
|
||||
|
||||
events = await store.list_events("t-lost-ack", "r-lost-ack")
|
||||
seqs = [event["seq"] for event in events]
|
||||
assert store.calls == 1, f"committed batch replayed {store.calls - 1} extra time(s)"
|
||||
assert len(seqs) == len(set(seqs)), "duplicate seq created by replay"
|
||||
assert [event["event_type"] for event in events] == ["A"]
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_background_completion_after_caller_cancel_is_reaped(self, caplog):
|
||||
class FailingStore(MemoryRunEventStore):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.started = asyncio.Event()
|
||||
self.finish = asyncio.Event()
|
||||
|
||||
async def put_batch(self, events):
|
||||
self.started.set()
|
||||
await self.finish.wait()
|
||||
raise RuntimeError("store failed after caller detached")
|
||||
|
||||
store = FailingStore()
|
||||
journal = RunJournal("r-bg-reap", "t-bg-reap", store, flush_threshold=100)
|
||||
journal._put(event_type="A", category="trace", content="one")
|
||||
|
||||
flush_task = asyncio.create_task(journal.flush())
|
||||
await store.started.wait()
|
||||
flush_task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await flush_task
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="deerflow.runtime.journal"):
|
||||
store.finish.set()
|
||||
for _ in range(5):
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert journal._owned_writes == {}, "background write was not reaped"
|
||||
assert len(journal._quarantined_batches) == 1, "background failure was lost"
|
||||
assert "has an unknown outcome" in caplog.text
|
||||
assert "store failed after caller detached" in caplog.text
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_caller_cancellation_does_not_cancel_owned_write(self):
|
||||
class HangingStore(MemoryRunEventStore):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.started = asyncio.Event()
|
||||
self.finish = asyncio.Event()
|
||||
self.calls = 0
|
||||
|
||||
async def put_batch(self, events):
|
||||
self.calls += 1
|
||||
self.started.set()
|
||||
await self.finish.wait()
|
||||
return await super().put_batch(events)
|
||||
|
||||
store = HangingStore()
|
||||
journal = RunJournal("r-caller-cancel", "t-caller-cancel", store, flush_threshold=100)
|
||||
journal._put(event_type="A", category="trace", content="one")
|
||||
|
||||
flush_task = asyncio.create_task(journal.flush())
|
||||
await store.started.wait()
|
||||
flush_task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(flush_task, timeout=1.0)
|
||||
|
||||
# The owned write is still outstanding and completes exactly once.
|
||||
store.finish.set()
|
||||
await asyncio.sleep(0)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
events = await store.list_events("t-caller-cancel", "r-caller-cancel")
|
||||
types = [event["event_type"] for event in events]
|
||||
assert store.calls == 1, "cancelled write was retried (possible duplicate)"
|
||||
assert types.count("A") == 1, f"A persisted {types.count('A')} times: {types}"
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_repeated_cancellation_and_concurrent_close_settle_write_once(self):
|
||||
class HangingStore(MemoryRunEventStore):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.started = asyncio.Event()
|
||||
self.finish = asyncio.Event()
|
||||
self.calls = 0
|
||||
|
||||
async def put_batch(self, events):
|
||||
self.calls += 1
|
||||
self.started.set()
|
||||
await self.finish.wait()
|
||||
return await super().put_batch(events)
|
||||
|
||||
store = HangingStore()
|
||||
journal = RunJournal("r-concurrent-close", "t-concurrent-close", store, flush_threshold=100)
|
||||
journal._put(event_type="A", category="trace", content="first")
|
||||
|
||||
caller1 = asyncio.create_task(journal.close())
|
||||
await store.started.wait()
|
||||
caller1.cancel()
|
||||
caller1.cancel()
|
||||
await asyncio.sleep(0)
|
||||
|
||||
caller2 = asyncio.create_task(journal.close())
|
||||
await asyncio.sleep(0)
|
||||
store.finish.set()
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await caller1
|
||||
await asyncio.wait_for(caller2, timeout=1.0)
|
||||
|
||||
events = await store.list_events("t-concurrent-close", "r-concurrent-close")
|
||||
types = [event["event_type"] for event in events]
|
||||
assert store.calls == 1, "concurrent close retried the write"
|
||||
assert types.count("A") == 1, f"A persisted {types.count('A')} times: {types}"
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_not_committed_batch_is_retried_once_per_flush(self):
|
||||
class NeverCommittingStore(MemoryRunEventStore):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.calls = 0
|
||||
|
||||
async def put_batch(self, events):
|
||||
self.calls += 1
|
||||
raise RunEventStoreNotCommittedError("store rejected the batch")
|
||||
|
||||
store = NeverCommittingStore()
|
||||
journal = RunJournal("r-no-busy-loop", "t-no-busy-loop", store, flush_threshold=100)
|
||||
journal._put(event_type="A", category="trace", content="one")
|
||||
|
||||
with pytest.raises(RunEventStoreNotCommittedError):
|
||||
await journal.flush()
|
||||
# One attempt for this flush; the batch is re-buffered for a later one.
|
||||
assert store.calls == 1, f"flush busy-looped on a not-committed batch: {store.calls} calls"
|
||||
assert [event["event_type"] for event in journal._buffer] == ["A"]
|
||||
|
||||
with pytest.raises(RunEventStoreNotCommittedError):
|
||||
await journal.flush()
|
||||
assert store.calls == 2
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_background_not_committed_write_is_logged(self, caplog):
|
||||
class DetachedNotCommittedStore(MemoryRunEventStore):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.started = asyncio.Event()
|
||||
self.finish = asyncio.Event()
|
||||
|
||||
async def put_batch(self, events):
|
||||
self.started.set()
|
||||
await self.finish.wait()
|
||||
raise RunEventStoreNotCommittedError("rejected before storage")
|
||||
|
||||
store = DetachedNotCommittedStore()
|
||||
journal = RunJournal("r-bg-nc", "t-bg-nc", store, flush_threshold=100)
|
||||
journal._put(event_type="A", category="trace", content="one")
|
||||
|
||||
flush_task = asyncio.create_task(journal.flush())
|
||||
await store.started.wait()
|
||||
flush_task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await flush_task
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="deerflow.runtime.journal"):
|
||||
store.finish.set()
|
||||
for _ in range(5):
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert journal._owned_writes == {}
|
||||
assert [event["event_type"] for event in journal._buffer] == ["A"]
|
||||
assert "came back not-committed" in caplog.text
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sync_write_failure_is_not_double_logged_by_reaper(self, caplog):
|
||||
class FailingStore(MemoryRunEventStore):
|
||||
async def put_batch(self, events):
|
||||
raise RuntimeError("store unavailable")
|
||||
|
||||
journal = RunJournal("r-sync-log", "t-sync-log", FailingStore(), flush_threshold=100)
|
||||
journal._put(event_type="A", category="trace", content="one")
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="deerflow.runtime.journal"):
|
||||
with pytest.raises(RuntimeError, match="store unavailable"):
|
||||
await journal.flush()
|
||||
|
||||
# The drain surfaces the failure to the caller; the reaper must stay quiet
|
||||
# so the same failure is not logged twice.
|
||||
assert "has an unknown outcome" not in caplog.text
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize("cancel_times", [1, 2])
|
||||
async def test_flush_cancellation_preserves_task_cancel_count(self, cancel_times):
|
||||
"""Cancelling an explicit flush must not mutate the caller's cancel count.
|
||||
|
||||
The write is started by a threshold flush first, so the explicit flush
|
||||
enters its already-in-flight drain (the branch that used to call
|
||||
``uncancel()``).
|
||||
"""
|
||||
|
||||
class HangingStore(MemoryRunEventStore):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.started = asyncio.Event()
|
||||
self.finish = asyncio.Event()
|
||||
|
||||
async def put_batch(self, events):
|
||||
self.started.set()
|
||||
await self.finish.wait()
|
||||
return await super().put_batch(events)
|
||||
|
||||
store = HangingStore()
|
||||
journal = RunJournal("r-cancel-count", "t-cancel-count", store, flush_threshold=100)
|
||||
journal._put(event_type="A", category="trace", content="one")
|
||||
journal._flush_sync() # start the owned write before the explicit flush
|
||||
await store.started.wait()
|
||||
|
||||
flush_task = asyncio.create_task(journal.flush())
|
||||
# flush() yields once before touching the buffer; pump until it is parked
|
||||
# on the already-in-flight write so the cancel lands in the settle branch.
|
||||
for _ in range(20):
|
||||
await asyncio.sleep(0)
|
||||
if journal._draining_tasks:
|
||||
break
|
||||
assert journal._draining_tasks, "flush did not enter its in-flight drain"
|
||||
|
||||
for _ in range(cancel_times):
|
||||
flush_task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
store.finish.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await flush_task
|
||||
|
||||
assert flush_task.cancelling() == cancel_times, "journal must not mutate the caller's cancel count"
|
||||
for _ in range(20):
|
||||
await asyncio.sleep(0)
|
||||
types = [event["event_type"] for event in await store.list_events("t-cancel-count", "r-cancel-count")]
|
||||
assert types.count("A") == 1
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_close_without_flush_supervises_abandoned_store_task(self, monkeypatch, caplog):
|
||||
"""A store that outlives bounded teardown must still have its result observed."""
|
||||
import gc
|
||||
|
||||
import deerflow.runtime.journal as journal_module
|
||||
|
||||
monkeypatch.setattr(journal_module, "_WRITE_SETTLE_TIMEOUT_SECONDS", 0.02, raising=False)
|
||||
started = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
|
||||
class UncooperativeStore(MemoryRunEventStore):
|
||||
async def put_batch(self, events):
|
||||
started.set()
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
except asyncio.CancelledError:
|
||||
# Swallow the teardown cancel and keep the "I/O" running.
|
||||
await release.wait()
|
||||
raise RuntimeError("late failure after the writer was abandoned")
|
||||
|
||||
journal = RunJournal("r-abandoned", "t-abandoned", UncooperativeStore(), flush_threshold=100)
|
||||
journal._put(event_type="A", category="trace", content="one")
|
||||
journal._flush_sync()
|
||||
await started.wait()
|
||||
|
||||
unretrieved: list[BaseException] = []
|
||||
loop = asyncio.get_running_loop()
|
||||
previous = loop.get_exception_handler()
|
||||
loop.set_exception_handler(lambda _loop, ctx: unretrieved.append(ctx.get("exception")))
|
||||
try:
|
||||
with caplog.at_level(logging.WARNING, logger="deerflow.runtime.journal"):
|
||||
await journal.close(flush=False)
|
||||
assert journal._closed is True
|
||||
assert len(journal._abandoned_store_tasks) == 1, "inner store task was not supervised"
|
||||
|
||||
release.set()
|
||||
for _ in range(10):
|
||||
await asyncio.sleep(0)
|
||||
assert journal._abandoned_store_tasks == {}
|
||||
gc.collect()
|
||||
await asyncio.sleep(0)
|
||||
finally:
|
||||
loop.set_exception_handler(previous)
|
||||
|
||||
assert unretrieved == [], f"abandoned store task exception was never retrieved: {unretrieved}"
|
||||
assert "after the writer was abandoned" in caplog.text
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_abandoned_store_success_is_logged_without_clearing_quarantine(self, monkeypatch, caplog):
|
||||
"""A late successful store return must be observable but change nothing else."""
|
||||
import deerflow.runtime.journal as journal_module
|
||||
|
||||
monkeypatch.setattr(journal_module, "_WRITE_SETTLE_TIMEOUT_SECONDS", 0.02, raising=False)
|
||||
started = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
|
||||
class LateSuccessStore(MemoryRunEventStore):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.put_batch_calls = 0
|
||||
|
||||
async def put_batch(self, events):
|
||||
self.put_batch_calls += 1
|
||||
started.set()
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
except asyncio.CancelledError:
|
||||
# Swallow the teardown cancel, keep the "I/O" running, then
|
||||
# complete successfully once released.
|
||||
await release.wait()
|
||||
return await super().put_batch(events)
|
||||
|
||||
store = LateSuccessStore()
|
||||
journal = RunJournal("r-late-ok", "t-late-ok", store, flush_threshold=100)
|
||||
journal._put(event_type="A", category="trace", content="one")
|
||||
journal._flush_sync()
|
||||
await started.wait()
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="deerflow.runtime.journal"):
|
||||
await asyncio.wait_for(journal.close(flush=False), timeout=2.0)
|
||||
assert journal._closed is True
|
||||
assert len(journal._abandoned_store_tasks) == 1, "inner store task was not supervised"
|
||||
assert len(journal._quarantined_batches) == 1, "writer teardown should have quarantined the batch"
|
||||
|
||||
release.set()
|
||||
for _ in range(10):
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert journal._abandoned_store_tasks == {}
|
||||
assert store.put_batch_calls == 1
|
||||
# Observability only: the late success must not clear quarantine, replay,
|
||||
# release successors, or make close(flush=False) report durability.
|
||||
assert len(journal._quarantined_batches) == 1, "late success must not clear quarantine"
|
||||
events = await store.list_events("t-late-ok", "r-late-ok")
|
||||
assert [event["event_type"] for event in events] == ["A"]
|
||||
assert "completed successfully after writer abandonment" in caplog.text
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_store_self_cancellation_is_quarantined_not_replayed(self):
|
||||
class SelfCancellingStore(MemoryRunEventStore):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.calls = 0
|
||||
|
||||
async def put_batch(self, events):
|
||||
self.calls += 1
|
||||
raise asyncio.CancelledError("store cancelled its own write")
|
||||
|
||||
store = SelfCancellingStore()
|
||||
journal = RunJournal("r-self-cancel", "t-self-cancel", store, flush_threshold=100)
|
||||
journal._put(event_type="A", category="trace", content="one")
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await journal.flush()
|
||||
# UNKNOWN: no replay, and no successor may overtake it.
|
||||
with pytest.raises(Exception):
|
||||
await journal.flush()
|
||||
assert store.calls == 1, f"self-cancelled store write replayed {store.calls - 1} time(s)"
|
||||
assert await store.list_events("t-self-cancel", "r-self-cancel") == []
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@ -454,6 +891,330 @@ def _combine_llm_responses(*responses):
|
||||
return response
|
||||
|
||||
|
||||
async def _pump_loop(ticks: int = 30) -> None:
|
||||
for _ in range(ticks):
|
||||
await asyncio.sleep(0)
|
||||
|
||||
|
||||
def _returned_normally(task: asyncio.Task) -> bool:
|
||||
if not task.done() or task.cancelled():
|
||||
return False
|
||||
return task.exception() is None
|
||||
|
||||
|
||||
class _GatedRunEventStore(MemoryRunEventStore):
|
||||
"""Records ``put_batch`` call order and blocks chosen event types on a gate."""
|
||||
|
||||
def __init__(self, gates: dict[str, asyncio.Event], failures: dict[str, BaseException] | None = None):
|
||||
super().__init__()
|
||||
self._gates = gates
|
||||
self._failures = failures or {}
|
||||
self.calls: list[str] = []
|
||||
|
||||
async def put_batch(self, events):
|
||||
event_type = events[0]["event_type"]
|
||||
self.calls.append(event_type)
|
||||
gate = self._gates.get(event_type)
|
||||
if gate is not None:
|
||||
await gate.wait()
|
||||
failure = self._failures.get(event_type)
|
||||
if failure is not None:
|
||||
raise failure
|
||||
return await super().put_batch(events)
|
||||
|
||||
|
||||
class TestOperationOwnership:
|
||||
"""Normal flush/close are one serialised operation; force close preempts it.
|
||||
|
||||
Explicit flush and close(flush=True) share an operation lock so a successor
|
||||
batch can never overtake an unsettled predecessor, every waiter observes the
|
||||
same settlement, and a normal close only detaches after the whole buffer has
|
||||
settled. close(flush=False) is a separate abort path that must not queue
|
||||
behind a stuck flush and must never let an in-flight flush report success.
|
||||
"""
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_concurrent_close_waits_for_successor_batch(self):
|
||||
gate_a = asyncio.Event()
|
||||
gate_b = asyncio.Event()
|
||||
store = _GatedRunEventStore({"A": gate_a, "B": gate_b})
|
||||
journal = RunJournal("r-cclose", "t-cclose", store, flush_threshold=1)
|
||||
|
||||
journal._put(event_type="A", category="trace", content="a")
|
||||
await _pump_loop()
|
||||
assert store.calls == ["A"]
|
||||
journal._put(event_type="B", category="trace", content="b")
|
||||
|
||||
close1 = asyncio.create_task(journal.close())
|
||||
close2 = asyncio.create_task(journal.close())
|
||||
await _pump_loop()
|
||||
|
||||
gate_a.set()
|
||||
await _pump_loop()
|
||||
|
||||
early = [i for i, task in enumerate((close1, close2)) if _returned_normally(task)]
|
||||
assert not early, f"close(s) {early} returned before successor batch settled; calls={store.calls}"
|
||||
|
||||
gate_b.set()
|
||||
results = await asyncio.gather(close1, close2)
|
||||
assert results == [None, None], "both normal closes must complete successfully"
|
||||
assert store.calls == ["A", "B"]
|
||||
events = await store.list_events("t-cclose", "r-cclose")
|
||||
assert [event["event_type"] for event in events] == ["A", "B"]
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_concurrent_flushes_do_not_overtake_unsettled_predecessor(self):
|
||||
gate_a = asyncio.Event()
|
||||
gate_b = asyncio.Event()
|
||||
store = _GatedRunEventStore({"A": gate_a, "B": gate_b})
|
||||
journal = RunJournal("r-cflush", "t-cflush", store, flush_threshold=1)
|
||||
|
||||
journal._put(event_type="A", category="trace", content="a")
|
||||
await _pump_loop()
|
||||
assert store.calls == ["A"]
|
||||
journal._put(event_type="B", category="trace", content="b")
|
||||
journal._put(event_type="C", category="trace", content="c")
|
||||
|
||||
f1 = asyncio.create_task(journal.flush())
|
||||
f2 = asyncio.create_task(journal.flush())
|
||||
await _pump_loop()
|
||||
|
||||
gate_a.set()
|
||||
await _pump_loop()
|
||||
|
||||
assert "C" not in store.calls, f"successor C overtook unsettled predecessor: calls={store.calls}"
|
||||
|
||||
gate_b.set()
|
||||
results = await asyncio.gather(f1, f2)
|
||||
assert results == [None, None], "both concurrent flushes must complete successfully"
|
||||
assert store.calls == ["A", "B", "C"], f"unexpected final submission order: {store.calls}"
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_concurrent_flushes_all_report_unknown(self):
|
||||
gate_a = asyncio.Event()
|
||||
store = _GatedRunEventStore({"A": gate_a}, failures={"A": RuntimeError("unknown outcome")})
|
||||
journal = RunJournal("r-cunknown", "t-cunknown", store, flush_threshold=1)
|
||||
|
||||
journal._put(event_type="A", category="trace", content="a")
|
||||
await _pump_loop()
|
||||
assert store.calls == ["A"]
|
||||
|
||||
f1 = asyncio.create_task(journal.flush())
|
||||
f2 = asyncio.create_task(journal.flush())
|
||||
await _pump_loop()
|
||||
gate_a.set()
|
||||
await asyncio.gather(f1, f2, return_exceptions=True)
|
||||
|
||||
reported = [not task.cancelled() and task.exception() is not None for task in (f1, f2)]
|
||||
assert all(reported), f"some waiter did not report the UNKNOWN outcome: f1={f1!r} f2={f2!r}"
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_finishing_one_flush_does_not_release_operation_guard(self):
|
||||
gate_a = asyncio.Event()
|
||||
gate_b = asyncio.Event()
|
||||
store = _GatedRunEventStore({"A": gate_a, "B": gate_b})
|
||||
journal = RunJournal("r-cguard", "t-cguard", store, flush_threshold=1)
|
||||
|
||||
journal._put(event_type="A", category="trace", content="a")
|
||||
await _pump_loop()
|
||||
journal._put(event_type="B", category="trace", content="b")
|
||||
|
||||
f1 = asyncio.create_task(journal.flush())
|
||||
f2 = asyncio.create_task(journal.flush())
|
||||
await _pump_loop()
|
||||
gate_a.set()
|
||||
await _pump_loop()
|
||||
|
||||
assert journal._operation_lock.locked() is True, "operation guard was released while another flush was active"
|
||||
|
||||
gate_b.set()
|
||||
results = await asyncio.gather(f1, f2)
|
||||
assert results == [None, None], "both concurrent flushes must succeed once the predecessor settles"
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_lease_loss_close_aborts_active_flush_without_success(self):
|
||||
gate_a = asyncio.Event()
|
||||
store = _GatedRunEventStore({"A": gate_a})
|
||||
journal = RunJournal("r-lease", "t-lease", store, flush_threshold=1)
|
||||
|
||||
journal._put(event_type="A", category="trace", content="a")
|
||||
await _pump_loop()
|
||||
|
||||
flush_task = asyncio.create_task(journal.flush())
|
||||
await _pump_loop()
|
||||
await journal.close(flush=False)
|
||||
await _pump_loop()
|
||||
|
||||
assert not _returned_normally(flush_task), "active flush reported success after lease-loss close"
|
||||
await asyncio.gather(flush_task, return_exceptions=True)
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_force_close_preempts_normal_close_without_success(self):
|
||||
gate_a = asyncio.Event()
|
||||
store = _GatedRunEventStore({"A": gate_a})
|
||||
journal = RunJournal("r-mixclose", "t-mixclose", store, flush_threshold=1)
|
||||
|
||||
journal._put(event_type="A", category="trace", content="a")
|
||||
await _pump_loop()
|
||||
assert store.calls == ["A"]
|
||||
|
||||
normal_close = asyncio.create_task(journal.close())
|
||||
await _pump_loop()
|
||||
await journal.close(flush=False)
|
||||
await _pump_loop()
|
||||
|
||||
assert not _returned_normally(normal_close), "normal close reported success after a force close"
|
||||
await asyncio.gather(normal_close, return_exceptions=True)
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_queued_normal_close_after_force_close_does_not_report_success(self):
|
||||
"""A normal close already queueing on the lock must not succeed after abort."""
|
||||
gate_a = asyncio.Event()
|
||||
store = _GatedRunEventStore({"A": gate_a})
|
||||
journal = RunJournal("r-queued-close", "t-queued-close", store, flush_threshold=1)
|
||||
|
||||
journal._put(event_type="A", category="trace", content="a")
|
||||
await _pump_loop()
|
||||
assert store.calls == ["A"]
|
||||
|
||||
normal_close = asyncio.create_task(journal.close())
|
||||
await _pump_loop()
|
||||
assert journal._operation_lock.locked(), "first close did not acquire the operation lock"
|
||||
|
||||
queued_close = asyncio.create_task(journal.close())
|
||||
await _pump_loop()
|
||||
assert not queued_close.done(), "second close did not queue behind the lock"
|
||||
|
||||
# Force close preempts without taking the lock (sets _closed + _abandoned).
|
||||
await journal.close(flush=False)
|
||||
|
||||
gate_a.set()
|
||||
await _pump_loop()
|
||||
|
||||
assert not _returned_normally(queued_close), "queued normal close reported success after a force close"
|
||||
await asyncio.gather(normal_close, queued_close, return_exceptions=True)
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_force_close_after_durable_close_does_not_re_mark_abandoned(self):
|
||||
"""Terminal state is monotonic: DURABLY_CLOSED must not become ABANDONED."""
|
||||
store = MemoryRunEventStore()
|
||||
journal = RunJournal("r-mono", "t-mono", store, flush_threshold=100)
|
||||
journal._put(event_type="A", category="trace", content="a")
|
||||
await journal.close()
|
||||
|
||||
assert journal._closed is True
|
||||
assert journal._abandoned is False
|
||||
|
||||
await journal.close(flush=False)
|
||||
|
||||
assert journal._abandoned is False, "a durably closed journal must not be re-marked ABORT"
|
||||
await journal.flush() # stays a silent no-op, does not raise
|
||||
events = await store.list_events("t-mono", "r-mono")
|
||||
assert [event["event_type"] for event in events] == ["A"]
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_double_force_close_is_characterised(self, monkeypatch):
|
||||
"""Characterisation: the 2nd force close returns while the 1st still cleans.
|
||||
|
||||
Force close makes no durability promise, so "returned" means the abort was
|
||||
issued, not that local cleanup finished. This is documented rather than
|
||||
changed here; whether callers should join one cleanup is a contract call.
|
||||
"""
|
||||
import deerflow.runtime.journal as journal_module
|
||||
|
||||
monkeypatch.setattr(journal_module, "_WRITE_SETTLE_TIMEOUT_SECONDS", 0.05, raising=False)
|
||||
started = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
|
||||
class UncooperativeStore(MemoryRunEventStore):
|
||||
async def put_batch(self, events):
|
||||
started.set()
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
except asyncio.CancelledError:
|
||||
await release.wait()
|
||||
return await super().put_batch(events)
|
||||
|
||||
journal = RunJournal("r-2force", "t-2force", UncooperativeStore(), flush_threshold=1)
|
||||
journal._put(event_type="A", category="trace", content="a")
|
||||
await _pump_loop()
|
||||
await started.wait()
|
||||
|
||||
first = asyncio.create_task(journal.close(flush=False))
|
||||
await _pump_loop()
|
||||
second = asyncio.create_task(journal.close(flush=False))
|
||||
await _pump_loop()
|
||||
|
||||
assert second.done(), "second force close did not return promptly"
|
||||
assert not first.done(), "first force close already finished its bounded settle"
|
||||
|
||||
release.set()
|
||||
await asyncio.gather(first, second, return_exceptions=True)
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_force_close_progress_reporter_unbounded_is_characterised(self):
|
||||
"""Characterisation of a pre-existing limitation, tracked separately.
|
||||
|
||||
close(flush=False) bounds the store write (see the abandoned-store tests)
|
||||
but still awaits the pending progress task without a deadline, so a
|
||||
progress reporter that swallows cancellation can block a force close.
|
||||
This reproduces identically on the pre-PR base, so it is documented here
|
||||
rather than fixed inside the write-ownership PR.
|
||||
"""
|
||||
started = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
|
||||
async def reporter(snapshot):
|
||||
started.set()
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
except asyncio.CancelledError:
|
||||
await release.wait()
|
||||
|
||||
journal = RunJournal("r-prog", "t-prog", MemoryRunEventStore(), flush_threshold=100, progress_reporter=reporter)
|
||||
journal._schedule_progress_flush()
|
||||
await started.wait()
|
||||
assert journal._pending_progress_task is not None
|
||||
|
||||
bounded = True
|
||||
try:
|
||||
await asyncio.wait_for(journal.close(flush=False), timeout=0.25)
|
||||
except TimeoutError:
|
||||
bounded = False
|
||||
finally:
|
||||
release.set()
|
||||
await _pump_loop()
|
||||
|
||||
assert bounded is False, "force close became bounded; please update this characterisation"
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_concurrent_flush_not_committed_retry_is_characterised(self):
|
||||
"""Contract characterisation: a concurrent flush must not duplicate A."""
|
||||
gate_a = asyncio.Event()
|
||||
store = _GatedRunEventStore({"A": gate_a}, failures={"A": RunEventStoreNotCommittedError("not committed")})
|
||||
journal = RunJournal("r-cnc", "t-cnc", store, flush_threshold=1)
|
||||
|
||||
journal._put(event_type="A", category="trace", content="a")
|
||||
await _pump_loop()
|
||||
assert store.calls == ["A"]
|
||||
|
||||
f1 = asyncio.create_task(journal.flush())
|
||||
f2 = asyncio.create_task(journal.flush())
|
||||
await _pump_loop()
|
||||
gate_a.set()
|
||||
await asyncio.gather(f1, f2, return_exceptions=True)
|
||||
await _pump_loop()
|
||||
|
||||
# Durable-safety invariants only. The exact number of definite retries
|
||||
# (and which waiter owns them) under concurrency is still an open
|
||||
# maintainer contract decision, so it is deliberately not asserted here.
|
||||
assert store.calls.count("A") >= 1, "A was never attempted"
|
||||
surfaced = [task for task in (f1, f2) if not task.cancelled() and task.exception() is not None]
|
||||
assert surfaced, "the not-committed failure was not surfaced to any waiter"
|
||||
events = await store.list_events("t-cnc", "r-cnc")
|
||||
assert sum(1 for event in events if event["event_type"] == "A") == 0, "A must not be persisted after a not-committed failure"
|
||||
|
||||
|
||||
class TestLlmCallbacks:
|
||||
@pytest.mark.anyio
|
||||
async def test_on_chat_model_start_persists_original_user_input_without_mutating_model_message(self, journal_setup):
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user