diff --git a/backend/app/gateway/app.py b/backend/app/gateway/app.py index f7af51a54c..a4efca879b 100644 --- a/backend/app/gateway/app.py +++ b/backend/app/gateway/app.py @@ -394,6 +394,28 @@ async def _shutdown_startup_trash_sweep(app: FastAPI) -> None: logger.exception("Startup trash sweep failed during shutdown") +async def _shutdown_scheduled_task_service(app: FastAPI) -> None: + """Bound scheduler stop so a stuck poll cannot block Gateway exit. + + ``asyncio.wait_for`` cancels the ``stop()`` coroutine at the deadline. + If the scheduler is already stuck inside an iteration, its owned poll task + may remain pending until event-loop teardown; this helper bounds Gateway + shutdown latency rather than promising completion of stuck scheduler work. + """ + service = getattr(app.state, "scheduled_task_service", None) + if service is None: + return + try: + await asyncio.wait_for(service.stop(), timeout=_SHUTDOWN_HOOK_TIMEOUT_SECONDS) + except TimeoutError: + logger.warning( + "Scheduled task service shutdown exceeded %.1fs; proceeding with worker exit.", + _SHUTDOWN_HOOK_TIMEOUT_SECONDS, + ) + except Exception: + logger.exception("Failed to stop scheduled task service") + + def _scheduled_task_notification_repos(startup_config: AppConfig): """Return ``(connection_repo, notification_repo)`` for the scheduled-run outbox (issue #4254). @@ -810,11 +832,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: except Exception: logger.exception("Failed to stop channel service") - if getattr(app.state, "scheduled_task_service", None) is not None: - try: - await app.state.scheduled_task_service.stop() - except Exception: - logger.exception("Failed to stop scheduled task service") + await _shutdown_scheduled_task_service(app) if getattr(app.state, "mcp_task_service", None) is not None: app.state.mcp_tasks_available = False diff --git a/backend/tests/_gateway_shutdown_budget.py b/backend/tests/_gateway_shutdown_budget.py index 5538873077..9ce3ee41f3 100644 --- a/backend/tests/_gateway_shutdown_budget.py +++ b/backend/tests/_gateway_shutdown_budget.py @@ -23,6 +23,7 @@ HOOKS_BOUNDED_BY_SHUTDOWN_HOOK_TIMEOUT: tuple[str, ...] = ( "startup trash sweep", "notification delivery worker stop", "channel service stop", + "scheduled task service stop", "browser session close", "MCP session pool close", ) @@ -61,7 +62,7 @@ def memory_flush_seconds() -> float: def lifespan_shutdown_seconds() -> float: """Worst-case lifespan shutdown: every bounded hook + retrieval-warm wait + run drain + memory flush. - Unbounded steps (OIDC close, scheduler / task-service stops, backend close) + Unbounded steps (OIDC close, MCP/subagent task-service stops, backend close) are expected to be quick and are not modeled; deployments keep slack on top of this figure for them. """ diff --git a/backend/tests/test_scheduler_shutdown_timeout.py b/backend/tests/test_scheduler_shutdown_timeout.py new file mode 100644 index 0000000000..d3d091b6b8 --- /dev/null +++ b/backend/tests/test_scheduler_shutdown_timeout.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +import asyncio +import inspect +import logging +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import FastAPI + +from app.gateway import app as gateway_app + + +@pytest.mark.asyncio +async def test_scheduled_task_service_shutdown_is_bounded(monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture) -> None: + async def hang() -> None: + await asyncio.sleep(3600) + + app = FastAPI() + service = MagicMock() + service.stop = AsyncMock(side_effect=hang) + app.state.scheduled_task_service = service + monkeypatch.setattr(gateway_app, "_SHUTDOWN_HOOK_TIMEOUT_SECONDS", 0.01) + + started = asyncio.get_running_loop().time() + with caplog.at_level(logging.WARNING, logger="app.gateway.app"): + await gateway_app._shutdown_scheduled_task_service(app) + elapsed = asyncio.get_running_loop().time() - started + + assert elapsed < 0.5 + service.stop.assert_awaited_once() + assert "Scheduled task service shutdown exceeded" in caplog.text + + +def test_lifespan_uses_bounded_scheduled_task_shutdown() -> None: + source = inspect.getsource(gateway_app.lifespan) + assert "await _shutdown_scheduled_task_service(app)" in source diff --git a/deploy/helm/deer-flow/README.md b/deploy/helm/deer-flow/README.md index 3f3552494b..464abfe829 100644 --- a/deploy/helm/deer-flow/README.md +++ b/deploy/helm/deer-flow/README.md @@ -270,7 +270,7 @@ kubectl -n deer-flow exec deploy/deer-flow-provisioner -- curl -s localhost:8002 `%40`). The chart uses an external `databaseUrl` verbatim and does not rewrite the DSN in a user-managed Secret. -- **Graceful shutdown & memory drain.** The gateway pod sets `terminationGracePeriodSeconds` (default 90s, overridable via `gateway.terminationGracePeriodSeconds`) plus an optional `preStop` sleep (`gateway.preStopSleepSeconds`, default 5s), and bounds uvicorn's `--timeout-graceful-shutdown` (`gateway.uvicornGracefulShutdownSeconds`, default 10s) so an idle SSE connection cannot hold up lifespan shutdown indefinitely. The grace period MUST exceed the Gateway's graceful-shutdown work — the preStop sleep, the uvicorn timeout, and the lifespan's worst case: five hooks bounded at 5s each (startup trash sweep, notification delivery worker, channel service, browser sessions, MCP session pool), the 1s retrieval-warm wait, the in-flight run drain (5s) and the memory-queue drain (`memory.shutdown_flush_timeout_seconds`, default 30s), about 61s in total, plus a buffer — because the drain runs on a daemon thread and K8s SIGKILLs anything still running at the end of the grace window. K8s defaults to 30s, which SIGKILLs the drain mid-flight and silently re-introduces the memory loss the drain is fixing. `backend/tests/_gateway_shutdown_budget.py` reads these bounds from the Gateway and pins the chart and compose budgets against them. **When you raise `memory.shutdown_flush_timeout_seconds`, raise `gateway.terminationGracePeriodSeconds` to match** (preStop + uvicorn timeout + ~31s of bounded hooks and drains + memory drain + buffer). +- **Graceful shutdown & memory drain.** The gateway pod sets `terminationGracePeriodSeconds` (default 90s, overridable via `gateway.terminationGracePeriodSeconds`) plus an optional `preStop` sleep (`gateway.preStopSleepSeconds`, default 5s), and bounds uvicorn's `--timeout-graceful-shutdown` (`gateway.uvicornGracefulShutdownSeconds`, default 10s) so an idle SSE connection cannot hold up lifespan shutdown indefinitely. The grace period MUST exceed the Gateway's graceful-shutdown work — the preStop sleep, the uvicorn timeout, and the lifespan's worst case: six hooks bounded at 5s each (startup trash sweep, notification delivery worker, channel service, scheduled task service, browser sessions, MCP session pool), the 1s retrieval-warm wait, the in-flight run drain (5s) and the memory-queue drain (`memory.shutdown_flush_timeout_seconds`, default 30s), about 66s in total, plus a buffer — because the drain runs on a daemon thread and K8s SIGKILLs anything still running at the end of the grace window. K8s defaults to 30s, which SIGKILLs the drain mid-flight and silently re-introduces the memory loss the drain is fixing. `backend/tests/_gateway_shutdown_budget.py` reads these bounds from the Gateway and pins the chart and compose budgets against them. **When you raise `memory.shutdown_flush_timeout_seconds`, raise `gateway.terminationGracePeriodSeconds` to match** (preStop + uvicorn timeout + ~36s of bounded hooks and drains + memory drain + buffer). - **Gateway replicas.** Run control is cross-pod-safe since the work tracked by [issue #3948](https://github.com/bytedance/deer-flow/issues/3948) landed (#4003, #4064, #4500): admission is a durable one-active-run-per-thread