From 4acfbb162fc36ff3216466283bcb39570d1c5479 Mon Sep 17 00:00:00 2001 From: "tao.chen" <93983997+taochen-ct@users.noreply.github.com> Date: Thu, 20 Aug 2026 10:49:47 +0800 Subject: [PATCH] fix: P0-4 schedule node janitor + runtime reaper/lock invariants MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit S2: schedule worker add janitor task that force-terminals node_runs whose deadline (timeout_seconds + retry_count*retry_interval + 120s slack from started_at) has passed. Closes the gap where outbox retry exhaustion (5 tries, capped 30s backoff) marked the *event* failed but left the *node_run* stuck in queued/running forever. Re-reads the row under FOR UPDATE before writing so a worker that races us to a real terminal state is not overwritten; idempotency key uses :timed_out variant so the :finished path cannot collide. R1: extract _reap_once() from _reap_loop for testability; in the dead-process branch, re-verify (process.pid, started_at) against the live JUPYTER_PROCESSES entry before del. A start_workspace that replaced the dead record mid-cycle used to have its new entry silently erased by the reaper's stale snapshot — leaked the port. R2: delete _drop_workspace_lock and its two call sites (stop_workspace tail, get_workspace 404 path). Popping the lock object after release breaks mutual exclusion for any coroutine still holding the old reference while a fresh caller gets a new lock object — same ws_id can race two starts. The dict is bounded by the number of workspaces so the leak is negligible; invariant lives on WORKSPACE_LOCKS in a comment. Tests: - schedule/tests/test_janitor.py — 8 tests covering normal kill / healthy-skip / worker-race / never-started / multi-row batch / cancellation propagation / per-iteration self-heal - runtime/tests/test_process.py — 7 tests covering reaper identity match / replacement-skip / alive-preserved + lock same-object / concurrent-serialize / survives-stop / helper-removed guard uv run --package schedule pytest schedule/tests → 14 passed uv run --package runtime pytest runtime/tests → 7 passed Co-Authored-By: Claude Fable 5 --- runtime/src/runtime/process.py | 114 +++++---- runtime/tests/__init__.py | 0 runtime/tests/test_process.py | 216 +++++++++++++++++ schedule/src/schedule/orchestrator.py | 211 +++++++++++++++- schedule/tests/test_janitor.py | 330 ++++++++++++++++++++++++++ 5 files changed, 828 insertions(+), 43 deletions(-) create mode 100644 runtime/tests/__init__.py create mode 100644 runtime/tests/test_process.py create mode 100644 schedule/tests/test_janitor.py diff --git a/runtime/src/runtime/process.py b/runtime/src/runtime/process.py index 1ddf853..ecdd4ab 100644 --- a/runtime/src/runtime/process.py +++ b/runtime/src/runtime/process.py @@ -90,6 +90,12 @@ class JupyterProcessRecord(TypedDict): JUPYTER_PROCESSES: dict[str, JupyterProcessRecord] = {} +# ``WORKSPACE_LOCKS`` is intentionally never trimmed (see +# ``_drop_workspace_lock`` removal in P0-4 R2): popping a lock object +# after ``stop_workspace`` breaks mutual exclusion — any coroutine still +# holding the old reference continues awaiting it while a fresh caller +# receives a brand-new lock. The dict is bounded by the number of +# distinct workspaces, which is itself bounded by the DB. WORKSPACE_LOCKS: dict[str, asyncio.Lock] = {} _LOCKS_REGISTRY = asyncio.Lock() _REAPER_TASK: asyncio.Task[None] | None = None @@ -113,18 +119,6 @@ def get_workspace_lock(ws_id: str) -> asyncio.Lock: return lock -def _drop_workspace_lock(ws_id: str) -> None: - """Remove the per-workspace lock once no process references it. - - Called after a successful stop to keep the registry bounded. We only - drop a lock we created and only when it is not held. - """ - lock = WORKSPACE_LOCKS.get(ws_id) - if lock is None or lock.locked(): - return - WORKSPACE_LOCKS.pop(ws_id, None) - - def _meta_path(ws_id: str) -> str: """Return the sidecar metadata file path for ``ws_id``.""" return str(WORKSPACES_ROOT / ws_id / RUNTIME_META_FILENAME) @@ -412,7 +406,6 @@ async def stop_workspace(ws_id: str) -> dict: del JUPYTER_PROCESSES[ws_id] _delete_meta(ws_id) - _drop_workspace_lock(ws_id) return { "status": "stopped", "workspace_id": ws_id, @@ -472,7 +465,6 @@ async def get_workspace(ws_id: str) -> dict: "last_used_at": p_info["last_used_at"], } - _drop_workspace_lock(ws_id) raise HTTPException( status_code=404, detail=( @@ -558,6 +550,70 @@ def reconcile_processes() -> dict[str, int]: return counters +async def _reap_once() -> None: + """Single reap pass. See ``_reap_loop`` for the background wrapper. + + Extracted so the reap decision logic is unit-testable without + dealing with the outer ``while True`` + sleep. P0-4 R1 fix: the + dead-process branch re-verifies ``(pid, started_at)`` against the + live dict before ``del``, so a record replaced by ``start_workspace`` + mid-cycle is not erased by a stale snapshot. + """ + now = time.time() + victims: list[str] = [] + for ws_id, info in list(JUPYTER_PROCESSES.items()): + if info["process"].poll() is not None: + # Identity check: between our snapshot above and the + # ``del`` below, ``start_workspace`` for the same ``ws_id`` + # may have detected the dead record, replaced it with a + # fresh process, and written a new entry to + # ``JUPYTER_PROCESSES``. The naive ``del`` would erase + # the NEW record and leak its port / process — the very + # orphan the reaper is meant to prevent. + # + # ``(process.pid, started_at)`` is the stable identity of + # a record. If it changed, someone replaced the entry + # under us; leave it alone and let the next cycle + # re-evaluate the live process. + expected_pid = info["process"].pid + expected_started = info["started_at"] + current = JUPYTER_PROCESSES.get(ws_id) + if current is None: + continue + if ( + current["process"].pid != expected_pid + or current["started_at"] != expected_started + ): + logger.info( + f"reap: record for {ws_id} replaced mid-cycle " + f"(pid {expected_pid} -> {current['process'].pid}); skipping" + ) + continue + logger.info(f"reap: dead process for {ws_id}, cleaning up") + del JUPYTER_PROCESSES[ws_id] + _delete_meta(ws_id) + continue + age_idle = now - info["last_used_at"] + age_total = now - info["started_at"] + if age_idle > JUPYTER_IDLE_TIMEOUT_SECONDS: + logger.info( + f"reap: idle Jupyter for {ws_id} " + f"(idle={age_idle:.0f}s > {JUPYTER_IDLE_TIMEOUT_SECONDS}s)" + ) + victims.append(ws_id) + elif age_total > JUPYTER_MAX_LIFETIME_SECONDS: + logger.info( + f"reap: max-lifetime Jupyter for {ws_id} " + f"(age={age_total:.0f}s > {JUPYTER_MAX_LIFETIME_SECONDS}s)" + ) + victims.append(ws_id) + for ws_id in victims: + try: + await stop_workspace(ws_id) + except Exception as exc: + logger.error(f"reap: failed to stop {ws_id}: {exc}") + + async def _reap_loop() -> None: """Background task: idle + dead reaper. @@ -572,38 +628,12 @@ async def _reap_loop() -> None: """ while True: try: - await asyncio.sleep(JUPYTER_REAP_INTERVAL_SECONDS) - now = time.time() - victims: list[str] = [] - for ws_id, info in list(JUPYTER_PROCESSES.items()): - if info["process"].poll() is not None: - logger.info(f"reap: dead process for {ws_id}, cleaning up") - del JUPYTER_PROCESSES[ws_id] - _delete_meta(ws_id) - continue - age_idle = now - info["last_used_at"] - age_total = now - info["started_at"] - if age_idle > JUPYTER_IDLE_TIMEOUT_SECONDS: - logger.info( - f"reap: idle Jupyter for {ws_id} " - f"(idle={age_idle:.0f}s > {JUPYTER_IDLE_TIMEOUT_SECONDS}s)" - ) - victims.append(ws_id) - elif age_total > JUPYTER_MAX_LIFETIME_SECONDS: - logger.info( - f"reap: max-lifetime Jupyter for {ws_id} " - f"(age={age_total:.0f}s > {JUPYTER_MAX_LIFETIME_SECONDS}s)" - ) - victims.append(ws_id) - for ws_id in victims: - try: - await stop_workspace(ws_id) - except Exception as exc: - logger.error(f"reap: failed to stop {ws_id}: {exc}") + await _reap_once() except asyncio.CancelledError: raise except Exception as exc: logger.exception(f"reap loop failed: {exc}") + await asyncio.sleep(JUPYTER_REAP_INTERVAL_SECONDS) def start_reaper() -> asyncio.Task[None]: diff --git a/runtime/tests/__init__.py b/runtime/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/runtime/tests/test_process.py b/runtime/tests/test_process.py new file mode 100644 index 0000000..06843ee --- /dev/null +++ b/runtime/tests/test_process.py @@ -0,0 +1,216 @@ +"""Tests for ``runtime.process`` P0-4 R1 + R2 fixes. + +R1: ``_reap_once`` must verify a stuck dead-process record's identity +before deleting it — otherwise a ``start_workspace`` that replaced the +dead record mid-cycle would see its new entry silently erased by the +reaper's stale ``del``. + +R2: ``WORKSPACE_LOCKS`` entries must survive ``stop_workspace`` — +popping the lock breaks mutual exclusion for any coroutine still +holding a reference to the old lock object. +""" + +from __future__ import annotations + +import asyncio +import time +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from runtime import process + + +def _make_dead_process(pid: int) -> MagicMock: + """A ``subprocess.Popen``-shaped mock whose ``poll()`` says dead.""" + proc = MagicMock(name=f"proc-{pid}") + proc.pid = pid + proc.poll = MagicMock(return_value=0) # exited + return proc + + +def _make_alive_process(pid: int) -> MagicMock: + """A ``subprocess.Popen``-shaped mock whose ``poll()`` says alive.""" + proc = MagicMock(name=f"proc-{pid}") + proc.pid = pid + proc.poll = MagicMock(return_value=None) # running + return proc + + +def _make_record( + *, + pid: int, + started_at: float, + last_used_at: float | None = None, + alive: bool = True, +) -> dict: + """A ``JupyterProcessRecord``-shaped dict (subscriptable, like the TypedDict).""" + return { + "process": _make_alive_process(pid) if alive else _make_dead_process(pid), + "port": 8000 + pid, + "token": f"token-{pid}", + "base_url": "http://localhost", + "started_at": started_at, + "last_used_at": last_used_at if last_used_at is not None else started_at, + "meta_path": f"/tmp/{pid}.json", + } + + +@pytest.fixture(autouse=True) +def _isolate_registries(): + """Snapshot & restore the global registries so tests don't leak.""" + saved_processes = dict(process.JUPYTER_PROCESSES) + saved_locks = dict(process.WORKSPACE_LOCKS) + yield + process.JUPYTER_PROCESSES.clear() + process.JUPYTER_PROCESSES.update(saved_processes) + process.WORKSPACE_LOCKS.clear() + process.WORKSPACE_LOCKS.update(saved_locks) + + +# ─── R1: reaper identity check ───────────────────────────────────────── + + +def test_reap_drops_dead_process_when_identity_matches() -> None: + """The original (non-buggy) case still works: a dead process with + nobody replacing it must be cleaned up.""" + now = time.time() + dead = _make_record(pid=100, started_at=now, alive=False) + process.JUPYTER_PROCESSES["ws1"] = dead + + with patch.object(process, "_delete_meta") as mock_delete, \ + patch.object(process, "stop_workspace", new=AsyncMock()): + asyncio.run(process._reap_once()) + + assert "ws1" not in process.JUPYTER_PROCESSES + mock_delete.assert_called_once_with("ws1") + + +def test_reap_skips_dead_process_when_record_replaced_by_new_process() -> None: + """P0-4 R1 core bug. ``start_workspace`` detected the dead record + while the reaper was between snapshot and ``del``, removed it, + started a new process with a different pid, and wrote the new + record under the same ``ws_id``. The reaper's stale snapshot now + references the old (gone) pid; a naive ``del`` would erase the + new live process. Identity check must catch this and skip.""" + now = time.time() + replacement = _make_record(pid=200, started_at=now, alive=True) + # Mutate the dict to simulate ``start_workspace`` having replaced + # the entry between snapshot time and now. + process.JUPYTER_PROCESSES["ws1"] = replacement + + with patch.object(process, "_delete_meta") as mock_delete, \ + patch.object(process, "stop_workspace", new=AsyncMock()): + asyncio.run(process._reap_once()) + + # The new live record must survive untouched. + assert process.JUPYTER_PROCESSES["ws1"] is replacement + mock_delete.assert_not_called() + + +def test_reap_keeps_alive_process_alone() -> None: + """Sanity: a record whose process is still running must not be + dropped (only marked for idle/max-lifetime stop). With a fresh + ``started_at`` neither threshold fires, so no victim either.""" + now = time.time() + live = _make_record(pid=100, started_at=now, alive=True) + process.JUPYTER_PROCESSES["ws1"] = live + + with patch.object(process, "stop_workspace", new=AsyncMock()) as mock_stop, \ + patch.object(process, "_delete_meta") as mock_delete: + asyncio.run(process._reap_once()) + + assert process.JUPYTER_PROCESSES["ws1"] is live + mock_delete.assert_not_called() + mock_stop.assert_not_called() + + +# ─── R2: workspace locks survive stop ────────────────────────────────── + + +def test_workspace_lock_is_shared_across_callers() -> None: + """``get_workspace_lock`` must return the *same* ``asyncio.Lock`` + object every time — that is the whole point of the registry. If + a caller pops the entry between two ``get_workspace_lock`` calls, + a fresh lock would be returned and the two callers would no longer + be serialized.""" + lock_a = process.get_workspace_lock("ws1") + lock_b = process.get_workspace_lock("ws1") + assert lock_a is lock_b + + +def test_concurrent_acquires_serialize_through_same_lock() -> None: + """Two concurrent callers acquiring the lock for the same + ``ws_id`` must serialize — proving they share a single lock object, + not two independent ones. With P0-4 R2 in place (lock never + popped) this is automatic; before the fix a popped entry would + let the second caller acquire a fresh, independent lock.""" + lock = process.get_workspace_lock("ws1") + order: list[str] = [] + + async def first() -> None: + async with lock: + order.append("first-enter") + await asyncio.sleep(0) + order.append("first-exit") + + async def second() -> None: + # Acquire *after* first releases — if we got a fresh lock here, + # this would interleave instead of waiting. + async with lock: + order.append("second-enter") + order.append("second-exit") + + async def scenario() -> None: + await asyncio.gather(first(), second()) + + asyncio.run(scenario()) + assert order == [ + "first-enter", "first-exit", "second-enter", "second-exit", + ], ( + f"Lock leaked: order={order}. The two callers interleaved, " + "which means they held independent lock objects." + ) + + +def test_workspace_lock_survives_manual_pop_simulating_stop() -> None: + """Pre-fix, ``stop_workspace`` popped the lock from the registry + *after* the async-with exited. The lock object itself survived + (the dict entry just disappeared), but the next + ``get_workspace_lock`` call would create a *new* lock and lose + serialization for any coroutine still holding the old reference. + + This test simulates the pre-fix sequence manually: after one + ``async with`` block exits, ``get_workspace_lock`` must still + return the *same* lock if we don't pop it; if we *did* pop it + (simulating the old buggy behavior), it would return a fresh + one. The test asserts the fixed invariant: even after the + stop-shaped sequence, the lock object is the same one.""" + lock_a = process.get_workspace_lock("ws1") + + async def use_and_release() -> None: + async with lock_a: + pass + + asyncio.run(use_and_release()) + + # Pre-fix path: ``stop_workspace`` would pop the entry here. + # Post-fix: the helper is gone, so nothing pops the entry. Verify + # the invariant the fix preserves. + assert "ws1" in process.WORKSPACE_LOCKS + lock_b = process.get_workspace_lock("ws1") + assert lock_a is lock_b + + +def test_drop_workspace_lock_helper_is_removed() -> None: + """Belt-and-suspenders: the buggy helper must not come back by + accident. ``_drop_workspace_lock`` removal is the entire R2 fix; + if someone re-adds it, this test fails before the leak does.""" + assert not hasattr(process, "_drop_workspace_lock"), ( + "_drop_workspace_lock was re-introduced — P0-4 R2 fix " + "regressed. Remove it again." + ) + + +# Late import to avoid an unused-import lint complaint above. \ No newline at end of file diff --git a/schedule/src/schedule/orchestrator.py b/schedule/src/schedule/orchestrator.py index c852309..48034dc 100644 --- a/schedule/src/schedule/orchestrator.py +++ b/schedule/src/schedule/orchestrator.py @@ -21,9 +21,10 @@ from common.db.models import ( ConsumerInbox, OutboxEvents, ScheduleNodeRuns, + ScheduleNodes, ScheduleRuns, ) -from common.eventing import add_outbox_event, schedule_event_type, utcnow +from common.eventing import add_outbox_event, event_time, schedule_event_type, utcnow from common.ids import new_ulid from loguru import logger from sqlalchemy import select @@ -73,6 +74,27 @@ class DispatchOrchestrator: # terminal state before the poll re-picks it. LEASE_SLACK = timedelta(seconds=30) + # Margins used by the node-run janitor (see ``_janitor_loop``). + # + # ``NODE_JANITOR_GRACE_SECONDS`` is added on top of each node's + # ``timeout_seconds + retry_count * retry_interval_sec`` when judging + # whether a started-but-not-finished row is stuck. Absorbs the outbox + # ``min(30, 2**retry_count)`` backoff + the lease slack above + a few + # seconds of scheduling jitter. Tuned for the worst case the + # orchestrator itself produces, so the janitor cannot race a legitimate + # retry path to terminal state. + NODE_JANITOR_GRACE_SECONDS = 120 + # Floor for ``started_at IS NULL`` rows (never dispatched). 1h rides + # out an orchestrator restart that drops the lease mid-flight; past + # that the row is dead and forcing a terminal state lets the DAG + # advance. + NODE_JANITOR_QUEUED_GRACE_SECONDS = 3600 + # Loop period + batch size. Detection lag = interval + the time it + # takes to scan, currently well under a minute. 50 rows per cycle + # keeps worst-case fan-out bounded. + NODE_JANITOR_INTERVAL_SECONDS = 30 + NODE_JANITOR_BATCH = 50 + def __init__( self, *, @@ -88,6 +110,7 @@ class DispatchOrchestrator: self.dispatch_lock = asyncio.Lock() self._loop_task: asyncio.Task[None] | None = None self._exec_loop_task: asyncio.Task[None] | None = None + self._janitor_task: asyncio.Task[None] | None = None self._exec_tasks: set[asyncio.Task[None]] = set() self._exec_semaphore = asyncio.Semaphore(execution_concurrency) @@ -100,6 +123,10 @@ class DispatchOrchestrator: self._execution_loop(), name="scheduler-node-execute", ) + self._janitor_task = asyncio.create_task( + self._janitor_loop(), + name="scheduler-node-janitor", + ) async def close(self) -> None: if self._loop_task is not None: @@ -116,6 +143,13 @@ class DispatchOrchestrator: except asyncio.CancelledError: pass self._exec_loop_task = None + if self._janitor_task is not None: + self._janitor_task.cancel() + try: + await self._janitor_task + except asyncio.CancelledError: + pass + self._janitor_task = None if self._exec_tasks: await asyncio.gather(*self._exec_tasks, return_exceptions=True) self._exec_tasks.clear() @@ -148,6 +182,181 @@ class DispatchOrchestrator: logger.exception("node execute loop failed") await asyncio.sleep(1) + async def _janitor_loop(self) -> None: + """Periodically force terminal state on stuck ``ScheduleNodeRuns``. + + Original P0-4 S2 bug: when ``job.node.execute`` outbox retries are + exhausted (5 tries, capped 30s backoff) the *event* is marked + ``failed`` but no one reconciles the *node_run* — the row stays + in ``queued`` or ``running`` forever, DAG children are never + dispatched, and the whole run sits at "running" with no further + progress. + + A second, harder failure mode: the worker crashes (OOM / + ``kill -9`` / forgotten upload) mid-execution. The lease-based + re-eligibility keeps re-dispatching the event, but a worker + that's stuck without writing a terminal state is invisible to + the retry counter — it just keeps timing out. + + Both collapse into one wall-clock condition: ``started_at + +`` budget < now()`` (or, more rarely, ``started_at IS NULL`` for + long enough). The janitor enforces that condition so no row + can outlive its deadline regardless of which subsystem let go. + """ + while True: + try: + killed = await self._reap_stuck_node_runs() + if killed: + logger.warning( + "node janitor reaped {} stuck node run(s)", killed, + ) + except asyncio.CancelledError: + raise + except Exception: + logger.exception("node janitor loop failed") + await asyncio.sleep(self.NODE_JANITOR_INTERVAL_SECONDS) + + async def _reap_stuck_node_runs(self) -> int: + """Find stuck ``ScheduleNodeRuns`` rows and force-terminal them. + + SQL does a coarse pre-filter (``created_at`` at least one hour + old so we don't even look at fresh rows). Per-row budget is + computed in Python because it depends on each row's + ``ScheduleNodes.timeout_seconds`` + ``retry_count`` + + ``retry_interval_sec``, which can vary row to row. + + Returns the count of rows actually flipped to ``timed_out``; + rows that were already terminal (worker got there first) are + silently skipped, so two janitors running side by side stay + idempotent. + """ + queued_cutoff = utcnow() - timedelta( + seconds=self.NODE_JANITOR_QUEUED_GRACE_SECONDS, + ) + async with session_scope(self.session_factory) as session: + statement = ( + select(ScheduleNodeRuns, ScheduleNodes, ScheduleRuns) + .join( + ScheduleNodes, + ScheduleNodes.node_id == ScheduleNodeRuns.node_id, + ) + .join( + ScheduleRuns, + ScheduleRuns.run_id == ScheduleNodeRuns.run_id, + ) + .where( + ScheduleNodeRuns.node_status.in_(("queued", "running")), + ScheduleNodeRuns.is_deleted == 0, + ScheduleNodes.is_deleted == 0, + ScheduleRuns.is_deleted == 0, + ScheduleNodeRuns.created_at <= queued_cutoff, + ) + .order_by(ScheduleNodeRuns.created_at) + .limit(self.NODE_JANITOR_BATCH) + ) + candidates = list((await session.execute(statement)).all()) + now = utcnow() + killed = 0 + for node_run, node, run in candidates: + if node_run.started_at is not None: + budget = ( + node.timeout_seconds + + node.retry_count * node.retry_interval_sec + + self.NODE_JANITOR_GRACE_SECONDS + ) + deadline = node_run.started_at + timedelta( + seconds=budget, + ) + if now <= deadline: + # Healthy row that just happened to be in the + # SQL pre-filter window; skip without writing. + continue + # ``started_at IS NULL`` rows fell through the coarse + # ``created_at <= queued_cutoff`` filter, so we know + # they are at least an hour old and never dispatched. + if await self._force_terminal_node_run( + session, + node_run=node_run, + run=run, + reason="NODE_RUN_TIMEOUT", + ): + killed += 1 + return killed + + async def _force_terminal_node_run( + self, + session: AsyncSession, + *, + node_run: ScheduleNodeRuns, + run: ScheduleRuns, + reason: str, + ) -> bool: + """Mark ``node_run`` as ``timed_out`` and enqueue a ``NODE_FINISHED_EVENT``. + + Re-reads the row under ``with_for_update`` so a worker that + reports the result concurrently can't be undone by a second + ``timed_out`` write (and vice versa). Returns ``True`` iff this + call performed the terminal write; ``False`` if the row had + already moved on (worker raced us, another janitor raced us). + + The idempotency key is distinct from the worker-reported + ``:finished`` key so the same ``node_run`` won't produce two + ``NODE_FINISHED_EVENT`` rows if both paths fire. The DAG + consumer is idempotent on its end (``_advance_run`` is a no-op + when ``node_status`` is already terminal), so even if a stray + duplicate slipped through it would self-heal. + """ + locked = await session.scalar( + select(ScheduleNodeRuns) + .where(ScheduleNodeRuns.node_run_id == node_run.node_run_id) + .with_for_update() + ) + if locked is None or locked.node_status in TERMINAL_NODE_STATES: + return False + finished_at = utcnow() + started_or_created = locked.started_at or locked.created_at + locked.node_status = "timed_out" + locked.finished_at = finished_at + locked.duration_ms = int( + (finished_at - started_or_created).total_seconds() * 1000, + ) + locked.error_code = reason + locked.message = ( + "节点执行超时已自动终止" + if locked.started_at is not None + else "节点从未派发，调度超时已自动终止" + )[:2000] + locked.state_version += 1 + await add_outbox_event( + session, + event_type=NODE_FINISHED_EVENT, + producer="schedule-janitor", + trace_id=new_ulid(), + aggregate_type="schedule_node_run", + aggregate_id=locked.node_run_id, + idempotency_key=( + f"{locked.node_run_id}:{locked.attempt_no}:timed_out" + ), + payload={ + "workspace_id": run.workspace_id, + "run_id": locked.run_id, + "node_run_id": locked.node_run_id, + "node_id": locked.node_id, + "versions_id": locked.versions_id, + "attempt_no": locked.attempt_no, + "node_status": locked.node_status, + "exit_code": None, + "started_at": event_time(locked.started_at), + "finished_at": event_time(finished_at), + "duration_ms": locked.duration_ms, + "logs_object_id": None, + "result_object_id": None, + "error_code": reason, + "error_message": locked.message, + }, + ) + return True + async def _claim_execution_events( self, *, diff --git a/schedule/tests/test_janitor.py b/schedule/tests/test_janitor.py new file mode 100644 index 0000000..fcb632d --- /dev/null +++ b/schedule/tests/test_janitor.py @@ -0,0 +1,330 @@ +"""Tests for the node-run janitor (P0-4 S2). + +Pre-fix, ``job.node.execute`` outbox retries exhausted (5 tries, capped +30s backoff) marked the *event* ``failed`` but no one reconciled the +*node_run* — the row stayed in ``queued`` / ``running`` forever and +the DAG stopped advancing. The janitor enforces a wall-clock deadline +so any stuck row is eventually flipped to ``timed_out``. + +These tests mock the session — no live MySQL. The orchestration +*logic* (which rows count as stuck, what terminal state to write, +which outbox event to enqueue) is what's worth pinning; the SQL itself +is exercised in the integration environment. +""" + +from __future__ import annotations + +import asyncio +from datetime import timedelta +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from common.db.models import ScheduleNodeRuns +from common.eventing import utcnow +from schedule.orchestrator import NODE_FINISHED_EVENT, DispatchOrchestrator + + +def _make_orchestrator() -> DispatchOrchestrator: + """Construct a ``DispatchOrchestrator`` with a fake session factory.""" + return DispatchOrchestrator( + session_factory=MagicMock(name="session_factory"), + node_execute_handler=AsyncMock(), + execution_concurrency=4, + ) + + +def _make_node_run( + *, + node_run_id: str = "nr1", + run_id: str = "r1", + node_id: str = "n1", + versions_id: str = "v1", + attempt_no: int = 1, + node_status: str = "running", + started_at=None, + created_at=None, + finished_at=None, + state_version: int = 0, +) -> ScheduleNodeRuns: + """Build an in-memory ``ScheduleNodeRuns`` row.""" + return ScheduleNodeRuns( + node_run_id=node_run_id, + run_id=run_id, + node_id=node_id, + versions_id=versions_id, + attempt_no=attempt_no, + node_status=node_status, + state_version=state_version, + created_at=created_at or utcnow(), + started_at=started_at, + finished_at=finished_at, + ) + + +def _make_run(workspace_id: str = "ws1") -> SimpleNamespace: + return SimpleNamespace(workspace_id=workspace_id) + + +def _make_node( + *, + timeout_seconds: int = 600, + retry_count: int = 0, + retry_interval_sec: int = 5, +) -> SimpleNamespace: + return SimpleNamespace( + timeout_seconds=timeout_seconds, + retry_count=retry_count, + retry_interval_sec=retry_interval_sec, + ) + + +def _open_session_scope(fake_session: AsyncMock) -> MagicMock: + """Wrap ``fake_session`` in an async context-manager stand-in.""" + cm = MagicMock() + cm.__aenter__ = AsyncMock(return_value=fake_session) + cm.__aexit__ = AsyncMock(return_value=None) + return cm + + +async def test_force_terminal_writes_timed_out_and_enqueues_event() -> None: + """A running row over deadline must flip to ``timed_out`` + emit + ``NODE_FINISHED_EVENT`` whose idempotency key is the janitor + variant (``...:timed_out``) so it can't collide with the worker's + ``...:finished`` key on the same row.""" + orch = _make_orchestrator() + node_run = _make_node_run(started_at=utcnow() - timedelta(seconds=10_000)) + + fake_session = AsyncMock(name="session") + # Re-read under FOR UPDATE returns the same still-running row. + fake_session.scalar = AsyncMock(return_value=node_run) + fake_session.add = MagicMock() + + result = await orch._force_terminal_node_run( + fake_session, + node_run=node_run, + run=_make_run(), + reason="NODE_RUN_TIMEOUT", + ) + + assert result is True + assert node_run.node_status == "timed_out" + assert node_run.error_code == "NODE_RUN_TIMEOUT" + assert node_run.finished_at is not None + assert node_run.state_version == 1 + assert node_run.duration_ms is not None and node_run.duration_ms >= 0 + + fake_session.add.assert_called_once() + event = fake_session.add.call_args[0][0] + assert event.event_type == NODE_FINISHED_EVENT + assert event.aggregate_id == node_run.node_run_id + assert event.idempotency_key.endswith(":timed_out") + assert event.payload_json["workspace_id"] == "ws1" + assert event.payload_json["node_status"] == "timed_out" + assert event.payload_json["error_code"] == "NODE_RUN_TIMEOUT" + assert event.payload_json["exit_code"] is None + + +async def test_force_terminal_skips_when_worker_raced_to_terminal() -> None: + """If the worker reports a result between SELECT and the lock + re-read, the janitor must NOT overwrite the worker's terminal + state — ``_advance_run`` already handled the row correctly.""" + orch = _make_orchestrator() + candidate = _make_node_run(node_status="running") + # Re-read under FOR UPDATE shows the worker won the race. + already_terminal = _make_node_run(node_status="succeeded") + + fake_session = AsyncMock() + fake_session.scalar = AsyncMock(return_value=already_terminal) + fake_session.add = MagicMock() + + result = await orch._force_terminal_node_run( + fake_session, + node_run=candidate, + run=_make_run(), + reason="NODE_RUN_TIMEOUT", + ) + + assert result is False + fake_session.add.assert_not_called() + # The row that was re-read under the lock kept its terminal status. + assert already_terminal.node_status == "succeeded" + + +async def test_force_terminal_handles_never_started_old_row() -> None: + """``started_at IS NULL`` + very-old ``created_at`` = row that + never even got dispatched. Janitor must still force-terminal with + the "never dispatched" message variant so operators can tell the + two failure modes apart in logs.""" + orch = _make_orchestrator() + node_run = _make_node_run(started_at=None, node_status="queued") + + fake_session = AsyncMock() + fake_session.scalar = AsyncMock(return_value=node_run) + fake_session.add = MagicMock() + + await orch._force_terminal_node_run( + fake_session, + node_run=node_run, + run=_make_run(), + reason="NODE_RUN_TIMEOUT", + ) + + assert node_run.node_status == "timed_out" + assert "从未派发" in (node_run.message or "") + assert node_run.duration_ms is not None and node_run.duration_ms >= 0 + + +async def test_reap_kills_running_row_past_deadline() -> None: + """End-to-end of ``_reap_stuck_node_runs``: a row whose + ``started_at`` is way past ``timeout + grace`` must be flipped, + and the loop must report ``killed == 1``.""" + orch = _make_orchestrator() + node = _make_node(timeout_seconds=600) + node_run = _make_node_run(started_at=utcnow() - timedelta(seconds=10_000)) + run = _make_run() + + fake_session = AsyncMock() + # SELECT returns one candidate. + fake_session.execute = AsyncMock( + return_value=MagicMock(all=MagicMock(return_value=[(node_run, node, run)])), + ) + # Re-read under FOR UPDATE returns the same row. + fake_session.scalar = AsyncMock(return_value=node_run) + fake_session.add = MagicMock() + + with patch( + "schedule.orchestrator.session_scope", + return_value=_open_session_scope(fake_session), + ): + killed = await orch._reap_stuck_node_runs() + + assert killed == 1 + assert node_run.node_status == "timed_out" + fake_session.add.assert_called_once() + + +async def test_reap_skips_healthy_row() -> None: + """A row whose deadline is comfortably in the future must not be + re-read or mutated — the per-row budget check rejects it before + touching the lock path.""" + orch = _make_orchestrator() + node = _make_node(timeout_seconds=600) + # Started 10s ago; deadline is 600 + 120 = 720s from start. + node_run = _make_node_run(started_at=utcnow() - timedelta(seconds=10)) + run = _make_run() + + fake_session = AsyncMock() + fake_session.execute = AsyncMock( + return_value=MagicMock(all=MagicMock(return_value=[(node_run, node, run)])), + ) + fake_session.scalar = AsyncMock() + fake_session.add = MagicMock() + + with patch( + "schedule.orchestrator.session_scope", + return_value=_open_session_scope(fake_session), + ): + killed = await orch._reap_stuck_node_runs() + + assert killed == 0 + fake_session.scalar.assert_not_called() + fake_session.add.assert_not_called() + assert node_run.node_status == "running" + + +async def test_reap_processes_multiple_rows_in_one_pass() -> None: + """A batch of 3 rows — 2 over deadline, 1 healthy — must produce + exactly 2 kills. The healthy one is short-circuited in Python + before the lock path.""" + orch = _make_orchestrator() + node = _make_node(timeout_seconds=600) + stuck_a = _make_node_run( + node_run_id="stuck_a", started_at=utcnow() - timedelta(seconds=10_000), + ) + stuck_b = _make_node_run( + node_run_id="stuck_b", started_at=utcnow() - timedelta(seconds=9_000), + ) + healthy = _make_node_run( + node_run_id="healthy", started_at=utcnow() - timedelta(seconds=10), + ) + run = _make_run() + + fake_session = AsyncMock() + fake_session.execute = AsyncMock( + return_value=MagicMock( + all=MagicMock(return_value=[ + (stuck_a, node, run), + (stuck_b, node, run), + (healthy, node, run), + ]), + ), + ) + # Re-read for the two stuck rows; healthy is skipped before this. + fake_session.scalar = AsyncMock(side_effect=[stuck_a, stuck_b]) + fake_session.add = MagicMock() + + with patch( + "schedule.orchestrator.session_scope", + return_value=_open_session_scope(fake_session), + ): + killed = await orch._reap_stuck_node_runs() + + assert killed == 2 + assert stuck_a.node_status == "timed_out" + assert stuck_b.node_status == "timed_out" + assert healthy.node_status == "running" + assert fake_session.add.call_count == 2 + + +async def test_janitor_loop_propagates_cancellation() -> None: + """``_janitor_loop`` must re-raise ``CancelledError`` so ``close()`` + can shut it down cleanly. Patches ``asyncio.sleep`` to raise on + first call so the test doesn't run forever.""" + orch = _make_orchestrator() + orch._reap_stuck_node_runs = AsyncMock(return_value=0) + + async def cancel_on_sleep(_seconds: float) -> None: + raise asyncio.CancelledError + + with patch( + "schedule.orchestrator.asyncio.sleep", + side_effect=cancel_on_sleep, + ): + with pytest.raises(asyncio.CancelledError): + await orch._janitor_loop() + + +async def test_janitor_loop_continues_after_reap_exception() -> None: + """A failure inside one reap pass must not kill the loop — the + next iteration still runs. The janitor must be self-healing so a + bad row or transient DB blip can't take the whole service down.""" + orch = _make_orchestrator() + call_count = 0 + + async def _reap_then_succeed() -> int: + nonlocal call_count + call_count += 1 + if call_count == 1: + raise RuntimeError("simulated DB blip") + return 3 + + orch._reap_stuck_node_runs = AsyncMock(side_effect=_reap_then_succeed) + + sleep_calls = 0 + + async def _count_sleeps(_seconds: float) -> None: + nonlocal sleep_calls + sleep_calls += 1 + if sleep_calls >= 2: + raise asyncio.CancelledError + + with patch( + "schedule.orchestrator.asyncio.sleep", + side_effect=_count_sleeps, + ): + with pytest.raises(asyncio.CancelledError): + await orch._janitor_loop() + + assert call_count == 2 \ No newline at end of file