fix(projects): drain trash reconciliation before cancellation returns (#5511)

This commit is contained in:
NanPan 2026-09-18 09:05:26 +08:00 committed by GitHub
parent 9f79ddf9b6
commit 408b015d5f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 120 additions and 25 deletions

View File

@ -65,9 +65,9 @@ logging.basicConfig(
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
# Upper bound (seconds) each lifespan shutdown hook is allowed to run. # Grace budget (seconds) lifespan shutdown hooks get for ordinary completion.
# Bounds worker exit time so uvicorn's reload supervisor does not keep # Hooks that own uncancellable executor work may spend additional time draining
# firing signals into a worker that is stuck waiting for shutdown cleanup. # already-started work after cancellation rather than detach it from teardown.
_SHUTDOWN_HOOK_TIMEOUT_SECONDS = 5.0 _SHUTDOWN_HOOK_TIMEOUT_SECONDS = 5.0
# The retrieval index is derived state, so shutdown only waits briefly for its # 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 Runs beside the lazy trigger on the trash listing no daemon, no
scheduler (§15.9). Sweeps every user (``user_id=None``) with the scheduler (§15.9). Sweeps every user (``user_id=None``) with the
configured retention window, including the full reconciliation. A sweep configured retention window, including the full reconciliation. A sweep
failure is logged and never blocks gateway readiness; the lifespan runs failure is logged and never blocks gateway readiness; on shutdown the
this as a background task and awaits it (bounded) on shutdown, cancelling lifespan gives it a bounded graceful-completion budget, then cancels an
it when the budget runs out. overrun while draining any already-started file reconciliation before
teardown continues.
""" """
try: try:
from deerflow.config.paths import get_paths 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: 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 Waits ``_SHUTDOWN_HOOK_TIMEOUT_SECONDS`` for graceful completion, then
cancels it when the budget runs out. The shield keeps that wait bounded cancels an overrun. That budget is not a hard upper bound for this hook:
without killing the sweep, so an overrun must be cancelled here: the reconciliation runs in executor threads that cannot be safely killed, so
all-users reconciliation reads through the document repo and the DB cancellation drains any already-started filesystem worker before returning.
engine, and leaving it running would have it walk rows and files while This keeps the sweep's file ownership intact until the teardown below can
the teardown below disposes both underneath it. safely dispose the document repo and DB engine.
""" """
task = getattr(app.state, "startup_trash_sweep_task", None) task = getattr(app.state, "startup_trash_sweep_task", None)
if task is None or task.done(): if task is None or task.done():
@ -246,17 +247,19 @@ async def _shutdown_startup_trash_sweep(app: FastAPI) -> None:
try: try:
await asyncio.wait_for(asyncio.shield(task), timeout=_SHUTDOWN_HOOK_TIMEOUT_SECONDS) await asyncio.wait_for(asyncio.shield(task), timeout=_SHUTDOWN_HOOK_TIMEOUT_SECONDS)
except TimeoutError: except TimeoutError:
# Cancellation lands at the sweep's next await; ``_run_startup_trash_sweep`` # Cancellation prevents later sweep stages from starting, but an
# only catches ``Exception``, so ``CancelledError`` propagates. A # executor-backed reconciliation that already started drains before
# ``cancel()`` that returns False means the sweep finished inside the # ``CancelledError`` reaches this task. The final await may therefore
# window between the deadline firing and this call — report that as # exceed the graceful-completion budget; detaching that worker would
# the late finish it is, not as a cancellation that never happened. # 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() cancelled = task.cancel()
with suppress(asyncio.CancelledError): with suppress(asyncio.CancelledError):
await task await task
if cancelled: if cancelled:
logger.warning( 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, _SHUTDOWN_HOOK_TIMEOUT_SECONDS,
) )
else: else:
@ -387,9 +390,9 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
# Runs after langgraph_runtime so app.state.project_document_repo is # Runs after langgraph_runtime so app.state.project_document_repo is
# available. The per-user reconciliation walks every row and file, so # available. The per-user reconciliation walks every row and file, so
# it is scheduled as a background task: gateway readiness never waits # it is scheduled as a background task: gateway readiness never waits
# on it, a failure is logged by the task itself, and shutdown awaits # on it, a failure is logged by the task itself, and shutdown gives the
# the in-flight sweep (bounded, cancelled on overrun) before the # in-flight sweep a bounded graceful budget before cancelling it; any
# runtime is torn down. # 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)) app.state.startup_trash_sweep_task = asyncio.create_task(_run_startup_trash_sweep(app, startup_config))
try: try:

View File

@ -30,7 +30,7 @@ from typing import TYPE_CHECKING
from deerflow.config.paths import Paths from deerflow.config.paths import Paths
from deerflow.projects.documents import _content_intact, check_document_content, converted_markdown_path, original_file_path 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 from deerflow.utils.time import coerce_iso
if TYPE_CHECKING: if TYPE_CHECKING:
@ -330,6 +330,6 @@ async def run_trash_retention_sweep(
if include_reconciliation: if include_reconciliation:
rows = await repo.list_all_for_sweep(user_id=user_id) rows = await repo.list_all_for_sweep(user_id=user_id)
guard_cutoff = now - _ORPHAN_GUARD 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 await_drained(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_rows, paths, rows=rows, guard_cutoff=guard_cutoff, report=report))
return report return report

View File

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