From a78ff4817f16947c666fb0e29381ba3fb3f6352d Mon Sep 17 00:00:00 2001 From: NanPan <111261006+poijygfdyy@users.noreply.github.com> Date: Tue, 6 Oct 2026 20:35:54 +0800 Subject: [PATCH] fix(extensions): drain pending notify dispatches on shutdown (#6357) * fix(extensions): drain pending notify dispatches on shutdown * fix(gateway): await extension notify drain before reset * test(extensions): retain pending notify dispatches on shutdown * fix(extensions): bound and close notify shutdown drain * test(extensions): cover bounded notify shutdown drain --- backend/app/gateway/deps.py | 17 +++-- .../harness/deerflow/extensions/notify.py | 33 +++++++++- ..._extension_notify_shutdown_cancellation.py | 64 +++++++++++++++++++ 3 files changed, 109 insertions(+), 5 deletions(-) create mode 100644 backend/tests/test_extension_notify_shutdown_cancellation.py diff --git a/backend/app/gateway/deps.py b/backend/app/gateway/deps.py index 77ac3ea486..0b8748bb65 100644 --- a/backend/app/gateway/deps.py +++ b/backend/app/gateway/deps.py @@ -535,6 +535,7 @@ async def langgraph_runtime(app: FastAPI, startup_config: AppConfig) -> AsyncGen # startup-failure and cancellation path below. try: from deerflow.extensions.notify import ( + drain_extension_notify_dispatches, reset_extension_notify_loop, set_extension_notify_loop, ) @@ -544,16 +545,24 @@ async def langgraph_runtime(app: FastAPI, startup_config: AppConfig) -> AsyncGen logger.exception("Failed to register the extension notify loop; sync observations will be dropped") else: - def reset_notify_loop_safely() -> None: + async def reset_notify_loop_safely() -> None: try: - reset_extension_notify_loop() + await drain_extension_notify_dispatches() except Exception: logger.debug( - "Failed to reset the extension notify loop (non-fatal)", + "Failed to drain pending extension notifications (non-fatal)", exc_info=True, ) + finally: + try: + reset_extension_notify_loop() + except Exception: + logger.debug( + "Failed to reset the extension notify loop (non-fatal)", + exc_info=True, + ) - stack.callback(reset_notify_loop_safely) + stack.push_async_callback(reset_notify_loop_safely) config = startup_config app.state.checkpoint_channel_mode = freeze_checkpoint_channel_mode(config.database.checkpoint_channel_mode) diff --git a/backend/packages/harness/deerflow/extensions/notify.py b/backend/packages/harness/deerflow/extensions/notify.py index f51cd7184a..b0d77d60c9 100644 --- a/backend/packages/harness/deerflow/extensions/notify.py +++ b/backend/packages/harness/deerflow/extensions/notify.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import concurrent.futures import logging import time from collections.abc import Awaitable, Callable, Coroutine, Mapping @@ -20,6 +21,7 @@ from deerflow_extension_api import ( ) from deerflow.extensions.registry import LoadedExtensions +from deerflow.utils.file_io import await_drained logger = logging.getLogger(__name__) @@ -182,7 +184,7 @@ async def _notify_each( # loops, but extension resources must always be touched on the loop where they # were started. _notify_loop: asyncio.AbstractEventLoop | None = None -_pending_dispatches: set[asyncio.Future[Any]] = set() +_pending_dispatches: set[concurrent.futures.Future[None]] = set() _warned_no_loop = False _system_observations_enabled = True @@ -195,6 +197,35 @@ def set_extension_notify_loop(loop: asyncio.AbstractEventLoop | None) -> None: _warned_no_loop = False +async def drain_extension_notify_dispatches(*, timeout: float = 5.0) -> None: + """Boundedly drain already-submitted fire-and-forget observations. + + New detached observations are suspended before the snapshot so shutdown + cannot open a submit-after-snapshot window. The bounded wait keeps a wedged + extension observer from holding Gateway teardown forever. + """ + suspend_extension_system_observations() + pending = tuple(_pending_dispatches) + if not pending: + return + + async def _wait() -> None: + try: + async with asyncio.timeout(timeout): + await asyncio.gather( + *(asyncio.wrap_future(future) for future in pending), + return_exceptions=True, + ) + except TimeoutError: + logger.warning( + "Timed out after %.1fs draining %d pending extension observation(s)", + timeout, + len(pending), + ) + + await await_drained(_wait()) + + def reset_extension_notify_loop() -> None: """Remove the process-wide loop binding during host shutdown or tests.""" global _notify_loop, _system_observations_enabled, _warned_no_loop diff --git a/backend/tests/test_extension_notify_shutdown_cancellation.py b/backend/tests/test_extension_notify_shutdown_cancellation.py new file mode 100644 index 0000000000..31a7ff4a71 --- /dev/null +++ b/backend/tests/test_extension_notify_shutdown_cancellation.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +import asyncio +from concurrent.futures import Future + +import pytest + +from deerflow.extensions.notify import ( + _pending_dispatches, + drain_extension_notify_dispatches, + reset_extension_notify_loop, +) + + +@pytest.mark.asyncio +async def test_extension_notify_shutdown_drains_pending_dispatch_across_repeated_cancellation() -> None: + reset_extension_notify_loop() + pending: Future[None] = Future() + _pending_dispatches.add(pending) + pending.add_done_callback(_pending_dispatches.discard) + + task = asyncio.create_task(drain_extension_notify_dispatches()) + try: + await asyncio.sleep(0) + task.cancel() + for _ in range(5): + await asyncio.sleep(0) + assert not task.done() + + task.cancel() + for _ in range(5): + await asyncio.sleep(0) + assert not task.done() + + pending.set_result(None) + with pytest.raises(asyncio.CancelledError): + await task + finally: + if not pending.done(): + pending.set_result(None) + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + reset_extension_notify_loop() + + assert pending not in _pending_dispatches + + +@pytest.mark.asyncio +async def test_extension_notify_shutdown_drain_is_bounded() -> None: + reset_extension_notify_loop() + pending: Future[None] = Future() + _pending_dispatches.add(pending) + pending.add_done_callback(_pending_dispatches.discard) + + try: + await asyncio.wait_for( + drain_extension_notify_dispatches(timeout=0.01), + timeout=0.5, + ) + finally: + if not pending.done(): + pending.cancel() + reset_extension_notify_loop()