app/backend/fluksio/flow/supervision.py

"""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)