diff --git a/README.md b/README.md index 4a1aaed8c1..dd4938670a 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/backend/docs/runtime-guidance-details.md b/backend/docs/runtime-guidance-details.md index 86493105f7..023237eb17 100644 --- a/backend/docs/runtime-guidance-details.md +++ b/backend/docs/runtime-guidance-details.md @@ -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 diff --git a/backend/packages/harness/deerflow/runtime/AGENTS.md b/backend/packages/harness/deerflow/runtime/AGENTS.md index 773760528c..e1945b7d81 100644 --- a/backend/packages/harness/deerflow/runtime/AGENTS.md +++ b/backend/packages/harness/deerflow/runtime/AGENTS.md @@ -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. diff --git a/backend/packages/harness/deerflow/runtime/runs/manager.py b/backend/packages/harness/deerflow/runtime/runs/manager.py index 434e5b0540..e5c90bc4ba 100644 --- a/backend/packages/harness/deerflow/runtime/runs/manager.py +++ b/backend/packages/harness/deerflow/runtime/runs/manager.py @@ -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 diff --git a/backend/tests/test_run_admission_isolation.py b/backend/tests/test_run_admission_isolation.py new file mode 100644 index 0000000000..04494c7a8a --- /dev/null +++ b/backend/tests/test_run_admission_isolation.py @@ -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) diff --git a/backend/tests/test_stale_active_renewal.py b/backend/tests/test_stale_active_renewal.py index 302831a140..77570d83d7 100644 --- a/backend/tests/test_stale_active_renewal.py +++ b/backend/tests/test_stale_active_renewal.py @@ -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)