fix(projects): make read_file_chunks cancellation-safe for worker handoffs (#7496)

* 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 <jayfeng77@outlook.com>

* 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 <jayfeng77@outlook.com>

* 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 <jayfeng77@outlook.com>

---------

Co-authored-by: Jay Feng <jayfeng77@outlook.com>
This commit is contained in:
jay 2026-10-25 23:02:38 +08:00 committed by GitHub
parent 9a92d2e9b7
commit f6da662c68
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 211 additions and 3 deletions

View File

@ -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:

View File

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