diff --git a/schedule/src/schedule/scheduling/__init__.py b/schedule/src/schedule/scheduling/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/schedule/src/schedule/scheduling/orchestrator.py b/schedule/src/schedule/scheduling/orchestrator.py new file mode 100644 index 0000000..69e921a --- /dev/null +++ b/schedule/src/schedule/scheduling/orchestrator.py @@ -0,0 +1,946 @@ +"""Outbox-driven DAG orchestrator. + +Polls the MySQL ``OutboxEvents`` table for ``schedule.run.requested`` and +``job.node.finished`` events and advances schedule runs accordingly. New +nodes are dispatched by writing ``job.node.execute`` rows to the Outbox +and letting the worker component consume them. + +This module owns no HTTP / boto3 / notebook execution dependencies — +those belong to the worker (see ``schedule.worker``) and the cron +post-back (see ``schedule.service.trigger_schedule``). +""" + +from __future__ import annotations + +import asyncio +from datetime import timedelta +from typing import Any + +from common.db import session_scope +from common.db.models import ( + ConsumerInbox, + OutboxEvents, + ScheduleNodeRuns, + ScheduleNodes, + ScheduleRuns, +) +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 +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from schedule.domain.context import ( + FAILED_NODE_STATES, + TERMINAL_NODE_STATES, + TERMINAL_RUN_STATES, +) + +SCHEDULE_RUN_REQUESTED_EVENT = schedule_event_type("schedule.run.requested") +NODE_EXECUTE_EVENT = schedule_event_type("job.node.execute") +NODE_FINISHED_EVENT = schedule_event_type("job.node.finished") + +class DispatchOrchestrator: + """Polls Outbox + advances DAG schedule runs. + + Two independent loops run in the same process: + + - ``_database_event_loop`` drains ``schedule.run.requested`` and + ``job.node.finished`` events under ``dispatch_lock``. These are + short, in-line DB transactions. + - ``_execution_loop`` claims ``job.node.execute`` events whose + ``available_at <= utcnow()`` and dispatches each as + ``asyncio.create_task`` so the polling path is never blocked by + notebook execution. A semaphore caps concurrent notebooks. + + Holds a ``dispatch_lock`` to keep two concurrent drain loops from + fighting over the same batch. + + Lease semantics live on the outbox row itself, not in the claim + step. A new ``job.node.execute`` event is immediately eligible; + the claim step atomically moves ``available_at`` to ``utcnow() + + node_timeout + LEASE_SLACK``, so a process crash mid-execution lets + the row become eligible again once the lease expires. A hard-coded + 30-minute lease was the + original P0-2 bug: a node with ``timeout_seconds = 86_400`` would + be re-claimed at 30 minutes and run twice; a node with + ``timeout_seconds = 60`` would have its lease expire 29 minutes + too early. Tieing the lease to the actual node timeout closes both + cases. + """ + + # Margin added on top of ``timeout_seconds`` when writing the lease + # ``available_at``. Gives the worker time to update the row to a + # 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, + *, + session_factory: async_sessionmaker[AsyncSession], + node_execute_handler, + execution_concurrency: int = 4, + ) -> None: + self.session_factory = session_factory + # Injected by the facade; routes ``job.node.execute`` outbox events + # to ``NodeExecutor.handle_node_execute``. Kept as a callable so this + # module can stay independent of worker module imports. + self._node_execute_handler = node_execute_handler + 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) + + def start(self) -> None: + self._loop_task = asyncio.create_task( + self._database_event_loop(), + name="scheduler-database-events", + ) + self._exec_loop_task = asyncio.create_task( + 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: + self._loop_task.cancel() + try: + await self._loop_task + except asyncio.CancelledError: + pass + self._loop_task = None + if self._exec_loop_task is not None: + self._exec_loop_task.cancel() + try: + await self._exec_loop_task + 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() + + async def _database_event_loop(self) -> None: + while True: + try: + processed = await self.process_pending_events(limit=20) + if processed: + logger.debug("database event loop processed {} events", processed) + if not processed: + await asyncio.sleep(0.25) + except asyncio.CancelledError: + raise + except Exception: + logger.exception("database event loop failed") + await asyncio.sleep(1) + + async def _execution_loop(self) -> None: + while True: + try: + claimed = await self._claim_execution_events(limit=10) + if claimed: + logger.debug("execution loop claimed {} events", claimed) + if not claimed: + await asyncio.sleep(0.25) + except asyncio.CancelledError: + raise + except Exception: + 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, + *, + limit: int, + ) -> int: + """Claim ``job.node.execute`` rows and dispatch them as tasks. + + Lease is owned by the row itself (the dispatcher sets + ``available_at = utcnow() + node.timeout_seconds + LEASE_SLACK`` + when the event is enqueued), so this method is a pure + read-and-dispatch — no DB writes in the claim step. If the + process dies before ``_run_node_execute`` finishes, the row + re-eligible once ``available_at`` falls back to now; the worker + handler is idempotent (short-circuits on terminal node state). + """ + async with session_scope(self.session_factory) as session: + statement = ( + select(OutboxEvents) + .where( + OutboxEvents.event_status == "pending", + OutboxEvents.event_type == NODE_EXECUTE_EVENT, + OutboxEvents.available_at <= utcnow(), + ) + .order_by(OutboxEvents.created_at) + .limit(limit) + .with_for_update(skip_locked=True) + ) + events = list((await session.scalars(statement)).all()) + now = utcnow() + for item in events: + timeout_seconds = max( + 1, + int(item.payload_json.get("timeout_seconds", 1)), + ) + item.available_at = ( + now + + timedelta(seconds=timeout_seconds) + + self.LEASE_SLACK + ) + claimed: list[tuple[dict[str, Any], str]] = [ + ( + { + "event_type": item.event_type, + "event_id": item.event_id, + "trace_id": item.trace_id, + "payload": item.payload_json, + }, + f"mysql:{item.event_id}", + ) + for item in events + ] + for envelope, message_id in claimed: + task = asyncio.create_task( + self._run_node_execute(envelope, message_id), + name=f"node-execute:{envelope['event_id']}", + ) + self._exec_tasks.add(task) + task.add_done_callback(self._exec_tasks.discard) + logger.debug("claimed {} node execute events", len(claimed)) + return len(claimed) + + async def _run_node_execute( + self, + envelope: dict[str, Any], + message_id: str, + ) -> None: + logger.debug("node execute start: event_id={}", envelope["event_id"][-12:]) + async with self._exec_semaphore: + exc: Exception | None = None + try: + await self._node_execute_handler(envelope, message_id) + except Exception as run_exc: # noqa: BLE001 + exc = run_exc + await self._update_execution_status(envelope["event_id"], exc) + + async def _update_execution_status( + self, + event_id: str, + exc: Exception | None, + ) -> None: + async with session_scope(self.session_factory) as session: + item = await session.scalar( + select(OutboxEvents) + .where(OutboxEvents.event_id == event_id) + .with_for_update() + ) + if item is None: + logger.warning("execution event {} disappeared", event_id) + return + if exc is None: + item.event_status = "published" + item.published_at = utcnow() + item.last_error = None + logger.info("node execute success: event_id={}", event_id[-12:]) + else: + item.retry_count += 1 + item.last_error = str(exc)[:2000] + if item.retry_count >= 5: + item.event_status = "failed" + logger.warning( + "node execute exhausted retries: event_id={} retry_count={}", + event_id[-12:], + item.retry_count, + ) + else: + item.available_at = utcnow() + timedelta( + seconds=min(30, 2 ** item.retry_count), + ) + logger.warning( + "node execute retry scheduled: event_id={} retry_count={} delay={}s", + event_id[-12:], + item.retry_count, + min(30, 2 ** item.retry_count), + ) + + async def process_pending_events( + self, + *, + limit: int = 20, + aggregate_id: str | None = None, + ) -> int: + async with self.dispatch_lock: + async with session_scope(self.session_factory) as session: + statement = ( + select(OutboxEvents) + .where( + OutboxEvents.event_status == "pending", + OutboxEvents.available_at <= utcnow(), + OutboxEvents.event_type.in_( + ( + SCHEDULE_RUN_REQUESTED_EVENT, + NODE_FINISHED_EVENT, + ), + ), + ) + .order_by(OutboxEvents.created_at) + .limit(limit) + ) + if aggregate_id: + statement = statement.where( + OutboxEvents.aggregate_id == aggregate_id + ) + events = list((await session.scalars(statement)).all()) + logger.debug("processed {} outbox events", len(events)) + for item in events: + try: + await self._process_outbox_event(item) + item.event_status = "published" + item.published_at = utcnow() + item.last_error = None + except Exception as exc: + item.retry_count += 1 + item.last_error = str(exc)[:2000] + if item.retry_count >= 5: + item.event_status = "failed" + else: + item.available_at = utcnow() + timedelta( + seconds=min(30, 2 ** item.retry_count) + ) + return len(events) + + async def _process_outbox_event(self, item: OutboxEvents) -> None: + handlers = { + SCHEDULE_RUN_REQUESTED_EVENT: self._handle_run_requested, + NODE_FINISHED_EVENT: self._handle_node_finished, + } + handler = handlers.get(item.event_type) + if handler is None: + raise ValueError(f"unsupported event type: {item.event_type}") + event = { + "event_type": item.event_type, + "event_id": item.event_id, + "trace_id": item.trace_id, + "payload": item.payload_json, + } + await handler(event, f"mysql:{item.event_id}") + + async def dispatch_run(self, run_id: str) -> int: + return await self.process_pending_events( + limit=50, + aggregate_id=run_id, + ) + + async def _start_inbox( + self, + session: AsyncSession, + *, + consumer_name: str, + event_id: str, + message_id: str, + ) -> tuple[ConsumerInbox, bool]: + item = await session.get( + ConsumerInbox, + (consumer_name, event_id), + with_for_update=True, + ) + if item is not None and item.process_status == "succeeded": + return item, False + if item is None: + item = ConsumerInbox( + consumer_name=consumer_name, + event_id=event_id, + process_status="processing", + message_id=message_id, + ) + session.add(item) + else: + item.process_status = "processing" + item.message_id = message_id + item.error_message = None + return item, True + + @staticmethod + def _finish_inbox(item: ConsumerInbox) -> None: + item.process_status = "succeeded" + item.processed_at = utcnow() + item.error_message = None + + async def _handle_run_requested( + self, + event: dict[str, Any], + message_id: str, + ) -> None: + if event.get("event_type") != SCHEDULE_RUN_REQUESTED_EVENT: + raise ValueError("unexpected event type") + payload = event["payload"] + logger.info( + "run requested: run={} trace={}", + payload["run_id"][-12:], + event["trace_id"][-12:], + ) + async with session_scope(self.session_factory) as session: + inbox, should_process = await self._start_inbox( + session, + consumer_name="schedule-orchestrator", + event_id=event["event_id"], + message_id=message_id, + ) + if not should_process: + return + run = await session.scalar( + select(ScheduleRuns) + .where(ScheduleRuns.run_id == payload["run_id"]) + .with_for_update() + ) + if run is None: + raise ValueError("schedule run does not exist") + if run.run_status not in TERMINAL_RUN_STATES: + if run.run_status == "queued": + run.run_status = "running" + run.started_at = utcnow() + run.state_version += 1 + await self._bootstrap_root_nodes( + session, + run, + trace_id=event["trace_id"], + ) + await self._advance_run( + session, + run, + trace_id=event["trace_id"], + ) + self._finish_inbox(inbox) + + async def _dispatch_node( + self, + session: AsyncSession, + *, + run: ScheduleRuns, + node: dict[str, Any], + attempt_no: int, + trace_id: str, + delay_seconds: int = 0, + ) -> ScheduleNodeRuns: + node_run = ScheduleNodeRuns( + node_run_id=new_ulid(), + run_id=run.run_id, + node_id=node["node_id"], + versions_id=node["versions_id"], + attempt_no=attempt_no, + node_status="queued", + state_version=0, + message=( + f"等待重试（{delay_seconds} 秒）" + if delay_seconds + else "等待 Worker 执行" + ), + ) + session.add(node_run) + # ``available_at`` is the first execution time here. The claim + # step moves it forward by timeout + slack to become the retry + # lease while the worker is running. + retry_at = ( + utcnow() + timedelta(seconds=delay_seconds) + if delay_seconds + else utcnow() + ) + logger.info( + "dispatch node: run={} node={} attempt={}", + run.run_id[-12:], + node["node_id"][-12:], + attempt_no, + ) + await add_outbox_event( + session, + event_type=NODE_EXECUTE_EVENT, + producer="schedule-orchestrator", + trace_id=trace_id, + aggregate_type="schedule_node_run", + aggregate_id=node_run.node_run_id, + idempotency_key=f"{node_run.node_run_id}:{attempt_no}", + available_at=retry_at, + payload={ + "workspace_id": run.workspace_id, + "run_id": run.run_id, + "node_run_id": node_run.node_run_id, + "node_id": node["node_id"], + "versions_id": node["versions_id"], + "attempt_no": attempt_no, + "script_type": node["script_type"], + "artifact_object_id": node["artifact_object_id"], + "artifact_path": node["artifact_path"], + "timeout_seconds": node["timeout_seconds"], + "arguments": node.get("arguments", []), + }, + ) + return node_run + + async def _bootstrap_root_nodes( + self, + session: AsyncSession, + run: ScheduleRuns, + *, + trace_id: str, + ) -> None: + """Persist the first runnable nodes before normal DAG advancement.""" + existing_node_id = await session.scalar( + select(ScheduleNodeRuns.node_run_id) + .where(ScheduleNodeRuns.run_id == run.run_id) + .limit(1) + ) + if existing_node_id is not None: + return + + snapshot = run.schedule_snapshot + nodes = snapshot.get("nodes", []) + target_node_ids = { + edge["target_node_id"] for edge in snapshot.get("edges", []) + } + roots = [ + node for node in nodes + if node["node_id"] not in target_node_ids + ] + max_concurrency = max(1, int(snapshot.get("max_concurrency", 1))) + selected_roots = roots[:max_concurrency] + if selected_roots: + logger.info( + "bootstrap roots: run={} roots_count={} max_concurrency={}", + run.run_id[-12:], + len(selected_roots), + max_concurrency, + ) + for node in selected_roots: + await self._dispatch_node( + session, + run=run, + node=node, + attempt_no=1, + trace_id=trace_id, + ) + # The shared session factory disables autoflush. Flush here so + # _advance_run observes the roots and cannot dispatch duplicates. + await session.flush() + + async def _advance_run( + self, + session: AsyncSession, + run: ScheduleRuns, + *, + trace_id: str, + ) -> None: + snapshot = run.schedule_snapshot + nodes = snapshot.get("nodes", []) + node_by_id = {node["node_id"]: node for node in nodes} + parents: dict[str, set[str]] = {node_id: set() for node_id in node_by_id} + for edge in snapshot.get("edges", []): + parents.setdefault(edge["target_node_id"], set()).add( + edge["source_node_id"] + ) + rows = list( + ( + await session.scalars( + select(ScheduleNodeRuns) + .where(ScheduleNodeRuns.run_id == run.run_id) + .order_by(ScheduleNodeRuns.attempt_no) + ) + ).all() + ) + latest: dict[str, ScheduleNodeRuns] = {} + for row in rows: + current = latest.get(row.node_id) + if current is None or row.attempt_no >= current.attempt_no: + latest[row.node_id] = row + + max_concurrency = max(1, int(snapshot.get("max_concurrency", 1))) + failure_policy = snapshot.get("failure_policy", "stop") + while True: + changed = False + active_count = sum( + item.node_status in {"queued", "running"} + for item in latest.values() + ) + for node in nodes: + current = latest.get(node["node_id"]) + if ( + current is not None + and current.node_status in FAILED_NODE_STATES + and current.attempt_no <= int(node.get("retry_count", 0)) + and active_count < max_concurrency + ): + retried = await self._dispatch_node( + session, + run=run, + node=node, + attempt_no=current.attempt_no + 1, + trace_id=trace_id, + delay_seconds=int(node.get("retry_interval_sec", 0)), + ) + latest[node["node_id"]] = retried + active_count += 1 + changed = True + + exhausted_failure = any( + item.node_status in FAILED_NODE_STATES + and item.attempt_no + > int(node_by_id[item.node_id].get("retry_count", 0)) + for item in latest.values() + ) + stop_all = ( + failure_policy == "stop" + and exhausted_failure + ) + for node in nodes: + node_id = node["node_id"] + if node_id in latest: + continue + parent_runs = [latest.get(parent) for parent in parents[node_id]] + # 只有 ``stop`` 策略才会在失败后跳过尚未启动的节点。 + # ``continue`` 表示“前一个节点失败也继续往后执行”，因此 + # 下游只需等待所有上游结束，不要求它们全部成功。 + parents_terminal = all( + item is not None + and item.node_status in TERMINAL_NODE_STATES + for item in parent_runs + ) + if stop_all: + skipped = ScheduleNodeRuns( + node_run_id=new_ulid(), + run_id=run.run_id, + node_id=node_id, + versions_id=node["versions_id"], + attempt_no=1, + node_status="skipped", + state_version=1, + finished_at=utcnow(), + duration_ms=0, + message="调度失败策略为 stop，未再启动", + ) + session.add(skipped) + latest[node_id] = skipped + changed = True + logger.debug( + "skipped node: run={} node={} reason={}", + run.run_id[-12:], + node_id[-12:], + "stop_policy", + ) + elif parents_terminal and active_count < max_concurrency: + dispatched = await self._dispatch_node( + session, + run=run, + node=node, + attempt_no=1, + trace_id=trace_id, + ) + latest[node_id] = dispatched + active_count += 1 + changed = True + if not changed: + break + + if nodes and len(latest) == len(nodes) and all( + item.node_status in TERMINAL_NODE_STATES + for item in latest.values() + ): + now = utcnow() + # Final run status depends on the schedule's + # ``failure_policy``. ``stop`` keeps the legacy rule — any + # non-success node fails the whole run. ``continue`` is + # more lenient: the run is a success when at least one + # root-level node succeeded and there is no remaining + # ``failed`` / ``cancelled`` / ``timed_out`` node that + # would have produced real artifacts had it run. Nodes + # marked ``skipped`` count as "decided to not run" and do + # not by themselves fail the run. + failure_policy = snapshot.get("failure_policy", "stop") + statuses = [item.node_status for item in latest.values()] + any_real_failure = any( + status in FAILED_NODE_STATES for status in statuses + ) + any_success = any( + status == "succeeded" for status in statuses + ) + if failure_policy == "continue": + # A run with mixed success/failure/skip outcomes is + # only "succeeded" when at least one node actually ran + # to completion and nothing hit a hard failure. A + # run where every node was skipped or failed is + # itself a failure. + succeeded = any_success and not any_real_failure + else: + succeeded = all( + status == "succeeded" for status in statuses + ) + logger.info( + "run finalized: run={} status={} failure_policy={} any_failure={} any_success={}", + run.run_id[-12:], + "succeeded" if succeeded else "failed", + failure_policy, + any_real_failure, + any_success, + ) + run.run_status = "succeeded" if succeeded else "failed" + run.error_code = None if succeeded else "SCHEDULE_NODE_FAILED" + run.error_message = ( + None + if succeeded + else "one or more schedule nodes did not succeed" + ) + run.finished_at = now + if run.started_at: + run.duration_ms = max( + 0, + int((now - run.started_at).total_seconds() * 1000), + ) + run.state_version += 1 + + async def _handle_node_finished( + self, + event: dict[str, Any], + message_id: str, + ) -> None: + if event.get("event_type") != NODE_FINISHED_EVENT: + raise ValueError("unexpected event type") + payload = event["payload"] + logger.info( + "node finished event: run={} node={} status={}", + payload["run_id"][-12:], + payload["node_id"][-12:], + payload.get("node_status"), + ) + async with session_scope(self.session_factory) as session: + inbox, should_process = await self._start_inbox( + session, + consumer_name="schedule-results", + event_id=event["event_id"], + message_id=message_id, + ) + if not should_process: + return + run = await session.scalar( + select(ScheduleRuns) + .where(ScheduleRuns.run_id == payload["run_id"]) + .with_for_update() + ) + if run is None: + raise ValueError("schedule run does not exist") + if run.run_status not in TERMINAL_RUN_STATES: + await self._advance_run( + session, + run, + trace_id=event["trace_id"], + ) + self._finish_inbox(inbox) + + +__all__ = ["DispatchOrchestrator"] diff --git a/schedule/src/schedule/scheduling/scheduler.py b/schedule/src/schedule/scheduling/scheduler.py new file mode 100644 index 0000000..bf57dc2 --- /dev/null +++ b/schedule/src/schedule/scheduling/scheduler.py @@ -0,0 +1,232 @@ +"""APScheduler-backed cron trigger layer. + +Owns the ``AsyncIOScheduler`` instance plus the periodic sync loop that +reconciles in-memory APScheduler jobs against the ``Schedules`` table in +MySQL. The cron tick callback delegates to ``SchedulerService.trigger_schedule`` +(via the ``on_trigger`` callable injected at construction), which posts +back to Backend; Backend then writes a ``schedule.run.requested`` Outbox +row that the orchestrator consumes. + +The two layers (this cron scheduler, the orchestrator) are decoupled +through MySQL — they share no in-memory state and survive independent +restarts. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable +from datetime import UTC, datetime, timedelta +from zoneinfo import ZoneInfo + +from apscheduler.schedulers.asyncio import AsyncIOScheduler +from apscheduler.triggers.cron import CronTrigger +from common.db import session_scope +from common.db.models import Schedules +from common.scheduler import build_sqlalchemy_jobstore +from loguru import logger +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from schedule.domain.context import naive_utc + +# APScheduler's persistent SQLAlchemy job store pickles each job. A bound +# ``CronScheduler`` method captures this instance (including SQLAlchemy engine +# state) and therefore cannot be pickled. Keep the persisted callable at +# module scope and resolve the process-local callback when the job fires. +_ACTIVE_TRIGGER: Callable[[str], Awaitable[None]] | None = None + + +async def dispatch_persisted_cron(schedule_id: str) -> None: + """Dispatch one persisted cron tick through the active service.""" + logger.debug("dispatch persisted cron: schedule={}", schedule_id[-12:]) + callback = _ACTIVE_TRIGGER + if callback is None: + raise RuntimeError("cron trigger callback is not initialized") + await callback(schedule_id) + + +class CronScheduler: + """Manages cron triggers in APScheduler, backed by a MySQL jobstore.""" + + def __init__( + self, + *, + session_factory: async_sessionmaker[AsyncSession], + database_url: str, + on_trigger: Callable[[str], Awaitable[None]], + ) -> None: + global _ACTIVE_TRIGGER + + self.session_factory = session_factory + self.scheduler = AsyncIOScheduler( + jobstores={"default": build_sqlalchemy_jobstore(database_url)}, + timezone=UTC, + ) + self._on_trigger = on_trigger + _ACTIVE_TRIGGER = on_trigger + self._sync_task: asyncio.Task[None] | None = None + # 仅在调度配置实际变化时才重置 APScheduler job。若每 5 秒都 + # reschedule，一旦恰好落在整分钟之后，就可能把本分钟的触发跳过。 + self._job_signatures: dict[str, tuple[str, str, int]] = {} + # APScheduler 的定时唤醒异常时，由 5 秒同步循环兜底。键保存的是 + # 已由兜底路径处理的“本地整分钟”，避免同一分钟重复提交。 + self._fallback_dispatched_minutes: dict[str, datetime] = {} + + def start(self) -> None: + """Start APScheduler and spawn the periodic sync loop.""" + logger.info("cron scheduler starting") + self.scheduler.start() + self._sync_task = asyncio.create_task( + self._sync_loop(), + name="scheduler-cron-sync", + ) + + async def close(self) -> None: + """Cancel the sync loop and shut APScheduler down.""" + logger.info("cron scheduler closing") + global _ACTIVE_TRIGGER + + if self._sync_task is not None: + self._sync_task.cancel() + try: + await self._sync_task + except asyncio.CancelledError: + pass + self._sync_task = None + if self.scheduler.running: + self.scheduler.shutdown(wait=False) + if _ACTIVE_TRIGGER is self._on_trigger: + _ACTIVE_TRIGGER = None + + async def trigger(self, schedule_id: str) -> None: + """APScheduler cron tick callback. + + The persisted job calls :func:`dispatch_persisted_cron`, which then + resolves this process-local callback. The + standard on_trigger is ``SchedulerService.trigger_schedule`` which + posts back to Backend; Backend then writes the Outbox row that the + orchestrator picks up. + """ + await self._on_trigger(schedule_id) + + async def _sync_loop(self) -> None: + while True: + try: + await self._sync_once() + logger.debug("cron sync tick ok") + except asyncio.CancelledError: + raise + except Exception: + logger.exception("cron job synchronization failed") + await asyncio.sleep(5) + + async def _sync_once(self) -> None: + """Reconcile APScheduler jobs against ``Schedules.cron_expression``. + + - adds jobs for enabled cron schedules present in MySQL + - removes jobs whose schedule has been disabled / soft-deleted + - updates ``Schedules.next_run_at`` from the Cron expression itself + """ + due_schedule_ids: list[str] = [] + async with session_scope(self.session_factory) as session: + schedules = list( + ( + await session.scalars( + select(Schedules).where( + Schedules.deleted_at.is_(None), + Schedules.enabled == 1, + Schedules.trigger_type == "cron", + Schedules.cron_expression.is_not(None), + ) + ) + ).all() + ) + active_job_ids: set[str] = set() + added_count = 0 + updated_count = 0 + for item in schedules: + job_id = f"schedule:{item.schedule_id}" + active_job_ids.add(job_id) + expression = (item.cron_expression or "").strip() + max_instances = max(1, item.max_concurrency) + signature = (expression, item.timezone, max_instances) + trigger = CronTrigger.from_crontab( + expression, + timezone=ZoneInfo(item.timezone), + ) + now = datetime.now(ZoneInfo(item.timezone)) + minute = now.replace(second=0, microsecond=0) + job_changed = False + if self.scheduler.get_job(job_id) is None: + self.scheduler.add_job( + dispatch_persisted_cron, + trigger=trigger, + args=[item.schedule_id], + id=job_id, + replace_existing=True, + coalesce=True, + max_instances=max_instances, + misfire_grace_time=60, + ) + added_count += 1 + job_changed = True + elif self._job_signatures.get(job_id) != signature: + # 服务重启后的首次同步也会走这里，确保持久化 job 与 + # 数据库当前配置一致；之后配置不变时保留原定时点。 + self.scheduler.reschedule_job(job_id, trigger=trigger) + self.scheduler.modify_job( + job_id, + max_instances=max_instances, + ) + updated_count += 1 + job_changed = True + self._job_signatures[job_id] = signature + # 以 CronTrigger 本身计算下次执行时间，不依赖 APScheduler 的 + # 内部唤醒状态；页面展示的「下次执行」也因此保持准确。 + item.next_run_at = naive_utc( + trigger.get_next_fire_time(None, now) + ) + + # 首次观察或刚修改表达式时，从下一个整分钟才开始兜底，符合 + # Cron 的常规语义，避免用户在本分钟中途保存后立刻多跑一次。 + if job_changed or job_id not in self._fallback_dispatched_minutes: + self._fallback_dispatched_minutes[job_id] = minute + + # 正常情况下 APScheduler 会在整分钟回调。实测其偶发漏唤醒时， + # 这里每 5 秒检查一次当前分钟是否命中表达式，并补发一次。 + due_at = trigger.get_next_fire_time( + minute - timedelta(minutes=1), + minute, + ) + if ( + due_at == minute + and self._fallback_dispatched_minutes.get(job_id) != minute + ): + self._fallback_dispatched_minutes[job_id] = minute + due_schedule_ids.append(item.schedule_id) + removed_count = 0 + for job in self.scheduler.get_jobs(): + if ( + job.id.startswith("schedule:") + and job.id not in active_job_ids + ): + self.scheduler.remove_job(job.id) + self._job_signatures.pop(job.id, None) + self._fallback_dispatched_minutes.pop(job.id, None) + removed_count += 1 + if added_count or updated_count or removed_count: + logger.info( + "cron sync reconciled: added={} updated={} removed={}", + added_count, + updated_count, + removed_count, + ) + # 在数据库同步事务提交后再创建运行记录，避免两个会话同时读取调度方案 + # 时发生不必要的锁等待。重复回调由运行记录的幂等键自动去重。 + for schedule_id in due_schedule_ids: + logger.debug("cron fallback dispatch: schedule={}", schedule_id[-12:]) + await self._on_trigger(schedule_id) + + +__all__ = ["CronScheduler"] diff --git a/schedule/src/schedule/service.py b/schedule/src/schedule/service.py index 86c2abb..7044ade 100644 --- a/schedule/src/schedule/service.py +++ b/schedule/src/schedule/service.py @@ -36,8 +36,8 @@ from common.storage import create_storage from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker -from schedule.orchestrator import DispatchOrchestrator -from schedule.scheduler import CronScheduler +from schedule.scheduling.orchestrator import DispatchOrchestrator +from schedule.scheduling.scheduler import CronScheduler from schedule.worker import NodeExecutor diff --git a/schedule/tests/test_janitor.py b/schedule/tests/test_janitor.py index fcb632d..386c9c4 100644 --- a/schedule/tests/test_janitor.py +++ b/schedule/tests/test_janitor.py @@ -23,7 +23,7 @@ import pytest from common.db.models import ScheduleNodeRuns from common.eventing import utcnow -from schedule.orchestrator import NODE_FINISHED_EVENT, DispatchOrchestrator +from schedule.scheduling.orchestrator import NODE_FINISHED_EVENT, DispatchOrchestrator def _make_orchestrator() -> DispatchOrchestrator: @@ -195,7 +195,7 @@ async def test_reap_kills_running_row_past_deadline() -> None: fake_session.add = MagicMock() with patch( - "schedule.orchestrator.session_scope", + "schedule.scheduling.orchestrator.session_scope", return_value=_open_session_scope(fake_session), ): killed = await orch._reap_stuck_node_runs() @@ -223,7 +223,7 @@ async def test_reap_skips_healthy_row() -> None: fake_session.add = MagicMock() with patch( - "schedule.orchestrator.session_scope", + "schedule.scheduling.orchestrator.session_scope", return_value=_open_session_scope(fake_session), ): killed = await orch._reap_stuck_node_runs() @@ -266,7 +266,7 @@ async def test_reap_processes_multiple_rows_in_one_pass() -> None: fake_session.add = MagicMock() with patch( - "schedule.orchestrator.session_scope", + "schedule.scheduling.orchestrator.session_scope", return_value=_open_session_scope(fake_session), ): killed = await orch._reap_stuck_node_runs() @@ -289,7 +289,7 @@ async def test_janitor_loop_propagates_cancellation() -> None: raise asyncio.CancelledError with patch( - "schedule.orchestrator.asyncio.sleep", + "schedule.scheduling.orchestrator.asyncio.sleep", side_effect=cancel_on_sleep, ): with pytest.raises(asyncio.CancelledError): @@ -321,7 +321,7 @@ async def test_janitor_loop_continues_after_reap_exception() -> None: raise asyncio.CancelledError with patch( - "schedule.orchestrator.asyncio.sleep", + "schedule.scheduling.orchestrator.asyncio.sleep", side_effect=_count_sleeps, ): with pytest.raises(asyncio.CancelledError):