mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-20 19:46:16 +00:00
fix(projects): drain trash reconciliation before cancellation returns (#5511)
This commit is contained in:
parent
9f79ddf9b6
commit
408b015d5f
@ -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:
|
||||||
|
|||||||
@ -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
|
||||||
|
|||||||
@ -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()
|
||||||
Loading…
x
Reference in New Issue
Block a user