From 6a94bef9086ff78cc599a0062ad52fcf3e8a2ae2 Mon Sep 17 00:00:00 2001 From: NanPan <111261006+poijygfdyy@users.noreply.github.com> Date: Thu, 17 Sep 2026 21:26:52 +0800 Subject: [PATCH] fix(gateway): keep run drain alive across repeated cancellation (#5487) * test(gateway): cover repeated cancellation during run drain * fix(gateway): keep run drain alive across repeated cancellation * test(gateway): harden run-drain cancellation coverage --- backend/app/gateway/deps.py | 41 +++++--- ...gateway_run_drain_repeated_cancellation.py | 96 +++++++++++++++++++ 2 files changed, 123 insertions(+), 14 deletions(-) create mode 100644 backend/tests/test_gateway_run_drain_repeated_cancellation.py diff --git a/backend/app/gateway/deps.py b/backend/app/gateway/deps.py index f4d834ea0..1e93f23f9 100644 --- a/backend/app/gateway/deps.py +++ b/backend/app/gateway/deps.py @@ -167,24 +167,37 @@ async def _drain_inflight_runs(run_manager: RunManager) -> None: """Drain in-flight runs before the checkpointer is torn down (issue #3373). Shields the (internally-bounded) drain so that even if the lifespan - coroutine is itself cancelled mid-shutdown — a second SIGINT or the server's - graceful-shutdown timeout, i.e. the same signal storm behind #3373 — the - checkpointer pool is not closed while run tasks are still writing - checkpoints. On such a cancellation we let the already-running drain finish - (it is bounded by ``RunManager.shutdown``'s own timeout) and then propagate - the cancellation. + coroutine is repeatedly cancelled mid-shutdown — e.g. signal escalation or + the server's graceful-shutdown timeout — the checkpointer pool is not closed + while run tasks are still writing checkpoints. Cancellation is remembered + and propagated only after the already-running drain reaches a safe terminal + point. """ drain = asyncio.create_task(run_manager.shutdown(timeout=_RUN_DRAIN_TIMEOUT_SECONDS)) - try: - await asyncio.shield(drain) - except asyncio.CancelledError: - # Re-shield so this second wait does not abandon the in-flight drain; - # it is bounded, so this cannot hang. Then re-raise to honour shutdown. + cancellation: asyncio.CancelledError | None = None + + while not drain.done(): try: await asyncio.shield(drain) + except asyncio.CancelledError as exc: + if cancellation is None: + cancellation = exc + except Exception: + if cancellation is not None: + logger.exception("In-flight run drain failed after shutdown cancellation") + raise cancellation + logger.exception("Failed to drain in-flight runs during shutdown") + return + + if cancellation is not None: + try: + drain.result() except Exception: logger.exception("In-flight run drain failed after shutdown cancellation") - raise + raise cancellation + + try: + drain.result() except Exception: logger.exception("Failed to drain in-flight runs during shutdown") @@ -736,8 +749,8 @@ def get_run_context(request: Request) -> RunContext: ``app_config`` field is resolved live so per-run fields (e.g. ``models[*].max_tokens``) follow ``config.yaml`` edits; the ``event_store`` / ``run_events_config`` pair stays frozen to the snapshot - captured in :func:`langgraph_runtime` so callers never see a store bound - to one backend paired with a config pointing at another. + captured in :func:`langgraph_runtime` so callers never see a store bound to + one backend paired with a config pointing at another. """ return RunContext( checkpointer=get_checkpointer(request), diff --git a/backend/tests/test_gateway_run_drain_repeated_cancellation.py b/backend/tests/test_gateway_run_drain_repeated_cancellation.py new file mode 100644 index 000000000..17b6f3828 --- /dev/null +++ b/backend/tests/test_gateway_run_drain_repeated_cancellation.py @@ -0,0 +1,96 @@ +"""Repeated-cancellation regressions for Gateway run draining on shutdown.""" + +from __future__ import annotations + +import asyncio + +import pytest + + +@pytest.mark.asyncio +async def test_run_drain_waits_for_shutdown_across_repeated_cancellation(): + """A second cancellation cannot let checkpointer teardown outrun the drain.""" + from app.gateway import deps + + shutdown_started = asyncio.Event() + allow_shutdown_finish = asyncio.Event() + shutdown_finished = asyncio.Event() + + class _RunManager: + async def shutdown(self, *, timeout: float) -> None: + assert timeout == deps._RUN_DRAIN_TIMEOUT_SECONDS + shutdown_started.set() + await allow_shutdown_finish.wait() + shutdown_finished.set() + + drain_task = asyncio.create_task(deps._drain_inflight_runs(_RunManager())) + await asyncio.wait_for(shutdown_started.wait(), timeout=1.0) + + drain_task.cancel("first shutdown cancellation") + # Let the helper observe and remember the first cancellation, then suspend + # again while the still-blocked RunManager.shutdown() task remains alive. + await asyncio.sleep(0) + assert not drain_task.done() + + # A second SIGINT / graceful-shutdown cancellation arrives while the helper + # is already draining after the first one. It must keep owning that wait so + # the surrounding AsyncExitStack cannot close the checkpointer underneath + # still-running run tasks. + drain_task.cancel("second shutdown cancellation") + await asyncio.sleep(0) + escaped_before_shutdown_finished = drain_task.done() + + allow_shutdown_finish.set() + with pytest.raises(asyncio.CancelledError) as exc_info: + await drain_task + + assert not escaped_before_shutdown_finished + assert shutdown_finished.is_set() + assert exc_info.value.args == ("first shutdown cancellation",) + + +@pytest.mark.asyncio +async def test_run_drain_logs_failure_and_returns_without_cancellation(caplog): + """A drain failure alone keeps the helper's existing best-effort contract.""" + from app.gateway import deps + + class _RunManager: + async def shutdown(self, *, timeout: float) -> None: + assert timeout == deps._RUN_DRAIN_TIMEOUT_SECONDS + raise RuntimeError("drain failed") + + await deps._drain_inflight_runs(_RunManager()) + + assert "Failed to drain in-flight runs during shutdown" in caplog.text + assert "drain failed" in caplog.text + + +@pytest.mark.asyncio +async def test_run_drain_failure_preserves_remembered_cancellation(caplog): + """A later drain error cannot replace cancellation already requested.""" + from app.gateway import deps + + shutdown_started = asyncio.Event() + allow_shutdown_failure = asyncio.Event() + + class _RunManager: + async def shutdown(self, *, timeout: float) -> None: + assert timeout == deps._RUN_DRAIN_TIMEOUT_SECONDS + shutdown_started.set() + await allow_shutdown_failure.wait() + raise RuntimeError("drain failed after cancellation") + + drain_task = asyncio.create_task(deps._drain_inflight_runs(_RunManager())) + await asyncio.wait_for(shutdown_started.wait(), timeout=1.0) + + drain_task.cancel("shutdown requested") + await asyncio.sleep(0) + assert not drain_task.done() + + allow_shutdown_failure.set() + with pytest.raises(asyncio.CancelledError) as exc_info: + await drain_task + + assert exc_info.value.args == ("shutdown requested",) + assert "In-flight run drain failed after shutdown cancellation" in caplog.text + assert "drain failed after cancellation" in caplog.text