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]] = {}
        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._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()
        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)