github/deer-flow
1
0
Fork 0
mirror of https://github.com/bytedance/deer-flow.git synced 2026-10-29 22:50:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(runtime): isolate run admission from unrelated heartbeats (#7664)

Browse Source 
Signed-off-by: LingYi-Liang <270316896+LingYi-Liang@users.noreply.github.com>
Co-authored-by: LingYi-Liang <270316896+LingYi-Liang@users.noreply.github.com>
This commit is contained in:
LingYi 2026-10-29 22:19:10 +08:00 committed by GitHub
parent 44b88151be
commit 30aa601f6a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 494 additions and 122 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
README.md
View File

@ -694,7 +694,11 @@ DeerFlow still uses `Forwarded` / `X-Forwarded-*` headers to recover the browser
> The producer-lease recovery behavior above is gated by `stream_bridge.producer_lease_recovery_enabled`. Keep it disabled while old Gateway binaries or their runs remain; enable it on every instance only after the writer rollout and drain complete.
>
> With run ownership heartbeats enabled, terminal success requires acknowledgment
> from that worker's own durable write. A matching database row alone cannot
> from that worker's own durable write. Run creation and replacement serialize
> per conversation: a slow admission does not hold the manager-wide lock needed
> by another conversation's cached run reads or heartbeat snapshots. Their own
> Store I/O can still encounter database or pool contention. New runs remain
> hidden until their initial Store write succeeds. A matching database row alone cannot
> confirm success, and completion metadata cannot retry a rejected terminal
> transition. Each run renews independently, so a stalled store operation does
> not delay other runs. If acknowledgment remains unavailable when the last

19
backend/docs/runtime-guidance-details.md
View File

@ -333,6 +333,25 @@ Checkpointer storage runs in one of two channel modes, selected by `checkpoint_c
## Deferred terminal commits and lease ownership
RunManager serializes creation, replacement and assembly publication per thread.
The global registry lock never spans Store I/O; unrelated cached reads and
heartbeat snapshots do not queue behind an admission gate. Their own Store I/O
can still encounter database or connection-pool contention. A pending-ID map
lets readers of preallocated run IDs join the same-thread gate before the record
is published. Failed or cancelled admissions remove that temporary membership;
normal creation registers only after its initial write succeeds. Heartbeats use
the short registry snapshot, not the thread gate, so an admitted run can renew
while its replacement awaits a Store result. Preserve terminal-attribution and
cancellation compensation below; changing lock granularity does not grant a new
worker authority over a durable row.
Assembly evidence persists a full row snapshot. Serialize that write with the
same run's lease renewal and local acknowledgment, so it cannot restore an old
lease after renewal. Lock order is thread gate, then snapshot gate, then the
short registry lock. Renewal uses only the snapshot gate and must release it
before signaling local cancellation, which acquires the thread gate. Keep the
lease-expiry timer armed while waiting for the snapshot gate.
**Deferred terminal commit:** With an event store, the worker stages its terminal
status locally and commits it only after finalization's receipt and duration
writes. `RunRecord.terminal_commit_pending` keeps `_renew_leases()` renewing that

7
backend/packages/harness/deerflow/runtime/AGENTS.md
View File

@ -13,6 +13,13 @@ errors preserve cache fallbacks and valid neighbors, including Redis outages/pur
Escape literal stems for SCAN; recheck byte/string keys before UNLINK. Preserve
namespace/thread bounds, cursor iteration and outage/TTL fallback.
### Run Admission
Run admission and assembly publication serialize per thread across Store I/O.
Keep registry locks short; heartbeat snapshots must not queue on an
admission gate. Preserve same-thread publication visibility, including reads of
preallocated run IDs. Details: [admission isolation](../../../../docs/runtime-guidance-details.md#deferred-terminal-commits-and-lease-ownership).
### Stream Bridge Heartbeats
Observed Redis streams removed during replay yield `StreamGap` with empty bounds; initially empty streams still wait. After this gap, `/wait` refreshes store-only status and completes only for a terminal run with no retained stream.

290
backend/packages/harness/deerflow/runtime/runs/manager.py
View File

@ -10,7 +10,7 @@ import socket
import sqlite3
import uuid
from collections.abc import AsyncIterator, Awaitable, Callable
from contextlib import asynccontextmanager
from contextlib import AsyncExitStack, asynccontextmanager
from dataclasses import dataclass, field
from datetime import UTC, datetime, timedelta
from enum import StrEnum
@ -19,6 +19,7 @@ from typing import TYPE_CHECKING, Any
from sqlalchemy.exc import IntegrityError as SAIntegrityError
from deerflow.runtime.events.terminal import RecoveryTerminalPublisher
from deerflow.runtime.keyed_lock import AsyncKeyedLockTable
from deerflow.runtime.user_context import AUTO, _AutoSentinel, get_current_user, resolve_user_id
from deerflow.utils.time import is_lease_expired
from deerflow.utils.time import now_iso as _now_iso
@ -353,6 +354,9 @@ class RunManager:
# preserving ``_runs`` iteration order (see ``_thread_records_locked``).
self._runs_by_thread: dict[str, dict[str, None]] = {}
self._lock = asyncio.Lock()
self._admission_locks = AsyncKeyedLockTable[str]()
self._snapshot_locks = AsyncKeyedLockTable[str]()
self._pending_admissions: dict[str, dict[str, int]] = {}
self._store = store
self._persistence_retry_policy = persistence_retry_policy or PersistenceRetryPolicy()
self._worker_id = worker_id or _generate_worker_id()
@ -370,6 +374,43 @@ class RunManager:
self._lease_operations: dict[str, asyncio.Task] = {}
self._lease_deadlines: dict[str, asyncio.TimerHandle] = {}
@asynccontextmanager
async def _admission_scope(self, thread_id: str, run_id: str) -> AsyncIterator[None]:
# Manager state is owned by one event loop. Register before the first
# await so readers of a preallocated ID also wait for its publication.
threads = self._pending_admissions.setdefault(run_id, {})
threads[thread_id] = threads.get(thread_id, 0) + 1
try:
async with self._admission_locks.hold(thread_id):
yield
finally:
threads[thread_id] -= 1
if not threads[thread_id]:
del threads[thread_id]
if not threads:
self._pending_admissions.pop(run_id, None)
@asynccontextmanager
async def _run_scope(self, run_id: str) -> AsyncIterator[None]:
threads = set(self._pending_admissions.get(run_id, ()))
record = self._runs.get(run_id)
if record is not None:
threads.add(record.thread_id)
async with AsyncExitStack() as stack:
for thread_id in sorted(threads):
await stack.enter_async_context(self._admission_locks.hold(thread_id))
yield
@asynccontextmanager
async def _run_registry_lock(self, run_id: str) -> AsyncIterator[None]:
async with self._run_scope(run_id), self._lock:
yield
@asynccontextmanager
async def _thread_registry_lock(self, thread_id: str) -> AsyncIterator[None]:
async with self._admission_locks.hold(thread_id), self._lock:
yield
def _index_run_locked(self, record: RunRecord) -> None:
"""Register *record* in the thread index. Caller must hold ``self._lock``."""
self._runs_by_thread.setdefault(record.thread_id, {})[record.run_id] = None
@ -500,8 +541,8 @@ class RunManager:
Initial run creation is part of the run visibility boundary: callers
should not observe a run in memory unless its backing store row exists.
Unlike follow-up status/model updates, failures are propagated so the
caller can treat creation as failed. Rollback is the caller's
responsibility after inserting the record into ``_runs``.
caller can treat creation as failed. The caller keeps the record private
until this method returns successfully.
"""
if self._store is None:
return
@ -793,7 +834,7 @@ class RunManager:
"""Persist token usage and completion data to the backing store."""
row_recovery_payload: dict[str, Any] | None = None
record: RunRecord | None = None
async with self._lock:
async with self._run_registry_lock(run_id):
record = self._runs.get(run_id)
if record is not None and record.ownership_lost:
logger.warning("Skipped completion persistence for run %s after lease ownership was lost", run_id)
@ -889,7 +930,7 @@ class RunManager:
async def update_run_progress(self, run_id: str, **kwargs) -> None:
"""Persist a running token/message snapshot without changing status."""
should_persist = True
async with self._lock:
async with self._run_registry_lock(run_id):
record = self._runs.get(run_id)
if record is not None:
should_persist = record.status == RunStatus.running and not record.ownership_lost
@ -907,7 +948,7 @@ class RunManager:
async def update_finalizing_progress(self, run_id: str, **kwargs) -> None:
"""Persist final fields while the durable row is deliberately active."""
should_persist = False
async with self._lock:
async with self._run_registry_lock(run_id):
record = self._runs.get(run_id)
if record is not None and not record.ownership_lost:
should_persist = record.status not in (RunStatus.pending, RunStatus.running)
@ -931,7 +972,7 @@ class RunManager:
async def seal_completed_evidence(self, run_id: str, receipt) -> bool:
"""Best-effort evidence metadata is independent of business completion."""
async with self._lock:
async with self._run_registry_lock(run_id):
record = self._runs.get(run_id)
if record is None or record.ownership_lost:
return False
@ -942,7 +983,7 @@ class RunManager:
return await seal(run_id, owner_worker_id=owner_worker_id, receipt=receipt)
async def mark_evidence_partial(self, run_id: str, *, error: str) -> bool:
async with self._lock:
async with self._run_registry_lock(run_id):
record = self._runs.get(run_id)
if record is None or record.ownership_lost:
return False
@ -998,21 +1039,15 @@ class RunManager:
evidence_origin=evidence_origin,
evidence_agent_id=evidence_agent_id,
)
async with self._lock:
self._runs[run_id] = record
self._index_run_locked(record)
persisted = False
async with self._admission_scope(thread_id, run_id):
try:
await self._persist_new_run_to_store(record)
persisted = True
except Exception:
logger.warning("Failed to persist run %s; rolled back in-memory record", run_id, exc_info=True)
logger.warning("Failed to persist run %s; record was not registered", run_id, exc_info=True)
raise
finally:
# Also covers cancellation, which bypasses ``except Exception``.
if not persisted:
self._runs.pop(run_id, None)
self._unindex_run_locked(run_id, record.thread_id)
async with self._lock:
self._runs[run_id] = record
self._index_run_locked(record)
logger.info("Run created: run_id=%s thread_id=%s", run_id, thread_id)
return record
@ -1031,7 +1066,7 @@ class RunManager:
raise_on_store_error: Propagate store hydration/mapping failures so
lifecycle callers can distinguish them from a missing run.
"""
async with self._lock:
async with self._run_registry_lock(run_id):
record = self._runs.get(run_id)
if record is not None:
return record
@ -1046,7 +1081,7 @@ class RunManager:
return None
# Re-check after store await: a concurrent create() may have inserted the
# in-memory record while the store call was in flight.
async with self._lock:
async with self._run_registry_lock(run_id):
record = self._runs.get(run_id)
if record is not None:
return record
@ -1115,7 +1150,7 @@ class RunManager:
def _page(records: list[RunRecord]) -> list[RunRecord]:
return sorted(records, key=lambda record: run_sort_key(record.created_at, record.run_id), reverse=True)[:limit]
async with self._lock:
async with self._thread_registry_lock(thread_id):
memory_records = [
record
for record in self._thread_records_locked(thread_id)
@ -1169,7 +1204,7 @@ class RunManager:
correct pagination.
"""
resolved_user_id = resolve_user_id(user_id, method_name="RunManager.list_successful_regenerate_sources")
async with self._lock:
async with self._thread_registry_lock(thread_id):
memory_records = [record for record in self._thread_records_locked(thread_id) if record.operation_kind == ThreadOperationKind.run and (resolved_user_id is None or record.user_id == resolved_user_id)]
sources = set(await self._store.list_successful_regenerate_sources(thread_id, user_id=resolved_user_id)) if self._store is not None else set()
@ -1240,7 +1275,7 @@ class RunManager:
continue
records_by_id[record.run_id] = record
async with self._lock:
async with self._thread_registry_lock(thread_id):
memory_records = [record for record in self._thread_records_locked(thread_id) if resolved_user_id is None or record.user_id == resolved_user_id]
for record in memory_records:
records_by_id[record.run_id] = record
@ -1249,13 +1284,13 @@ class RunManager:
async def try_start(self, run_id: str) -> RunStartOutcome:
"""Transition an uncancelled pending run to running before building the agent."""
async with self._lock:
async with self._run_registry_lock(run_id):
record = self._runs.get(run_id)
if record is None:
raise RunStartupError(f"Cannot start unknown run {run_id}")
async with record.start_lock:
async with self._lock:
async with self._run_registry_lock(run_id):
if record.abort_event.is_set() or record.status != RunStatus.pending:
return RunStartOutcome.cancelled
@ -1269,7 +1304,7 @@ class RunManager:
except Exception as exc:
raise RunStartupError(f"Failed to start run {run_id}: {exc}") from exc
if updated is False:
async with self._lock:
async with self._run_registry_lock(run_id):
if record.status == RunStatus.pending:
record.status = RunStatus.interrupted
record.terminal_status_persisted = False
@ -1277,7 +1312,7 @@ class RunManager:
record.updated_at = _now_iso()
return RunStartOutcome.cancelled
async with self._lock:
async with self._run_registry_lock(run_id):
if record.abort_event.is_set() or record.status != RunStatus.pending:
restore_status = record.status
restore_error = record.error
@ -1299,7 +1334,7 @@ class RunManager:
async def fail_start_if_pending(self, run_id: str, *, error: str) -> bool:
"""Mark an admitted run as failed if its worker task could not be attached."""
async with self._lock:
async with self._run_registry_lock(run_id):
record = self._runs.get(run_id)
if record is None or record.status != RunStatus.pending:
return False
@ -1314,7 +1349,7 @@ class RunManager:
RunStatus.error,
error=error,
)
async with self._lock:
async with self._run_registry_lock(run_id):
if self._runs.get(run_id) is record:
record.terminal_status_persisted = persisted
return True
@ -1330,7 +1365,7 @@ class RunManager:
if not run_ids:
return {}
resolved_user_id = resolve_user_id(user_id, method_name="RunManager.get_many_by_thread")
async with self._lock:
async with self._thread_registry_lock(thread_id):
records_by_id = {
record.run_id: record for record in self._thread_records_locked(thread_id) if record.operation_kind == ThreadOperationKind.run and record.run_id in run_ids and (resolved_user_id is None or record.user_id == resolved_user_id)
}
@ -1365,7 +1400,7 @@ class RunManager:
persist: bool = True,
) -> None:
"""Transition a run to a new status."""
async with self._lock:
async with self._run_registry_lock(run_id):
record = self._runs.get(run_id)
if record is None:
logger.warning("set_status called for unknown run %s", run_id)
@ -1401,7 +1436,7 @@ class RunManager:
if persist:
persisted = await self._persist_status(record, status, error=error, stop_reason=stop_reason)
if status not in (RunStatus.pending, RunStatus.running):
async with self._lock:
async with self._run_registry_lock(run_id):
record.terminal_status_persisted = self._has_final_terminal_ack(record) if self.heartbeat_enabled and self._store is not None else persisted
if not persisted and self.heartbeat_enabled and status == RunStatus.success and not record.ownership_lost:
await self._mark_ownership_lost(
@ -1415,7 +1450,7 @@ class RunManager:
async def persist_current_status(self, run_id: str) -> bool:
"""Persist the status already staged on the in-memory run record."""
async with self._lock:
async with self._run_registry_lock(run_id):
record = self._runs.get(run_id)
if record is None:
logger.warning("persist_current_status called for unknown run %s", run_id)
@ -1425,7 +1460,7 @@ class RunManager:
stop_reason = record.stop_reason
persisted = await self._persist_status(record, status, error=error, stop_reason=stop_reason)
if status not in (RunStatus.pending, RunStatus.running):
async with self._lock:
async with self._run_registry_lock(run_id):
record.terminal_status_persisted = self._has_final_terminal_ack(record) if self.heartbeat_enabled and self._store is not None else persisted
if not persisted and self.heartbeat_enabled and status == RunStatus.success and not record.ownership_lost:
await self._mark_ownership_lost(
@ -1457,7 +1492,7 @@ class RunManager:
)
return None
async with self._lock:
async with self._run_registry_lock(run_id):
record = self._runs.get(run_id)
if record is None or record.store_only or record.ownership_lost:
return None
@ -1507,7 +1542,7 @@ class RunManager:
async def set_finalizing(self, run_id: str, finalizing: bool) -> None:
"""Mark whether a run is performing post-cancel cleanup."""
async with self._lock:
async with self._run_registry_lock(run_id):
record = self._runs.get(run_id)
if record is None:
logger.warning("set_finalizing called for unknown run %s", run_id)
@ -1525,7 +1560,7 @@ class RunManager:
) -> None:
"""Wait until older same-thread runs have finished post-cancel cleanup."""
while True:
async with self._lock:
async with self._thread_registry_lock(thread_id):
found_current = False
prior_finalizing = False
for record in self._thread_records_locked(thread_id):
@ -1549,7 +1584,7 @@ class RunManager:
async def has_later_run(self, thread_id: str, run_id: str) -> bool:
"""Return whether a newer in-memory run has been admitted for the thread."""
async with self._lock:
async with self._thread_registry_lock(thread_id):
seen_current = False
for record in self._thread_records_locked(thread_id):
if record.run_id == run_id:
@ -1561,7 +1596,7 @@ class RunManager:
async def has_later_started_run(self, thread_id: str, run_id: str) -> bool:
"""Return whether a newer same-thread run may have already advanced state."""
async with self._lock:
async with self._thread_registry_lock(thread_id):
seen_current = False
for record in self._thread_records_locked(thread_id):
if record.run_id == run_id:
@ -1588,22 +1623,23 @@ class RunManager:
"""Publish only host assembly evidence, once, before execution."""
from deerflow.runtime.assembly_evidence import EVIDENCE_KEY
async with self._lock:
record = self._runs.get(run_id)
if record is None or evidence.get("run_id") != run_id or evidence.get("thread_id") != record.thread_id:
raise ValueError("Assembly/run binding mismatch")
if EVIDENCE_KEY in record.metadata:
if record.metadata[EVIDENCE_KEY] != evidence:
raise ValueError("Run assembly evidence is immutable")
return
record.metadata = {**record.metadata, EVIDENCE_KEY: copy.deepcopy(evidence)}
async with self._run_scope(run_id), self._snapshot_locks.hold(run_id):
async with self._lock:
record = self._runs.get(run_id)
if record is None or evidence.get("run_id") != run_id or evidence.get("thread_id") != record.thread_id:
raise ValueError("Assembly/run binding mismatch")
if EVIDENCE_KEY in record.metadata:
if record.metadata[EVIDENCE_KEY] != evidence:
raise ValueError("Run assembly evidence is immutable")
return
record.metadata = {**record.metadata, EVIDENCE_KEY: copy.deepcopy(evidence)}
if not await self._persist_to_store(record):
record.metadata.pop(EVIDENCE_KEY, None)
raise RuntimeError("Unable to persist assembly evidence")
async def update_model_name(self, run_id: str, model_name: str | None) -> None:
"""Update the model name for a run."""
async with self._lock:
async with self._run_registry_lock(run_id):
record = self._runs.get(run_id)
if record is None:
logger.warning("update_model_name called for unknown run %s", run_id)
@ -1688,7 +1724,7 @@ class RunManager:
action: str,
) -> None:
"""Set process-local abort state without status persistence or cleanup."""
async with self._lock:
async with self._run_registry_lock(run_id):
record = self._runs.get(run_id)
if record is None or (record.status not in (RunStatus.pending, RunStatus.running) and not self._awaits_terminal_commit(record)) or record.abort_event.is_set():
return
@ -1746,7 +1782,7 @@ class RunManager:
# ------------------------------------------------------------------
# Local path — this worker owns the run in-memory.
# ------------------------------------------------------------------
async with self._lock:
async with self._run_registry_lock(run_id):
record = self._runs.get(run_id)
if record is not None:
if record.status == RunStatus.interrupted:
@ -1771,7 +1807,7 @@ class RunManager:
elif outcome != CancelOutcome.lease_valid_elsewhere:
return outcome
async with self._lock:
async with self._run_registry_lock(run_id):
record = self._runs.get(run_id)
if record is not None:
if record.status == RunStatus.interrupted or record.abort_event.is_set():
@ -1797,7 +1833,7 @@ class RunManager:
# cancellation; ordinary runs retain immediate persistence.
return CancelOutcome.cancelled
persisted = await self._persist_status(record, RunStatus.interrupted)
async with self._lock:
async with self._run_registry_lock(run_id):
if self._runs.get(run_id) is record:
record.terminal_status_persisted = self._has_final_terminal_ack(record) if self.heartbeat_enabled and self._store is not None else persisted
if not persisted and self._store is not None:
@ -2001,14 +2037,14 @@ class RunManager:
raise RuntimeError(f"Cancelled admission {record.run_id} remains active in the run store")
if stored is None:
async with self._lock:
async with self._run_registry_lock(record.run_id):
if self._runs.get(record.run_id) is record:
self._runs.pop(record.run_id, None)
self._unindex_run_locked(record.run_id, record.thread_id)
return
stored_status = RunStatus(stored.get("status") or RunStatus.pending.value)
async with self._lock:
async with self._run_registry_lock(record.run_id):
if self._runs.get(record.run_id) is record:
record.status = stored_status
record.terminal_status_persisted = stored_status not in (
@ -2044,7 +2080,7 @@ class RunManager:
already has a pending/running run. For ``interrupt``/``rollback``,
cancels inflight runs before creating.
Lock ordering invariant: the local ``self._lock`` is held across
Lock ordering invariant: the per-thread admission gate is held across
the local check, the store insert, and the local register, so the
store insert can never succeed while a same-worker ConflictError
is about to fire (which would leak a pending row in the store).
@ -2093,52 +2129,53 @@ class RunManager:
idempotency_request=idempotency_request,
)
async with self._lock:
if idempotency_key is not None and execution_fence is None:
for existing in self._runs.values():
if existing.idempotency_key != idempotency_key:
continue
async with self._admission_scope(thread_id, run_id):
async with self._lock:
if idempotency_key is not None and execution_fence is None:
for existing in self._runs.values():
if existing.idempotency_key != idempotency_key:
continue
if existing.thread_id != thread_id or existing.user_id != user_id:
raise RunIdempotencyKeyMismatch("Run idempotency key resolved to a different thread or user")
existing.idempotency_reused = True
return existing
def reuse_idempotent_run(conflict: RunIdempotencyConflict) -> RunRecord:
# A locally held record for this key already returned above, so
# the conflicting row belongs to a peer or to a run this worker
# has cleaned up. Return a store-only handle without registering
# it: nothing here finalizes or cleans up that record, so a
# registered copy would keep its admission-time status, reject
# later admissions for the thread, and shadow the durable row
# for get(), cancel(), and orphan reconciliation.
existing = self._record_from_store(conflict.existing)
if existing.thread_id != thread_id or existing.user_id != user_id:
raise RunIdempotencyKeyMismatch("Run idempotency key resolved to a different thread or user")
raise RunIdempotencyKeyMismatch("Run idempotency key resolved to a different thread or user") from conflict
existing.idempotency_reused = True
return existing
def reuse_idempotent_run(conflict: RunIdempotencyConflict) -> RunRecord:
# A locally held record for this key already returned above, so
# the conflicting row belongs to a peer or to a run this worker
# has cleaned up. Return a store-only handle without registering
# it: nothing here finalizes or cleans up that record, so a
# registered copy would keep its admission-time status, reject
# later admissions for the thread, and shadow the durable row
# for get(), cancel(), and orphan reconciliation.
existing = self._record_from_store(conflict.existing)
if existing.thread_id != thread_id or existing.user_id != user_id:
raise RunIdempotencyKeyMismatch("Run idempotency key resolved to a different thread or user") from conflict
existing.idempotency_reused = True
return existing
# 1) Local inflight check (same-worker guard; cross-worker is the
# store's partial unique index below).
local_inflight = [r for r in self._thread_records_locked(thread_id) if r.status in (RunStatus.pending, RunStatus.running) or r.finalizing or r.scheduled_goal_cleanup_pending]
# 1) Local inflight check (same-worker guard; cross-worker is the
# store's partial unique index below).
local_inflight = [r for r in self._thread_records_locked(thread_id) if r.status in (RunStatus.pending, RunStatus.running) or r.finalizing or r.scheduled_goal_cleanup_pending]
if any(r.scheduled_goal_cleanup_pending for r in local_inflight):
raise ConflictError(f"Thread {thread_id} has a scheduled goal awaiting finalization")
if any(r.scheduled_goal_cleanup_pending for r in local_inflight):
raise ConflictError(f"Thread {thread_id} has a scheduled goal awaiting finalization")
if multitask_strategy in ("interrupt", "rollback") and any(record.operation_kind != ThreadOperationKind.run for record in local_inflight):
raise ConflictError(f"Thread {thread_id} has an active checkpoint write")
if multitask_strategy in ("interrupt", "rollback") and any(record.operation_kind != ThreadOperationKind.run for record in local_inflight):
raise ConflictError(f"Thread {thread_id} has an active checkpoint write")
if multitask_strategy == "reject" and local_inflight:
raise ConflictError(f"Thread {thread_id} already has an active run")
if multitask_strategy == "reject" and local_inflight:
raise ConflictError(f"Thread {thread_id} already has an active run")
if multitask_strategy in ("interrupt", "rollback") and local_inflight:
logger.info(
"Preparing to cancel %d inflight run(s) on thread %s (strategy=%s)",
len(local_inflight),
thread_id,
multitask_strategy,
)
if multitask_strategy in ("interrupt", "rollback") and local_inflight:
logger.info(
"Preparing to cancel %d inflight run(s) on thread %s (strategy=%s)",
len(local_inflight),
thread_id,
multitask_strategy,
)
# 2) Persist to store while still holding the local lock. The
# 2) Persist with only the same-thread admission gate held. The
# store is the source of truth for cross-process atomicity.
if self._store is not None:
if idempotency_request is not None and not _store_accepts_idempotency_parameter(self._store, "idempotency_request"):
@ -2242,27 +2279,28 @@ class RunManager:
# rows as interrupted in the same transaction; no extra
# store write is needed for them.
# 3) Only now safe to register locally — store insert succeeded.
self._runs[run_id] = record
self._index_run_locked(record)
self._arm_lease_deadline(record)
async with self._lock:
# 3) Only now safe to register locally — store insert succeeded.
self._runs[run_id] = record
self._index_run_locked(record)
self._arm_lease_deadline(record)
# 4) Cancel local in-memory inflight (interrupt/rollback). The
# store-side counterparts were already cancelled in step 2.
if self._store is None and multitask_strategy in ("interrupt", "rollback"):
for r in local_inflight:
if r.finalizing:
continue
r.abort_action = multitask_strategy
r.abort_event.set()
task_active = r.task is not None and not r.task.done()
r.finalizing = task_active
if task_active:
r.task.cancel()
r.status = RunStatus.interrupted
r.terminal_status_persisted = self._store is None
r.updated_at = now
interrupted_records.append(r)
# 4) Cancel local in-memory inflight (interrupt/rollback). The
# store-side counterparts were already cancelled in step 2.
if self._store is None and multitask_strategy in ("interrupt", "rollback"):
for r in local_inflight:
if r.finalizing:
continue
r.abort_action = multitask_strategy
r.abort_event.set()
task_active = r.task is not None and not r.task.done()
r.finalizing = task_active
if task_active:
r.task.cancel()
r.status = RunStatus.interrupted
r.terminal_status_persisted = self._store is None
r.updated_at = now
interrupted_records.append(r)
# Outside the lock: persist interrupted status for locally-cancelled
# runs. Store-side claimed rows are already finalised. Cancellation at
@ -2328,7 +2366,7 @@ class RunManager:
if reservation_task is None:
raise RuntimeError("Thread operation reservation requires an active asyncio task")
lease_lost = True
async with self._lock:
async with self._run_registry_lock(record.run_id):
if self._runs.get(record.run_id) is record:
record.task = reservation_task
lease_lost = record.abort_event.is_set()
@ -2355,7 +2393,7 @@ class RunManager:
exc_info=True,
)
finally:
async with self._lock:
async with self._run_registry_lock(record.run_id):
removed = self._runs.pop(record.run_id, None)
if removed is not None:
self._unindex_run_locked(record.run_id, removed.thread_id)
@ -2472,7 +2510,7 @@ class RunManager:
async def has_inflight(self, thread_id: str) -> bool:
"""Return ``True`` if *thread_id* has a pending or running run."""
async with self._lock:
async with self._thread_registry_lock(thread_id):
return any(r.operation_kind == ThreadOperationKind.run and (r.status in (RunStatus.pending, RunStatus.running) or r.finalizing) for r in self._thread_records_locked(thread_id))
async def cleanup(self, run_id: str, *, delay: float = 300) -> None:
@ -2488,7 +2526,7 @@ class RunManager:
return
if delay > 0:
await asyncio.sleep(delay)
async with self._lock:
async with self._run_registry_lock(run_id):
record = self._runs.pop(run_id, None)
if record is not None:
self._unindex_run_locked(run_id, record.thread_id)
@ -2505,7 +2543,7 @@ class RunManager:
async def owns_active_admission(self, record: RunRecord) -> bool:
"""Confirm the existing run slot without holding a SQL writer transaction."""
async with self._lock:
async with self._run_registry_lock(record.run_id):
if self._runs.get(record.run_id) is not record or record.ownership_lost:
return False
if self._store is None:
@ -2841,6 +2879,18 @@ class RunManager:
return loop.time() < deadline and self._rollback_retry_authorized(record, initial_allow_interrupted=initial_allow_interrupted)
async def _renew_one_lease(self, record: RunRecord) -> None:
# Assembly publication writes a complete snapshot. Keep the Store
# renewal and its local ACK together so that snapshot cannot restore
# the previous lease. Waiting for this gate retains the expiry timer.
self._arm_lease_deadline(record)
async with self._snapshot_locks.hold(record.run_id):
action = await self._renew_one_lease_locked(record)
if action is not None:
# Do not acquire the thread gate while holding the snapshot gate:
# assembly publication takes them in the opposite order.
await self._signal_local_cancel(record.run_id, action=action)
async def _renew_one_lease_locked(self, record: RunRecord) -> str | None:
assert self._store is not None and self._run_ownership_config is not None
if self._has_final_terminal_ack(record):
await self._release_confirmed_terminal_barrier(record, timeout=0)
@ -2915,7 +2965,7 @@ class RunManager:
if renewal.cancel_action is not None:
action = renewal.cancel_action if renewal.cancel_action in ("interrupt", "rollback") else "interrupt"
# Signal here, before joining any other run's decision.
await self._signal_local_cancel(record.run_id, action=action)
return action
return
if (
retry_index == 0

280
backend/tests/test_run_admission_isolation.py Normal file
View File

@ -0,0 +1,280 @@
"""An admission waiting on its store must not block unrelated run control."""
import asyncio
import pytest
from deerflow.config.run_ownership_config import RunOwnershipConfig
from deerflow.runtime.runs.manager import ConflictError, RunManager
from deerflow.runtime.runs.store.memory import MemoryRunStore
class HeldAdmissionStore(MemoryRunStore):
def __init__(self):
super().__init__()
self.entered = asyncio.Event()
self.release = asyncio.Event()
self.renewed = asyncio.Event()
async def put(self, run_id, *, thread_id, **kwargs):
if thread_id == "slow-thread":
self.entered.set()
await self.release.wait()
return await super().put(run_id, thread_id=thread_id, **kwargs)
async def create_thread_operation_atomic(self, *, thread_id, idempotency_key=None, **kwargs):
if thread_id == "slow-thread":
self.entered.set()
await self.release.wait()
return await super().create_thread_operation_atomic(thread_id=thread_id, idempotency_key=idempotency_key, **kwargs)
async def renew_lease(self, run_id, **kwargs):
result = await super().renew_lease(run_id, **kwargs)
if result.renewed:
self.renewed.set()
return result
@pytest.mark.anyio
@pytest.mark.parametrize("entrypoint", ["create", "create_or_reject", "record_assembly_evidence"])
async def test_admission_wait_does_not_block_another_threads_read_and_cancel(entrypoint):
store = HeldAdmissionStore()
manager = RunManager(store=store)
record = await manager.create_or_reject("other-thread", user_id="owner")
await manager.try_start(record.run_id)
started = asyncio.Event()
async def work():
started.set()
await asyncio.Event().wait()
record.task = asyncio.create_task(work())
await started.wait()
if entrypoint == "record_assembly_evidence":
store.release.set()
other = await manager.create_or_reject("slow-thread", user_id="owner")
store.entered.clear()
store.release.clear()
admission = asyncio.create_task(manager.record_assembly_evidence(other.run_id, {"run_id": other.run_id, "thread_id": other.thread_id}))
else:
admission = asyncio.create_task(getattr(manager, entrypoint)("slow-thread", user_id="owner"))
readers = []
try:
await asyncio.wait_for(store.entered.wait(), timeout=5)
readers = [asyncio.create_task(manager.get(record.run_id, user_id="owner")), asyncio.create_task(manager.cancel(record.run_id))]
result = await asyncio.wait_for(asyncio.gather(*readers), timeout=5)
assert result[0] is record
await asyncio.wait_for(asyncio.gather(record.task, return_exceptions=True), timeout=5)
assert record.task.cancelled()
assert not admission.done()
finally:
store.release.set()
await asyncio.wait_for(asyncio.gather(admission, return_exceptions=True), timeout=5)
record.task.cancel()
await asyncio.gather(record.task, *readers, return_exceptions=True)
@pytest.mark.anyio
async def test_admission_wait_does_not_stop_healthy_run_heartbeats():
store = HeldAdmissionStore()
manager = RunManager(store=store, run_ownership_config=RunOwnershipConfig(heartbeat_enabled=True, lease_seconds=5, grace_seconds=0))
record = await manager.create_or_reject("healthy-thread", user_id="owner")
await manager.try_start(record.run_id)
record.task = asyncio.create_task(asyncio.Event().wait())
initial_lease = record.lease_expires_at
store.renewed.clear()
admission = asyncio.create_task(manager.create_or_reject("slow-thread", user_id="owner"))
try:
await asyncio.wait_for(store.entered.wait(), timeout=5)
await manager.start_heartbeat()
await asyncio.wait_for(store.renewed.wait(), timeout=5)
# The Store operation and its Manager decision run in separate tasks.
# Observe the Manager's acknowledgment, not merely the Store return.
async with asyncio.timeout(5):
while record.lease_expires_at == initial_lease:
await asyncio.sleep(0)
assert record.lease_expires_at > initial_lease
assert not record.ownership_lost
assert not record.task.done()
assert not admission.done()
finally:
store.release.set()
admitted = await asyncio.wait_for(asyncio.gather(admission, return_exceptions=True), timeout=5)
if not isinstance(admitted[0], BaseException):
await manager.cancel(admitted[0].run_id)
await manager.cancel(record.run_id)
await manager.stop_heartbeat(timeout=3)
record.task.cancel()
await asyncio.gather(record.task, return_exceptions=True)
@pytest.mark.anyio
async def test_same_thread_admissions_and_preallocated_id_reads_wait_for_publication():
store = HeldAdmissionStore()
manager = RunManager(store=store)
first = asyncio.create_task(manager.create_or_reject("slow-thread", run_id="first-run"))
tasks = [first]
started = asyncio.Event()
async def second_admission():
started.set()
return await manager.create_or_reject("slow-thread", run_id="second-run")
try:
await asyncio.wait_for(store.entered.wait(), timeout=5)
second = asyncio.create_task(second_admission())
read = asyncio.create_task(manager.get("first-run"))
history = asyncio.create_task(manager.list_by_thread("slow-thread"))
tasks.extend([second, read, history])
await asyncio.wait_for(started.wait(), timeout=5)
assert not second.done() and not read.done() and not history.done()
assert await store.get("first-run") is None
store.release.set()
record = await asyncio.wait_for(first, timeout=5)
with pytest.raises(ConflictError):
await asyncio.wait_for(second, timeout=5)
assert await asyncio.wait_for(read, timeout=5) is record
assert await asyncio.wait_for(history, timeout=5) == [record]
assert await store.get("second-run") is None
finally:
store.release.set()
await asyncio.gather(*tasks, return_exceptions=True)
@pytest.mark.anyio
async def test_cancelled_waiter_does_not_release_an_active_admission():
store = HeldAdmissionStore()
manager = RunManager(store=store)
owner = asyncio.create_task(manager.create_or_reject("slow-thread", run_id="owner-run"))
tasks = [owner]
waiter_entered = asyncio.Event()
async def waiter():
waiter_entered.set()
return await manager.create_or_reject("slow-thread", run_id="cancelled-run")
try:
await asyncio.wait_for(store.entered.wait(), timeout=5)
queued = asyncio.create_task(waiter())
tasks.append(queued)
await asyncio.wait_for(waiter_entered.wait(), timeout=5)
queued.cancel()
with pytest.raises(asyncio.CancelledError):
await queued
assert not owner.done()
store.release.set()
first = await asyncio.wait_for(owner, timeout=5)
assert await manager.get("cancelled-run") is None
assert await manager.list_by_thread("slow-thread") == [first]
with pytest.raises(ConflictError):
await manager.create_or_reject("slow-thread")
finally:
store.release.set()
await asyncio.gather(*tasks, return_exceptions=True)
@pytest.mark.anyio
async def test_same_thread_keyed_retry_reuses_the_committed_admission():
store = HeldAdmissionStore()
manager = RunManager(store=store)
first = asyncio.create_task(manager.create_or_reject("slow-thread", idempotency_key="request-key"))
tasks = [first]
try:
await asyncio.wait_for(store.entered.wait(), timeout=5)
retry = asyncio.create_task(manager.create_or_reject("slow-thread", idempotency_key="request-key"))
tasks.append(retry)
store.release.set()
original, reused = await asyncio.wait_for(asyncio.gather(first, retry), timeout=5)
assert reused.run_id == original.run_id
assert reused.idempotency_reused
assert len(await manager.list_by_thread("slow-thread")) == 1
finally:
store.release.set()
await asyncio.gather(*tasks, return_exceptions=True)
@pytest.mark.anyio
async def test_assembly_snapshot_does_not_overwrite_an_acknowledged_renewal():
store = MemoryRunStore()
manager = RunManager(store=store, run_ownership_config=RunOwnershipConfig(heartbeat_enabled=True, lease_seconds=5, grace_seconds=0))
record = await manager.create_or_reject("thread", user_id="owner")
# Assembly is published before execution; pending runs also renew leases.
renewed = asyncio.Event()
release_ack = asyncio.Event()
assembly_started = asyncio.Event()
original_renew = store.renew_lease
renewal_targets = []
async def delayed_ack(run_id, **kwargs):
result = await original_renew(run_id, **kwargs)
renewal_targets.append(kwargs["lease_expires_at"])
renewed.set()
await release_ack.wait()
return result
async def publish_assembly():
assembly_started.set()
await manager.record_assembly_evidence(record.run_id, {"run_id": record.run_id, "thread_id": record.thread_id})
store.renew_lease = delayed_ack
publication = None
try:
await manager.start_heartbeat()
await asyncio.wait_for(renewed.wait(), timeout=5)
publication = asyncio.create_task(publish_assembly())
await asyncio.wait_for(assembly_started.wait(), timeout=5)
release_ack.set()
await manager.stop_heartbeat(timeout=5)
await asyncio.wait_for(publication, timeout=5)
stored = await store.get(record.run_id)
assert stored["lease_expires_at"] >= renewal_targets[0]
assert stored["lease_expires_at"] == record.lease_expires_at
finally:
release_ack.set()
if publication is not None:
await asyncio.gather(publication, return_exceptions=True)
await manager.cancel(record.run_id)
await manager.stop_heartbeat(timeout=3)
@pytest.mark.anyio
async def test_renewal_cancel_does_not_hold_snapshot_gate_while_waiting_for_assembly():
store = MemoryRunStore()
manager = RunManager(store=store, run_ownership_config=RunOwnershipConfig(heartbeat_enabled=True, lease_seconds=5, grace_seconds=0))
record = await manager.create_or_reject("thread", user_id="owner")
await store.request_cancel(record.run_id, action="interrupt")
renewed, release_ack, assembly_started = asyncio.Event(), asyncio.Event(), asyncio.Event()
original_renew = store.renew_lease
async def delayed_ack(run_id, **kwargs):
result = await original_renew(run_id, **kwargs)
assert result.cancel_action == "interrupt"
renewed.set()
await release_ack.wait()
return result
async def publish_assembly():
assembly_started.set()
await manager.record_assembly_evidence(record.run_id, {"run_id": record.run_id, "thread_id": record.thread_id})
store.renew_lease = delayed_ack
renewal = asyncio.create_task(manager._renew_leases())
publication = None
try:
await asyncio.wait_for(renewed.wait(), timeout=5)
publication = asyncio.create_task(publish_assembly())
await asyncio.wait_for(assembly_started.wait(), timeout=5)
release_ack.set()
await asyncio.wait_for(asyncio.shield(renewal), timeout=2)
await asyncio.wait_for(publication, timeout=2)
assert record.abort_event.is_set()
assert record.abort_action == "interrupt"
assert not record.ownership_lost
finally:
release_ack.set()
if publication is not None:
publication.cancel()
await asyncio.gather(publication, return_exceptions=True)
await manager.stop_heartbeat(timeout=3)
await asyncio.gather(renewal, return_exceptions=True)
await manager.cancel(record.run_id)

14
backend/tests/test_stale_active_renewal.py
View File

@ -9,6 +9,7 @@ from test_terminal_store_contract import store as store
from deerflow.config.run_ownership_config import RunOwnershipConfig
from deerflow.runtime import RunManager, RunStatus
from deerflow.runtime.runs import manager as manager_module
def manager_for(run_store):
@ -41,7 +42,15 @@ async def finish(manager, record, *tasks):
@pytest.mark.anyio
async def test_pre_rollback_renewal_retries_after_interrupted_ack(store):
async def test_pre_rollback_renewal_retries_after_interrupted_ack(store, monkeypatch):
now = datetime.now(UTC)
class Clock(datetime):
@classmethod
def now(cls, tz=None):
return now if tz is not None else now.replace(tzinfo=None)
monkeypatch.setattr(manager_module, "datetime", Clock)
manager = manager_for(store)
record = await running_record(manager)
entered, release = asyncio.Event(), asyncio.Event()
@ -66,6 +75,9 @@ async def test_pre_rollback_renewal_retries_after_interrupted_ack(store):
assert record.rollback_finalization_pending
old_expiry = record.lease_expires_at
# Memory calls may finish within one Windows wall-clock tick. Pin a
# later renewal instant instead of relying on incidental test latency.
now += timedelta(seconds=1)
release.set()
await asyncio.wait_for(renewal, 2)