From 2e3bfb7bbbca2748c6becde8a205ad8b33aa0127 Mon Sep 17 00:00:00 2001 From: Renn F Date: Sun, 28 Jun 2026 18:05:41 +0200 Subject: [PATCH] [F070] drain fire-and-forget _bg_tasks on shutdown (bounded, data-preserving) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Orchestrator.stop() cancelled only the named loop tasks + agents, then returned, abandoning in-flight _schedule_bg work. An in-flight _persist_respawn_record upsert dropped at shutdown meant the last few gate-mutation strikes never reached the DB; restore_respawn_tracker() on the next start repopulated a stale lower count and the dispatcher re-burned the full 4-spawn strike threshold against a still-wedged task — the exact re-burn the durable tracker exists to stop. Audit-log writes (load-bearing for cycle-time/rework metrics) were similarly dropped. Add _drain_bg_tasks(): bounded wait (5s default) lets short DB writes commit before exit (data preserved), then cancels any stuck task past the deadline so a hang can't wedge shutdown. return_exceptions=True so one failing bg task doesn't crash the drain. Wrap the stop_agent loop in try/except + logger.exception so one bad agent can't skip the drain (re-introducing the data-loss tail). Floor test pins the deadline >= 3s so a too-short change can't silently drop a legitimate slow write. --- roboco/runtime/orchestrator.py | 51 +++++- .../test_orchestrator_shutdown_drain.py | 150 ++++++++++++++++++ 2 files changed, 199 insertions(+), 2 deletions(-) create mode 100644 tests/unit/runtime/test_orchestrator_shutdown_drain.py diff --git a/roboco/runtime/orchestrator.py b/roboco/runtime/orchestrator.py index 89ab43ef..7a80b83c 100644 --- a/roboco/runtime/orchestrator.py +++ b/roboco/runtime/orchestrator.py @@ -98,6 +98,14 @@ _PROBE_TIMEOUT_SECONDS = 10.0 # short enough that a hang degrades one tick, not the whole fleet. _DOCKER_INSPECT_TIMEOUT_SECONDS = 10.0 _DOCKER_EXEC_TIMEOUT_SECONDS = 30.0 +# Deadline for draining fire-and-forget ``_bg_tasks`` on shutdown. Short DB +# writes (a respawn_tracker upsert, an audit-log row) finish before the +# process exits — preserving the durable PM-respawn counter and the +# metrics-bearing audit trail — while a stuck task can't hang shutdown: past +# this deadline the still-pending tasks are cancelled. Generous enough that a +# legitimate slow write under load commits rather than being dropped (the +# exact data-loss tail the durable tracker exists to prevent). +_SHUTDOWN_DRAIN_TIMEOUT_SECONDS = 5.0 _HTTP_TOO_MANY_REQUESTS = 429 _HTTP_OK = 200 _HTTP_MULTIPLE_CHOICES = 300 # first non-2xx status; 2xx == [_HTTP_OK, this) @@ -913,6 +921,32 @@ class AgentOrchestrator: with contextlib.suppress(asyncio.CancelledError): await task + async def _drain_bg_tasks(self) -> None: + """Let fire-and-forget ``_bg_tasks`` finish before the process exits. + + Short DB writes (a respawn_tracker upsert, an audit-log row) get a + bounded window to commit — preserving the durable PM-respawn counter + and the metrics-bearing audit trail — while a stuck task can't hang + shutdown: past ``_SHUTDOWN_DRAIN_TIMEOUT_SECONDS`` the still-pending + tasks are cancelled. ``return_exceptions=True`` so one failing bg task + doesn't crash the drain (a failed write already degraded to in-memory; + logging it here would just be noise). No-op when nothing is pending. + """ + pending = [t for t in self._bg_tasks if not t.done()] + if not pending: + return + try: + await asyncio.wait_for( + asyncio.gather(*pending, return_exceptions=True), + timeout=_SHUTDOWN_DRAIN_TIMEOUT_SECONDS, + ) + except TimeoutError: + for task in pending: + if not task.done(): + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await asyncio.gather(*pending, return_exceptions=True) + async def stop(self) -> None: """Stop the orchestrator and all agents.""" self._running = False @@ -932,9 +966,22 @@ class AgentOrchestrator: ): await self._cancel_background_task(task) - # Stop all agents + # Stop all agents. One agent's stop error must not skip the drain + # below — that would re-introduce the data-loss tail for every in-flight + # bg write, so log-and-continue rather than propagate. for agent_id in list(self._instances.keys()): - await self.stop_agent(agent_id) + try: + await self.stop_agent(agent_id) + except Exception: + logger.exception( + "stop_agent raised during shutdown; continuing to drain", + agent_id=agent_id, + ) + + # Drain fire-and-forget bg writes so short DB commits finish before the + # process exits (respawn_tracker upserts, audit-log rows). Bounded so a + # stuck task can't hang shutdown — it is cancelled past the deadline. + await self._drain_bg_tasks() logger.info("Orchestrator stopped") diff --git a/tests/unit/runtime/test_orchestrator_shutdown_drain.py b/tests/unit/runtime/test_orchestrator_shutdown_drain.py new file mode 100644 index 00000000..628037a4 --- /dev/null +++ b/tests/unit/runtime/test_orchestrator_shutdown_drain.py @@ -0,0 +1,150 @@ +"""F070 — fire-and-forget ``_bg_tasks`` (respawn_tracker upserts, audit-log +writes, intake first-message delivery) were never cancelled or drained on +shutdown. ``Orchestrator.stop()`` cancelled only the named loop tasks and the +agents, then returned, abandoning any in-flight ``_schedule_bg`` work. + +The data-loss tail: an in-flight ``_persist_respawn_record`` upsert dropped at +shutdown means the last few gate-mutation strikes never reach the DB. The +in-memory counter dies with the process; ``restore_respawn_tracker()`` on the +next start repopulates a stale lower count and the dispatcher re-burns the +full strike threshold (4 spawns) against a still-wedged task — the exact +re-burn the durable tracker exists to stop. Audit-log writes (load-bearing for +the cycle-time / rework metrics) are similarly dropped. + +The fix DRAINs ``_bg_tasks`` with a bounded timeout on shutdown — short DB +writes finish before the process exits (data preserved), while a stuck task +can't hang shutdown (it is cancelled once the drain deadline passes). Cancels +outright would lose the data (the opposite of the goal), so the drain tries to +let work complete first. The ``stop_agent`` loop is also wrapped so one agent's +stop error can't skip the drain (which would still drop the data). +""" + +from __future__ import annotations + +import asyncio +from typing import Any +from unittest.mock import patch + +import pytest +from roboco.models.runtime import AgentInstance +from roboco.runtime.orchestrator import ( + _SHUTDOWN_DRAIN_TIMEOUT_SECONDS, + AgentOrchestrator, +) + +# Floor encoding the logical-regression guard: a drain deadline below this +# would risk dropping a legitimate short DB write (an upsert that needs a +# second under load) before it commits — the exact data loss this fix targets. +# Named (not magic) for ruff PLR2004. +_MIN_DRAIN_TIMEOUT = 3.0 + + +def _make_orchestrator() -> AgentOrchestrator: + """AgentOrchestrator with constructor I/O skipped; stop() deps ready. + + ``stop()`` cancels the named loop tasks (all None here → no-op) and the + agents in ``_instances`` (empty here), then must drain ``_bg_tasks``. + """ + with patch.object(AgentOrchestrator, "__init__", return_value=None): + orch = AgentOrchestrator.__new__(AgentOrchestrator) + orch._instances = {} + orch._bg_tasks = set() + # Every named background loop ``stop()`` cancels — None makes each a no-op + # so the test exercises ONLY the _bg_tasks drain. + for attr in ( + "_health_task", + "_dispatcher_task", + "_sweeper_task", + "_rate_limit_probe_task", + "_strategy_engine_task", + "_external_pr_poll_task", + "_self_heal_task", + "_ci_watch_task", + "_dep_update_task", + "_release_manager_task", + ): + setattr(orch, attr, None) + return orch + + +def test_shutdown_drain_timeout_is_named_module_constant() -> None: + assert isinstance(_SHUTDOWN_DRAIN_TIMEOUT_SECONDS, int | float) + assert _SHUTDOWN_DRAIN_TIMEOUT_SECONDS > 0 + + +def test_shutdown_drain_timeout_is_generous() -> None: + """A short DB upsert under load can legitimately take a moment; the drain + deadline must not drop it. This guards the logical regression: a too-short + drain would silently lose the very writes it exists to preserve.""" + assert _SHUTDOWN_DRAIN_TIMEOUT_SECONDS >= _MIN_DRAIN_TIMEOUT + + +@pytest.mark.asyncio +async def test_stop_drains_completing_bg_task_before_returning() -> None: + """A bg task that finishes quickly MUST complete (its side effect observed) + before ``stop()`` returns. Without the drain, ``stop()`` returns immediately + and the task is abandoned mid-flight — the data-loss tail.""" + orch = _make_orchestrator() + ran: list[bool] = [] + + async def _completes() -> None: + await asyncio.sleep(0.01) + ran.append(True) + + orch._bg_tasks.add(asyncio.create_task(_completes())) + + await asyncio.wait_for(orch.stop(), timeout=5.0) + + assert ran == [True], "completing bg task was abandoned by stop()" + + +@pytest.mark.asyncio +async def test_stop_does_not_hang_on_stuck_bg_task( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A bg task that never completes MUST NOT hang shutdown past the drain + deadline — it is cancelled once the drain times out. Without the drain, + a stuck bg task would let ``stop()`` (and thus the process) hang forever. + + Deterministic: the drain deadline is patched tiny so a bounded fail-close is + asserted in well under a second, never relying on the real 5s default.""" + monkeypatch.setattr( + "roboco.runtime.orchestrator._SHUTDOWN_DRAIN_TIMEOUT_SECONDS", 0.05 + ) + orch = _make_orchestrator() + + async def _hangs() -> None: + await asyncio.Future() # never resolves + + stuck = asyncio.create_task(_hangs()) + orch._bg_tasks.add(stuck) + + await asyncio.wait_for(orch.stop(), timeout=2.0) + + assert stuck.cancelled(), "stuck bg task was not cancelled by the drain" + + +@pytest.mark.asyncio +async def test_stop_failing_agent_does_not_skip_drain() -> None: + """If one agent's ``stop_agent`` raises, the drain must still run — + otherwise a single bad agent would re-introduce the data-loss tail for every + in-flight bg write. The completing bg task should still finish.""" + orch = _make_orchestrator() + orch._instances["bad-agent"] = AgentInstance(agent_id="bad-agent") + + async def _raises(_aid: str, **_kw: Any) -> None: + raise RuntimeError("boom") + + patch.object(orch, "stop_agent", _raises).start() + + ran: list[bool] = [] + + async def _completes() -> None: + await asyncio.sleep(0.01) + ran.append(True) + + orch._bg_tasks.add(asyncio.create_task(_completes())) + + await asyncio.wait_for(orch.stop(), timeout=5.0) + + assert ran == [True], "failing stop_agent skipped the drain (data lost)"