From 26be90f60ddba190d4f2b11a9b772e572eac2655 Mon Sep 17 00:00:00 2001 From: Rajat Ahuja Date: Wed, 26 Aug 2026 14:49:06 -0400 Subject: [PATCH] fix(deriver): bound the shutdown drain and release claims before exit Shutdown waited on asyncio.gather(*active_tasks) with no timeout, and released work-unit claims only via the finally in initialize(). Both assume the process gets as long as it needs. It does not. A platform can terminate the process on a budget shorter than a single work unit takes -- a deriver LLM call alone can run for tens of seconds -- and when that budget expires the process is killed outright. The unbounded wait is then still in progress, initialize() never returns, its finally never runs, and the claims are left held. The work is not lost, but nothing may touch it until cleanup_stale_work_units notices, which takes STALE_SESSION_TIMEOUT_MINUTES plus a jittered poll interval. So shutdown is a race to hand the claims back, not an opportunity to finish the work. Two changes: Drain is bounded by SHUTDOWN_DRAIN_TIMEOUT_SECONDS (default 8). Work that fits in the budget still completes; work that does not is cancelled rather than waited on, since the work is reclaimable and the remaining budget is not. cleanup() is called at the end of shutdown() rather than relying on the finally in initialize(). Claims are given up while the process is still alive, so another worker repicks the work on its next poll instead of waiting out the stale sweep. cleanup() clears worker_ownership, so the existing call in the finally becomes a no-op and the ordering is idempotent. The drain deliberately runs before the release rather than after. Releasing first would leave a window in which a task that completes during shutdown has already had its claim taken by another worker, and the same queue items get derived twice. Cancelling first and releasing after keeps that window closed, and both steps fit inside a short budget. Cancelling also keeps the error reporting honest without special-casing it. asyncio.CancelledError derives from BaseException rather than Exception, so an abandoned task passes straight through the `except Exception` handlers that report to Sentry. Shutdown stops generating spurious errors on its own, while a task that genuinely fails during the drain, or a cleanup() that cannot reach the database, still reports as it should. Tests cover a task that outlives the budget being abandoned with cleanup still awaited, and a task that fits in the budget still finishing. --- src/config.py | 5 ++ src/deriver/queue_manager.py | 26 ++++++++-- tests/deriver/test_shutdown_drain.py | 73 ++++++++++++++++++++++++++++ 3 files changed, 101 insertions(+), 3 deletions(-) create mode 100644 tests/deriver/test_shutdown_drain.py diff --git a/src/config.py b/src/config.py index 993f9bfb..290531ce 100644 --- a/src/config.py +++ b/src/config.py @@ -908,6 +908,11 @@ class DeriverSettings(HonchoSettings): # to 0.0 to disable. POLLING_JITTER_RATIO: Annotated[float, Field(default=0.5, ge=0.0, le=1.0)] = 0.5 STALE_SESSION_TIMEOUT_MINUTES: Annotated[int, Field(default=5, gt=0, le=1440)] = 5 + # How long shutdown waits for in-flight work before abandoning it. Must + # leave room in the shutdown budget for the claim release that follows. + SHUTDOWN_DRAIN_TIMEOUT_SECONDS: Annotated[ + float, Field(default=8.0, ge=0.0, le=600.0) + ] = 8.0 # Minimum (jittered) spacing between stale-work-unit cleanup runs STALE_WORK_UNIT_CLEANUP_INTERVAL_SECONDS: Annotated[ float, Field(default=60.0, ge=0.0, le=3600.0) diff --git a/src/deriver/queue_manager.py b/src/deriver/queue_manager.py index 1c19c131..8acf250d 100644 --- a/src/deriver/queue_manager.py +++ b/src/deriver/queue_manager.py @@ -224,7 +224,12 @@ class QueueManager: await self.cleanup() async def shutdown(self, sig: signal.Signals) -> None: - """Handle graceful shutdown""" + """Hand work-unit claims back before the process dies. + + The shutdown budget can be shorter than one work unit takes, so finishing + the work is not the goal. A claim still held at exit strands its work + until the stale sweep reclaims it. + """ logger.info(f"Received exit signal {sig.name}...") self.shutdown_event.set() @@ -236,9 +241,24 @@ class QueueManager: if self.active_tasks: logger.info( - f"Waiting for {len(self.active_tasks)} active tasks to complete..." + "Waiting up to %ss for %d active tasks to complete...", + settings.DERIVER.SHUTDOWN_DRAIN_TIMEOUT_SECONDS, + len(self.active_tasks), ) - await asyncio.gather(*self.active_tasks, return_exceptions=True) + drain = asyncio.gather(*self.active_tasks, return_exceptions=True) + try: + await asyncio.wait_for( + drain, timeout=settings.DERIVER.SHUTDOWN_DRAIN_TIMEOUT_SECONDS + ) + except (TimeoutError, asyncio.TimeoutError): + # wait_for already cancelled the gather. The work is + # reclaimable; the remaining budget is not. + logger.warning( + "Drain timed out; abandoning in-flight work units for reclaim" + ) + + # Not left to the finally in initialize(): a hard kill skips it. + await self.cleanup() async def cleanup(self) -> None: """Clean up owned work units""" diff --git a/tests/deriver/test_shutdown_drain.py b/tests/deriver/test_shutdown_drain.py new file mode 100644 index 00000000..ba89cea4 --- /dev/null +++ b/tests/deriver/test_shutdown_drain.py @@ -0,0 +1,73 @@ +"""Shutdown must bound the drain and give claims back before the process dies. + +A claim still held at exit strands its work until the stale sweep reclaims it. +""" + +import asyncio +import signal +from unittest.mock import AsyncMock, patch + +import pytest + +from src.deriver.queue_manager import QueueManager + + +@pytest.mark.asyncio +async def test_drain_is_bounded_and_claims_are_released() -> None: + """A task outliving the budget is abandoned; cleanup still runs.""" + qm = QueueManager() + + started = asyncio.Event() + + async def never_finishes() -> None: + started.set() + await asyncio.sleep(3600) + + task = asyncio.create_task(never_finishes()) + qm.add_task(task) + await started.wait() + + with ( + patch.object(qm, "cleanup", new=AsyncMock()) as cleanup, + patch.object(qm.dream_scheduler, "shutdown", new=AsyncMock()), + patch.object(qm.reconciler_scheduler, "shutdown", new=AsyncMock()), + patch("src.deriver.queue_manager.settings") as settings, + ): + settings.DERIVER.SHUTDOWN_DRAIN_TIMEOUT_SECONDS = 0.05 + + loop = asyncio.get_running_loop() + began = loop.time() + await qm.shutdown(signal.SIGTERM) + elapsed = loop.time() - began + + assert elapsed < 2, f"shutdown blocked on the stuck task for {elapsed:.1f}s" + assert qm.shutdown_event.is_set() + cleanup.assert_awaited_once() # claims released before exit + + assert task.cancelled() or task.done() + + +@pytest.mark.asyncio +async def test_fast_work_still_completes_before_claims_are_released() -> None: + """Bounding the drain must not cut short work that fits.""" + qm = QueueManager() + finished = False + + async def quick() -> None: + nonlocal finished + await asyncio.sleep(0.01) + finished = True + + qm.add_task(asyncio.create_task(quick())) + + with ( + patch.object(qm, "cleanup", new=AsyncMock()) as cleanup, + patch.object(qm.dream_scheduler, "shutdown", new=AsyncMock()), + patch.object(qm.reconciler_scheduler, "shutdown", new=AsyncMock()), + patch("src.deriver.queue_manager.settings") as settings, + ): + settings.DERIVER.SHUTDOWN_DRAIN_TIMEOUT_SECONDS = 5.0 + await qm.shutdown(signal.SIGTERM) + + assert finished + cleanup.assert_awaited_once()