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