From 408b015d5fbc8dcf893ad5efbbd80d4bfc150cc3 Mon Sep 17 00:00:00 2001 From: NanPan <111261006+poijygfdyy@users.noreply.github.com> Date: Fri, 18 Sep 2026 09:05:26 +0800 Subject: [PATCH] fix(projects): drain trash reconciliation before cancellation returns (#5511) --- backend/app/gateway/app.py | 47 +++++----- .../harness/deerflow/projects/trash.py | 6 +- ...oject_trash_reconciliation_cancellation.py | 92 +++++++++++++++++++ 3 files changed, 120 insertions(+), 25 deletions(-) create mode 100644 backend/tests/test_project_trash_reconciliation_cancellation.py diff --git a/backend/app/gateway/app.py b/backend/app/gateway/app.py index dbe5b0c90..4f8a6ebcc 100644 --- a/backend/app/gateway/app.py +++ b/backend/app/gateway/app.py @@ -65,9 +65,9 @@ logging.basicConfig( logger = logging.getLogger(__name__) -# Upper bound (seconds) each lifespan shutdown hook is allowed to run. -# Bounds worker exit time so uvicorn's reload supervisor does not keep -# firing signals into a worker that is stuck waiting for shutdown cleanup. +# Grace budget (seconds) lifespan shutdown hooks get for ordinary completion. +# Hooks that own uncancellable executor work may spend additional time draining +# already-started work after cancellation rather than detach it from teardown. _SHUTDOWN_HOOK_TIMEOUT_SECONDS = 5.0 # The retrieval index is derived state, so shutdown only waits briefly for its @@ -202,9 +202,10 @@ async def _run_startup_trash_sweep(app: FastAPI, startup_config) -> None: Runs beside the lazy trigger on the trash listing — no daemon, no scheduler (§15.9). Sweeps every user (``user_id=None``) with the configured retention window, including the full reconciliation. A sweep - failure is logged and never blocks gateway readiness; the lifespan runs - this as a background task and awaits it (bounded) on shutdown, cancelling - it when the budget runs out. + failure is logged and never blocks gateway readiness; on shutdown the + lifespan gives it a bounded graceful-completion budget, then cancels an + overrun while draining any already-started file reconciliation before + teardown continues. """ try: from deerflow.config.paths import get_paths @@ -231,14 +232,14 @@ async def _run_startup_trash_sweep(app: FastAPI, startup_config) -> None: async def _shutdown_startup_trash_sweep(app: FastAPI) -> None: - """Bounded shutdown wait for the background startup sweep (§8.3). + """Grace-budgeted shutdown for the background startup sweep (§8.3). - Waits ``_SHUTDOWN_HOOK_TIMEOUT_SECONDS`` for an in-flight sweep and - cancels it when the budget runs out. The shield keeps that wait bounded - without killing the sweep, so an overrun must be cancelled here: the - all-users reconciliation reads through the document repo and the DB - engine, and leaving it running would have it walk rows and files while - the teardown below disposes both underneath it. + Waits ``_SHUTDOWN_HOOK_TIMEOUT_SECONDS`` for graceful completion, then + cancels an overrun. That budget is not a hard upper bound for this hook: + reconciliation runs in executor threads that cannot be safely killed, so + cancellation drains any already-started filesystem worker before returning. + This keeps the sweep's file ownership intact until the teardown below can + safely dispose the document repo and DB engine. """ task = getattr(app.state, "startup_trash_sweep_task", None) if task is None or task.done(): @@ -246,17 +247,19 @@ async def _shutdown_startup_trash_sweep(app: FastAPI) -> None: try: await asyncio.wait_for(asyncio.shield(task), timeout=_SHUTDOWN_HOOK_TIMEOUT_SECONDS) except TimeoutError: - # Cancellation lands at the sweep's next await; ``_run_startup_trash_sweep`` - # only catches ``Exception``, so ``CancelledError`` propagates. A - # ``cancel()`` that returns False means the sweep finished inside the - # window between the deadline firing and this call — report that as - # the late finish it is, not as a cancellation that never happened. + # Cancellation prevents later sweep stages from starting, but an + # executor-backed reconciliation that already started drains before + # ``CancelledError`` reaches this task. The final await may therefore + # exceed the graceful-completion budget; detaching that worker would + # reintroduce teardown/file-mutation overlap. A ``cancel()`` that + # returns False means the sweep finished inside the window between the + # deadline firing and this call — report that as the late finish it is. cancelled = task.cancel() with suppress(asyncio.CancelledError): await task if cancelled: logger.warning( - "Startup trash sweep exceeded %.1fs during shutdown; cancelled and proceeding with worker exit.", + "Startup trash sweep exceeded %.1fs during shutdown; cancelled and drained before proceeding with worker exit.", _SHUTDOWN_HOOK_TIMEOUT_SECONDS, ) else: @@ -387,9 +390,9 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: # Runs after langgraph_runtime so app.state.project_document_repo is # available. The per-user reconciliation walks every row and file, so # it is scheduled as a background task: gateway readiness never waits - # on it, a failure is logged by the task itself, and shutdown awaits - # the in-flight sweep (bounded, cancelled on overrun) before the - # runtime is torn down. + # on it, a failure is logged by the task itself, and shutdown gives the + # in-flight sweep a bounded graceful budget before cancelling it; any + # already-started file worker drains before the runtime is torn down. app.state.startup_trash_sweep_task = asyncio.create_task(_run_startup_trash_sweep(app, startup_config)) try: diff --git a/backend/packages/harness/deerflow/projects/trash.py b/backend/packages/harness/deerflow/projects/trash.py index 6ead65b56..fdafe15f6 100644 --- a/backend/packages/harness/deerflow/projects/trash.py +++ b/backend/packages/harness/deerflow/projects/trash.py @@ -30,7 +30,7 @@ from typing import TYPE_CHECKING from deerflow.config.paths import Paths from deerflow.projects.documents import _content_intact, check_document_content, converted_markdown_path, original_file_path -from deerflow.utils.file_io import run_file_io +from deerflow.utils.file_io import await_drained, run_file_io from deerflow.utils.time import coerce_iso if TYPE_CHECKING: @@ -330,6 +330,6 @@ async def run_trash_retention_sweep( if include_reconciliation: rows = await repo.list_all_for_sweep(user_id=user_id) guard_cutoff = now - _ORPHAN_GUARD - await run_file_io(_reconcile_storage, paths, user_id=user_id, rows=rows, guard_cutoff=guard_cutoff, report=report) - await run_file_io(_reconcile_rows, paths, rows=rows, guard_cutoff=guard_cutoff, report=report) + await await_drained(run_file_io(_reconcile_storage, paths, user_id=user_id, rows=rows, guard_cutoff=guard_cutoff, report=report)) + await await_drained(run_file_io(_reconcile_rows, paths, rows=rows, guard_cutoff=guard_cutoff, report=report)) return report diff --git a/backend/tests/test_project_trash_reconciliation_cancellation.py b/backend/tests/test_project_trash_reconciliation_cancellation.py new file mode 100644 index 000000000..fa4ea70e4 --- /dev/null +++ b/backend/tests/test_project_trash_reconciliation_cancellation.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +import asyncio +import threading +from types import SimpleNamespace + +import pytest +from fastapi import FastAPI + +import deerflow.projects.trash as trash_mod + + +class _EmptySweepRepo: + async def purge_candidates(self, *args, **kwargs): + return [] + + async def list_all_for_sweep(self, *args, **kwargs): + return [] + + +@pytest.mark.anyio +@pytest.mark.parametrize("worker_name", ["_reconcile_storage", "_reconcile_rows"]) +async def test_cancelled_retention_sweep_drains_reconciliation_worker(monkeypatch, worker_name: str) -> None: + """Cancellation cannot outlive an already-started reconciliation worker.""" + started = threading.Event() + finish = threading.Event() + + def blocking_reconcile(*args, **kwargs) -> None: + started.set() + finish.wait(5) + + if worker_name == "_reconcile_rows": + monkeypatch.setattr(trash_mod, "_reconcile_storage", lambda *args, **kwargs: None) + monkeypatch.setattr(trash_mod, worker_name, blocking_reconcile) + + task = asyncio.create_task( + trash_mod.run_trash_retention_sweep( + _EmptySweepRepo(), + object(), + retention_days=30, + user_id=None, + ) + ) + assert await asyncio.to_thread(started.wait, 5) + + task.cancel() + try: + for _ in range(10): + await asyncio.sleep(0) + assert not task.done(), "the sweep released ownership while its file-io worker was still running" + finally: + finish.set() + + with pytest.raises(asyncio.CancelledError): + await task + + +@pytest.mark.anyio +async def test_shutdown_hook_drains_running_reconciliation_worker_after_grace_budget(monkeypatch) -> None: + """The shutdown grace budget does not detach already-started file work.""" + import app.gateway.app as gateway_app + + started = threading.Event() + finish = threading.Event() + + def blocking_reconcile(*args, **kwargs) -> None: + started.set() + finish.wait(5) + + monkeypatch.setattr(trash_mod, "_reconcile_storage", blocking_reconcile) + monkeypatch.setattr(gateway_app, "_SHUTDOWN_HOOK_TIMEOUT_SECONDS", 0.01) + + app = FastAPI() + app.state.project_document_repo = _EmptySweepRepo() + sweep_task = asyncio.create_task(gateway_app._run_startup_trash_sweep(app, SimpleNamespace(projects=None))) + app.state.startup_trash_sweep_task = sweep_task + assert await asyncio.to_thread(started.wait, 5) + + shutdown_task = asyncio.create_task(gateway_app._shutdown_startup_trash_sweep(app)) + try: + for _ in range(100): + if sweep_task.cancelling(): + break + await asyncio.sleep(0.001) + assert sweep_task.cancelling(), "shutdown never cancelled the over-budget startup sweep" + assert not sweep_task.done(), "sweep cancellation detached its running file-io worker" + assert not shutdown_task.done(), "shutdown hook returned while reconciliation was still running" + finally: + finish.set() + + await shutdown_task + assert sweep_task.cancelled()