From f6da662c68be0d1701aa064a8e93df0cbe071528 Mon Sep 17 00:00:00 2001 From: jay <1569941651@qq.com> Date: Sun, 25 Oct 2026 23:02:38 +0800 Subject: [PATCH] fix(projects): make read_file_chunks cancellation-safe for worker handoffs (#7496) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(projects): make read_file_chunks cancellation-safe for worker handoffs Cancelling a project document read had three unsafe windows around the dedicated file-I/O worker: (1) a cancellation landing between the open worker completing and the result reaching the coroutine leaked the handle (the cleanup path never saw it); (2) the cleanup close could start while an admitted read was still running; (3) a repeated cancellation let the caller escape before the close worker settled (#7491). Gate every admitted operation with the existing await_drained primitive: cancellation is delivered only after the operation settles, the open result is retained in a list cell the moment the worker returns it (so the cleanup close always has the handle), and waiting for the next chunk via the yield stays immediately cancellable. Regression tests gate the real executor with asyncio events at the worker boundary (real files, no mocks): all three windows are red on unmodified main and green with the fix; the happy path is unchanged. The dispatch-spy promotion tests that import app.gateway.routers need the full backend dependency set and could not run in the minimal local env; they fail identically on unmodified main (verified). Fixes #7491 Co-authored-by: Jay Feng * test: fix _gated_offload return annotation to the unpacked tuple Review nit: the helper returns (calls, gated) and every call site unpacks it as a pair, but the annotation claimed list[str]. Type it as the tuple it actually returns. Co-authored-by: Jay Feng * test: pin missing-path FileNotFoundError in read_file_chunks Convergence item from the duplicate-pair verdict (#7493): the open-failure path must surface FileNotFoundError as-is — the cleanup finally skips the close because retained is empty (no handle was ever opened), rather than raising AttributeError from closing None. Co-authored-by: Jay Feng --------- Co-authored-by: Jay Feng --- .../harness/deerflow/projects/documents.py | 22 +- .../test_read_file_chunks_cancellation.py | 192 ++++++++++++++++++ 2 files changed, 211 insertions(+), 3 deletions(-) create mode 100644 backend/tests/blocking_io/test_read_file_chunks_cancellation.py diff --git a/backend/packages/harness/deerflow/projects/documents.py b/backend/packages/harness/deerflow/projects/documents.py index 553978e01e..246e4524be 100644 --- a/backend/packages/harness/deerflow/projects/documents.py +++ b/backend/packages/harness/deerflow/projects/documents.py @@ -593,17 +593,33 @@ async def read_file_chunks(path: Path, *, chunk_size: int = 65536) -> AsyncItera Shared by the from-thread promote path (thread file → shelf staging) and the attach path (staged shelf copy → thread-upload ingestion). + + Cancellation contract: once an open/read/close has been admitted to the + file-I/O worker, the caller's cancellation is delivered only after that + operation settles (:func:`await_drained`). The opened handle is retained + the moment the worker returns it, so a cancellation racing the result + handoff cannot leak it, and the cleanup close never races an in-flight + read. Waiting for the next chunk (the ``yield`` itself) stays immediately + cancellable. (#7491) """ def _open() -> Any: return open(path, "rb") - handle = await run_file_io(_open) + retained: list[Any] = [] + + async def _open_and_retain() -> Any: + handle = await run_file_io(_open) + retained.append(handle) + return handle + try: - while chunk := await run_file_io(handle.read, chunk_size): + handle = await await_drained(_open_and_retain()) + while chunk := await await_drained(run_file_io(handle.read, chunk_size)): yield chunk finally: - await run_file_io(handle.close) + for open_handle in retained: + await await_drained(run_file_io(open_handle.close)) def _copy_original_under_lock(paths: Paths, staging_path: Path, *, user_id: str, row: dict) -> None: diff --git a/backend/tests/blocking_io/test_read_file_chunks_cancellation.py b/backend/tests/blocking_io/test_read_file_chunks_cancellation.py new file mode 100644 index 0000000000..144ed6a28b --- /dev/null +++ b/backend/tests/blocking_io/test_read_file_chunks_cancellation.py @@ -0,0 +1,192 @@ +"""Cancellation-safety regression anchors for ``read_file_chunks`` (#7491). + +Each test gates the real ``run_file_io`` dispatch with an ``asyncio.Event`` at +the worker boundary (real executor, real temporary files) to reproduce one of +the three unsafe windows: a cancellation racing the open-result handoff, a +close racing an in-flight read, and a repeated cancellation trying to escape a +pending close. The expected contract: once an operation has been admitted to +the worker, the caller's cancellation is delivered only after it settles, the +handle always reaches the cleanup path, and the close never races a read. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Callable, Coroutine +from typing import Any + +import pytest + +from deerflow.projects.documents import read_file_chunks +from deerflow.utils.file_io import run_file_io as _real_run_file_io + +pytestmark = [ + pytest.mark.asyncio, + pytest.mark.allow_blocking_io, +] + + +class _Gate: + def __init__(self, mode: str) -> None: + self.mode = mode + self.admitted = asyncio.Event() + self.hold = asyncio.Event() + self.released = False + + def release(self) -> None: + self.released = True + self.hold.set() + + +def _gated_offload( + gates: dict[str, _Gate], + on_call, + on_done, +) -> tuple[list[str], Callable[..., Coroutine[Any, Any, Any]]]: + """Wrap ``documents.run_file_io`` call-through with worker-boundary gates. + + ``on_call`` fires when a dispatch enters the wrapper (admission), + ``on_done`` fires after the real worker call settles. + """ + calls: list[str] = [] + + async def gated(func, /, *args, **kwargs): + name = getattr(func, "__name__", "") + calls.append(name) + on_call(name) + gate = gates.get(name) + if gate is None: + result = await _real_run_file_io(func, *args, **kwargs) + elif gate.mode == "before": + gate.admitted.set() + await gate.hold.wait() + result = await _real_run_file_io(func, *args, **kwargs) + else: + result = await _real_run_file_io(func, *args, **kwargs) + gate.admitted.set() + await gate.hold.wait() + on_done(name) + return result + + return calls, gated + + +async def _settle() -> None: + for _ in range(10): + await asyncio.sleep(0) + + +async def test_cancel_racing_open_handoff_still_closes_handle(tmp_path, monkeypatch) -> None: + """Window 1: the open worker returns, then cancellation lands before the + result reaches the coroutine — the handle must still be closed.""" + path = tmp_path / "doc.bin" + path.write_bytes(b"payload") + open_gate = _Gate("after") + state: dict[str, bool] = {"closed": False} + + def on_call(name: str) -> None: ... + def on_done(name: str) -> None: + if name == "close": + state["closed"] = True + + calls, gated = _gated_offload({"_open": open_gate}, on_call, on_done) + monkeypatch.setattr("deerflow.projects.documents.run_file_io", gated) + + agen = read_file_chunks(path, chunk_size=4) + task = asyncio.create_task(agen.__anext__()) + await asyncio.wait_for(open_gate.admitted.wait(), 5) + task.cancel() + await _settle() + assert "close" not in calls, "cleanup must wait for the open handoff to settle" + + open_gate.release() + with pytest.raises(asyncio.CancelledError): + await task + assert state["closed"], "handle retained at handoff must reach the cleanup close" + + +async def test_cancel_during_read_does_not_race_close(tmp_path, monkeypatch) -> None: + """Window 2: cancellation while a read is admitted must not start the + cleanup close before that read settles.""" + path = tmp_path / "doc.bin" + path.write_bytes(b"abcdefgh") + read_gate = _Gate("before") + state: dict[str, bool] = {"close_raced": False, "read_settled": False} + + def on_call(name: str) -> None: + if name == "close" and not state["read_settled"]: + state["close_raced"] = True + + def on_done(name: str) -> None: + if name == "read" and read_gate.released: + state["read_settled"] = True + + calls, gated = _gated_offload({"read": read_gate}, on_call, on_done) + monkeypatch.setattr("deerflow.projects.documents.run_file_io", gated) + + agen = read_file_chunks(path, chunk_size=4) + task = asyncio.create_task(agen.__anext__()) + await asyncio.wait_for(read_gate.admitted.wait(), 5) + task.cancel() + await _settle() + assert "close" not in calls, "close must wait for the admitted read to settle" + + read_gate.release() + with pytest.raises(asyncio.CancelledError): + await task + assert state["read_settled"], "the admitted read must run to completion" + assert not state["close_raced"], "close must not precede or overlap the in-flight read" + + +async def test_repeated_cancellation_cannot_escape_pending_close(tmp_path, monkeypatch) -> None: + """Window 3: a cancellation arriving while the cleanup close is admitted + must be delivered only after the close settles — repeated cancellation + included.""" + path = tmp_path / "doc.bin" + path.write_bytes(b"abc") + close_gate = _Gate("before") + state: dict[str, bool] = {"close_completed": False} + + def on_call(name: str) -> None: ... + + def on_done(name: str) -> None: + if name == "close" and close_gate.released: + state["close_completed"] = True + + calls, gated = _gated_offload({"close": close_gate}, on_call, on_done) + monkeypatch.setattr("deerflow.projects.documents.run_file_io", gated) + + async def consume() -> list[bytes]: + return [chunk async for chunk in read_file_chunks(path, chunk_size=8)] + + task = asyncio.create_task(consume()) + await asyncio.wait_for(close_gate.admitted.wait(), 5) + task.cancel() + await _settle() + task.cancel() + await _settle() + close_gate.release() + with pytest.raises(asyncio.CancelledError): + await task + assert state["close_completed"], "the caller may escape only after the close settles" + + +async def test_read_file_chunks_happy_path_unchanged(tmp_path, monkeypatch) -> None: + """Sanity: full consumption still yields the exact bytes.""" + path = tmp_path / "doc.bin" + path.write_bytes(b"abcdefgh") + calls, gated = _gated_offload({}, lambda name: None, lambda name: None) + monkeypatch.setattr("deerflow.projects.documents.run_file_io", gated) + + chunks = [chunk async for chunk in read_file_chunks(path, chunk_size=4)] + + assert chunks == [b"abcd", b"efgh"] + assert calls[0] == "_open" and calls[-1] == "close" + + +async def test_missing_path_raises_file_not_found(tmp_path) -> None: + """Pins the open-failure path: the error surfaces as-is and the cleanup + skips the close (``retained`` is empty — no handle was ever opened).""" + with pytest.raises(FileNotFoundError): + async for _ in read_file_chunks(tmp_path / "absent.bin", chunk_size=4): + pass