"""Supervision for the long-lived tasks a pipeline starts. A node's subscription, schedule or poll loop is an `asyncio.Task`, and a task that raises is simply gone — the node stays listed as running while nothing listens any more. The supervisor restarts those tasks with a growing delay and gives up on a flow that keeps failing, because a flow crash-looping every second is worse than a flow that is visibly stopped. Deliberately a plain task registry rather than a `TaskGroup`: a group cancels its siblings when one member fails, which is the opposite of what supervision means here. """ from __future__ import annotations import asyncio import logging import time from collections import deque from collections.abc import Callable, Coroutine from typing import Any from fluksio.flow.events import EventBus logger = logging.getLogger(__name__) BACKOFF = (1.0, 5.0, 30.0, 60.0) # A flow that burns through this many restarts in the window is not going to # recover by being restarted again. FAILURE_BUDGET = 5 FAILURE_WINDOW = 300.0 # How long a cancelled task gets to notice. A loop that is still waiting after # this is not going to stop on its own — a client closing a socket the broker # no longer answers on is the case seen in the wild — and the rebuild asking # for it must not wait on that forever. CANCEL_GRACE = 5.0 TaskFactory = Callable[[], Coroutine[Any, Any, None]] class Supervisor: """Keeps the pipeline's background tasks alive, or admits it cannot.""" def __init__(self, events: EventBus | None = None) -> None: self._events = events self._tasks: dict[str, asyncio.Task[None]] = {} #: Which flow each task belongs to. Recorded rather than read off the #: task's name, because names are a flow and a node joined by a dot and #: matching on that prefix would let 'hea' cancel 'heating'. self._flows: dict[str, str] = {} self._failures: dict[str, deque[float]] = {} self.quarantined: set[str] = set() def spawn(self, name: str, flow: str, factory: TaskFactory) -> None: """Run `factory()` and keep running it until told to stop.""" if name in self._tasks: return self._flows[name] = flow self._tasks[name] = asyncio.create_task( self._supervise(name, flow, factory), name=f"supervised:{name}" ) async def _supervise(self, name: str, flow: str, factory: TaskFactory) -> None: attempt = 0 while True: try: await factory() except asyncio.CancelledError: raise except Exception as exc: if not self._record_failure(name, flow, exc): return else: # A clean return means the loop decided it was done. return delay = BACKOFF[min(attempt, len(BACKOFF) - 1)] attempt += 1 await asyncio.sleep(delay) def _record_failure(self, name: str, flow: str, exc: Exception) -> bool: """Note the crash; False when the flow has spent its budget.""" logger.warning("Supervised task '%s' crashed: %s", name, exc, exc_info=True) self._publish( { "type": "task_crashed", "task": name, "flow": flow, "error": f"{type(exc).__name__}: {exc}", "ts": time.time(), } ) now = time.monotonic() window = self._failures.setdefault(flow, deque()) window.append(now) while window and window[0] < now - FAILURE_WINDOW: window.popleft() if len(window) < FAILURE_BUDGET: return True logger.error( "Flow '%s' crashed %d times in %.0fs — quarantined", flow, len(window), FAILURE_WINDOW, ) self.quarantined.add(flow) self._publish( { "type": "flow_quarantined", "flow": flow, "error": f"{type(exc).__name__}: {exc}", "ts": time.time(), } ) return False async def cancel_all(self) -> None: """Stop supervising. Idempotent, and safe to call mid-restart.""" tasks = list(self._tasks.values()) self._tasks.clear() self._flows.clear() await self._cancel(tasks) async def cancel_flow(self, flow: str) -> None: """Stop one flow's supervised tasks and give it a clean slate. Rebuilding the whole pipeline throws the supervisor away and builds another, so a flow quarantined by the last build gets another chance. Rebuilding one flow has to say the same thing about that flow alone, or every other flow's quarantine would go with it. """ names = [name for name, owner in self._flows.items() if owner == flow] tasks = [self._tasks.pop(name) for name in names if name in self._tasks] for name in names: del self._flows[name] await self._cancel(tasks) # The clean slate: whatever this flow spent before, the build that # follows starts its budget again. self.quarantined.discard(flow) self._failures.pop(flow, None) async def _cancel(self, tasks: list[asyncio.Task[None]]) -> None: """Ask these tasks to stop, and wait no longer than the grace period.""" for task in tasks: task.cancel() if not tasks: return # `wait` hands back what is still going instead of waiting on it, and # lets a cancellation aimed at *this* coroutine through — the # `except CancelledError` it replaces swallowed that, which left # whoever asked for the teardown holding their lock and unkillable. done, pending = await asyncio.wait(tasks, timeout=CANCEL_GRACE) for task in done: if not task.cancelled(): # Retrieved so a crash on the way out is not reported at exit; # the supervisor has already said what it was. task.exception() for task in pending: # Cancelled once and still running means its shutdown is waiting on # something that is not answering. A second cancellation interrupts # that wait; whether it takes is no longer the rebuild's problem. task.cancel() logger.warning( "Supervised task '%s' did not stop within %.0fs — abandoned", task.get_name(), CANCEL_GRACE, ) def _publish(self, event: dict[str, Any]) -> None: if self._events is not None: self._events.publish(event)