From 3bb23fbd3ec5ced5a4c609c23a005a4fd1cc62ca Mon Sep 17 00:00:00 2001 From: stroblme Date: Sun, 23 Aug 2026 17:43:15 +0200 Subject: [PATCH] Let the supervisor stop one flow's tasks and clear its quarantine Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01StpRc2C6au1WJ1EUU7fsfu --- backend/fluksio/flow/supervision.py | 28 ++++++++++++++++ backend/tests/flow/test_supervision.py | 46 ++++++++++++++++++++++++++ 2 files changed, 74 insertions(+) diff --git a/backend/fluksio/flow/supervision.py b/backend/fluksio/flow/supervision.py index 5a1404f..2e72678 100644 --- a/backend/fluksio/flow/supervision.py +++ b/backend/fluksio/flow/supervision.py @@ -44,6 +44,10 @@ class Supervisor: 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() @@ -51,6 +55,7 @@ class Supervisor: """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}" ) @@ -116,6 +121,29 @@ class Supervisor: """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: diff --git a/backend/tests/flow/test_supervision.py b/backend/tests/flow/test_supervision.py index c0b9e00..4424c52 100644 --- a/backend/tests/flow/test_supervision.py +++ b/backend/tests/flow/test_supervision.py @@ -124,3 +124,49 @@ def test_cancel_all_is_idempotent(): await supervisor.cancel_all() asyncio.run(scenario()) + + +def test_cancelling_one_flow_leaves_another_flows_tasks_running(): + """Rebuilding one flow must not take every other flow's loops with it.""" + + async def loop() -> None: + await asyncio.sleep(3600) + + async def scenario() -> None: + supervisor = Supervisor() + supervisor.spawn("gone.sub", "gone", loop) + supervisor.spawn("fine.sub", "fine", loop) + await _settle() + + survivor = supervisor._tasks["fine.sub"] + await supervisor.cancel_flow("gone") + + assert list(supervisor._tasks) == ["fine.sub"] + assert not survivor.done() + await supervisor.cancel_all() + + asyncio.run(scenario()) + + +def test_cancelling_one_flow_lifts_only_its_quarantine(): + async def failing() -> None: + raise RuntimeError("nope") + + async def scenario() -> None: + supervisor = Supervisor() + supervisor.spawn("a", "broken", failing) + supervisor.spawn("b", "also-broken", failing) + await _settle() + + assert supervisor.quarantined == {"broken", "also-broken"} + + await supervisor.cancel_flow("broken") + + assert supervisor.quarantined == {"also-broken"} + # And its budget, or the flow would be quarantined again on the first + # crash after the rebuild that was meant to fix it. + assert "broken" not in supervisor._failures + assert "also-broken" in supervisor._failures + await supervisor.cancel_all() + + asyncio.run(scenario())