mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
fix(uploads): retain filesystem ownership across request cancellation (#7170)
* fix(uploads): retain filesystem ownership across cancellation * fix(uploads): preserve cancellation across cleanup failures
This commit is contained in:
parent
19e861d0c3
commit
f4eaa51a73
@ -5,6 +5,7 @@ import os
|
||||
import stat
|
||||
import tempfile
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import suppress
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import BinaryIO
|
||||
@ -190,7 +191,8 @@ def _cleanup_uploaded_paths(paths: list[os.PathLike[str] | str]) -> None:
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
except Exception:
|
||||
logger.warning("Failed to clean up upload path after rejected request: %s", path, exc_info=True)
|
||||
with suppress(Exception):
|
||||
logger.warning("Failed to clean up upload path after rejected request: %s", path, exc_info=True)
|
||||
|
||||
|
||||
def _pure_destination(uploads_dir: os.PathLike[str] | str, display_filename: str) -> Path:
|
||||
@ -320,7 +322,8 @@ def _remove_staged_file(staged_path: os.PathLike[str] | str) -> None:
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
except OSError:
|
||||
logger.warning("Failed to remove staged upload file: %s", staged_path, exc_info=True)
|
||||
with suppress(Exception):
|
||||
logger.warning("Failed to remove staged upload file: %s", staged_path, exc_info=True)
|
||||
|
||||
|
||||
def _write_upload_chunk(upload_temp: _UploadTempFile, chunk: bytes) -> None:
|
||||
@ -330,6 +333,9 @@ def _write_upload_chunk(upload_temp: _UploadTempFile, chunk: bytes) -> None:
|
||||
def _abort_upload_temp(upload_temp: _UploadTempFile) -> None:
|
||||
try:
|
||||
upload_temp.handle.close()
|
||||
except Exception:
|
||||
with suppress(Exception):
|
||||
logger.warning("Failed to close staged upload file: %s", upload_temp.temp_path, exc_info=True)
|
||||
finally:
|
||||
# Best-effort, not ``os.unlink``: an abandoned duplication worker can
|
||||
# still hold this inode open (Windows refuses to remove it then), and
|
||||
|
||||
@ -35,6 +35,7 @@ import shutil
|
||||
import tempfile
|
||||
import uuid
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import suppress
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
@ -42,6 +43,7 @@ from fastapi import HTTPException
|
||||
|
||||
from deerflow.config.app_config import AppConfig
|
||||
from deerflow.uploads.companions import register_companion, unregister_companion
|
||||
from deerflow.utils.atomic_write import cleanup_failed_atomic_write
|
||||
from deerflow.utils.file_io import await_drained, run_file_io
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@ -60,7 +62,8 @@ def _close_fd(fd: int | None) -> None:
|
||||
try:
|
||||
os.close(fd)
|
||||
except OSError:
|
||||
logger.warning("Failed to close upload conversion descriptor", exc_info=True)
|
||||
with suppress(Exception):
|
||||
logger.warning("Failed to close upload conversion descriptor", exc_info=True)
|
||||
|
||||
|
||||
def _close_abandoned_fd(duplication: asyncio.Future) -> None:
|
||||
@ -160,7 +163,11 @@ class ThreadUploadIngestionService:
|
||||
self._total_size = 0
|
||||
|
||||
async def __aenter__(self) -> ThreadUploadIngestionService:
|
||||
await self.open()
|
||||
try:
|
||||
await self.open()
|
||||
except BaseException:
|
||||
await self.aclose()
|
||||
raise
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
|
||||
@ -193,18 +200,23 @@ class ThreadUploadIngestionService:
|
||||
sandbox_provider = uploads.get_sandbox_provider()
|
||||
self._sync_to_sandbox = not uploads._uses_thread_data_mounts(sandbox_provider)
|
||||
if self._sync_to_sandbox:
|
||||
self._sandbox_lease = await uploads.try_acquire_sandbox_for_request(
|
||||
self._request,
|
||||
sandbox_provider,
|
||||
self._thread_id,
|
||||
user_id=self._user_id,
|
||||
app_config=self._config,
|
||||
owner_prefix="gateway:upload",
|
||||
release_on_last=True,
|
||||
)
|
||||
self._sandbox = self._sandbox_lease.sandbox
|
||||
if not self._sandbox_lease.denied and self._sandbox is None:
|
||||
raise HTTPException(status_code=500, detail="Failed to acquire sandbox")
|
||||
|
||||
async def acquire_lease() -> None:
|
||||
self._sandbox_lease = await uploads.try_acquire_sandbox_for_request(
|
||||
self._request,
|
||||
sandbox_provider,
|
||||
self._thread_id,
|
||||
user_id=self._user_id,
|
||||
app_config=self._config,
|
||||
owner_prefix="gateway:upload",
|
||||
release_on_last=True,
|
||||
)
|
||||
self._sandbox = self._sandbox_lease.sandbox
|
||||
if not self._sandbox_lease.denied and self._sandbox is None:
|
||||
raise HTTPException(status_code=500, detail="Failed to acquire sandbox")
|
||||
|
||||
# Retain the returned holder even if acquisition's awaiter is gone.
|
||||
await await_drained(acquire_lease())
|
||||
self._auto_convert = uploads._auto_convert_documents_enabled(self._config)
|
||||
|
||||
async def _link_commit_with_retry(self, uploads: Any, staged_path: Path, claimed_name: str) -> tuple[str, Path]:
|
||||
@ -220,14 +232,37 @@ class ThreadUploadIngestionService:
|
||||
the link makes the final name visible, fully formed.
|
||||
"""
|
||||
assert self._uploads_dir is not None, "open() must run before committing destinations"
|
||||
name = claimed_name
|
||||
while True:
|
||||
try:
|
||||
return name, await run_file_io(uploads._link_staged_no_overwrite, staged_path, self._uploads_dir, name)
|
||||
except FileExistsError:
|
||||
name = uploads.claim_unique_filename(name, self._seen_filenames)
|
||||
|
||||
async def publish() -> tuple[str, Path]:
|
||||
name = claimed_name
|
||||
while True:
|
||||
try:
|
||||
path = await run_file_io(uploads._link_staged_no_overwrite, staged_path, self._uploads_dir, name)
|
||||
self._written_paths.append(path)
|
||||
return name, path
|
||||
except FileExistsError:
|
||||
name = uploads.claim_unique_filename(name, self._seen_filenames)
|
||||
|
||||
# Publication and recording its cleanup owner are one admitted step.
|
||||
return await await_drained(publish())
|
||||
|
||||
async def ingest_chunks(self, chunks: AsyncIterator[bytes], *, display_name: str) -> dict[str, Any]:
|
||||
"""Ingest a cancellable stream, reclaiming session files on cancellation.
|
||||
|
||||
Only admitted filesystem steps and their ownership handoffs are drained;
|
||||
waiting for the next incoming chunk remains immediately cancellable.
|
||||
"""
|
||||
try:
|
||||
return await self._ingest_chunks(chunks, display_name=display_name)
|
||||
except asyncio.CancelledError:
|
||||
try:
|
||||
await self.cleanup_written()
|
||||
except Exception:
|
||||
with suppress(Exception):
|
||||
logger.warning("Failed to reclaim cancelled upload session files", exc_info=True)
|
||||
raise
|
||||
|
||||
async def _ingest_chunks(self, chunks: AsyncIterator[bytes], *, display_name: str) -> dict[str, Any]:
|
||||
"""Ingest one file from a chunk stream; return its wire metadata dict.
|
||||
|
||||
Owns the per-file pipeline exactly as the ordinary endpoint defines
|
||||
@ -252,8 +287,42 @@ class ThreadUploadIngestionService:
|
||||
upload_temp = None
|
||||
convert_source_fd: int | None = None
|
||||
deferred_staged_path: Path | None = None
|
||||
try:
|
||||
|
||||
async def prepare() -> None:
|
||||
nonlocal upload_temp
|
||||
upload_temp = await run_file_io(uploads._prepare_upload_destination, self._uploads_dir, safe_filename)
|
||||
|
||||
async def publish() -> Path:
|
||||
nonlocal safe_filename, upload_temp, deferred_staged_path
|
||||
while True:
|
||||
try:
|
||||
path = await run_file_io(
|
||||
uploads._commit_upload_temp_no_overwrite,
|
||||
upload_temp,
|
||||
self._uploads_dir,
|
||||
safe_filename,
|
||||
unlink_staged=convert_source_fd is None,
|
||||
)
|
||||
break
|
||||
except FileExistsError:
|
||||
safe_filename = uploads.claim_unique_filename(safe_filename, self._seen_filenames)
|
||||
self._written_paths.append(path)
|
||||
if convert_source_fd is not None:
|
||||
deferred_staged_path = upload_temp.temp_path
|
||||
upload_temp = None
|
||||
return path
|
||||
|
||||
async def abort() -> None:
|
||||
nonlocal convert_source_fd
|
||||
_close_fd(convert_source_fd)
|
||||
convert_source_fd = None
|
||||
if upload_temp is not None:
|
||||
await run_file_io(uploads._abort_upload_temp, upload_temp)
|
||||
if deferred_staged_path is not None:
|
||||
await run_file_io(uploads._remove_staged_file, deferred_staged_path)
|
||||
|
||||
try:
|
||||
await await_drained(prepare())
|
||||
async for chunk in chunks:
|
||||
file_size += len(chunk)
|
||||
self._total_size += len(chunk)
|
||||
@ -261,7 +330,7 @@ class ThreadUploadIngestionService:
|
||||
raise HTTPException(status_code=413, detail=f"File too large: {safe_filename}")
|
||||
if self._total_size > self._limits.max_total_size:
|
||||
raise HTTPException(status_code=413, detail="Total upload size too large")
|
||||
await run_file_io(uploads._write_upload_chunk, upload_temp, chunk)
|
||||
await await_drained(run_file_io(uploads._write_upload_chunk, upload_temp, chunk))
|
||||
if self._auto_convert and Path(safe_filename).suffix.lower() in uploads.CONVERTIBLE_EXTENSIONS:
|
||||
# Conversion must read the bytes this request staged. Once the
|
||||
# name is committed a sandbox process can replace it with a
|
||||
@ -278,37 +347,17 @@ class ThreadUploadIngestionService:
|
||||
# Publishing is therefore split from removing the staged name:
|
||||
# this request keeps ownership of the staged path and removes it
|
||||
# once the descriptor is released below.
|
||||
while True:
|
||||
try:
|
||||
file_path = await run_file_io(
|
||||
uploads._commit_upload_temp_no_overwrite,
|
||||
upload_temp,
|
||||
self._uploads_dir,
|
||||
safe_filename,
|
||||
unlink_staged=convert_source_fd is None,
|
||||
)
|
||||
break
|
||||
except FileExistsError:
|
||||
safe_filename = uploads.claim_unique_filename(safe_filename, self._seen_filenames)
|
||||
if convert_source_fd is not None:
|
||||
deferred_staged_path = upload_temp.temp_path
|
||||
upload_temp = None
|
||||
file_path = await await_drained(publish())
|
||||
except uploads.UnsafeUploadPathError as exc:
|
||||
_close_fd(convert_source_fd)
|
||||
if upload_temp is not None:
|
||||
await run_file_io(uploads._abort_upload_temp, upload_temp)
|
||||
await await_drained(abort())
|
||||
raise UnsafeUploadDestinationError(safe_filename) from exc
|
||||
except BaseException:
|
||||
# BaseException, not Exception: a cancellation between the
|
||||
# duplication and the commit would otherwise leave the descriptor
|
||||
# open, and nothing downstream owns it yet.
|
||||
_close_fd(convert_source_fd)
|
||||
convert_source_fd = None
|
||||
if upload_temp is not None:
|
||||
await run_file_io(uploads._abort_upload_temp, upload_temp)
|
||||
await await_drained(abort())
|
||||
raise
|
||||
|
||||
self._written_paths.append(file_path)
|
||||
virtual_path = uploads.upload_virtual_path(safe_filename)
|
||||
if self._sync_to_sandbox:
|
||||
self._sync_targets.append((file_path, virtual_path))
|
||||
@ -340,8 +389,22 @@ class ThreadUploadIngestionService:
|
||||
# the descriptor is already owned here — leaking it would also hold
|
||||
# the unlinked staged bytes until the process exits.
|
||||
private_dir: Path | None = None
|
||||
try:
|
||||
|
||||
async def prepare_conversion() -> None:
|
||||
nonlocal private_dir
|
||||
private_dir = Path(await run_file_io(tempfile.mkdtemp, "-deerflow-convert"))
|
||||
|
||||
def cleanup_conversion() -> None:
|
||||
# Release the pinned descriptor before removing either name
|
||||
# on Windows, and cover allocation/publication cancellations.
|
||||
cleanup_failed_atomic_write(md_staging, convert_source_fd)
|
||||
if deferred_staged_path is not None:
|
||||
uploads._remove_staged_file(deferred_staged_path)
|
||||
if private_dir is not None:
|
||||
shutil.rmtree(private_dir, ignore_errors=True)
|
||||
|
||||
try:
|
||||
await await_drained(prepare_conversion())
|
||||
conversion_source = private_dir / safe_filename
|
||||
# Hand the descriptor over before the call: the copy closes it
|
||||
# even when it fails, so this scope must not close it again and
|
||||
@ -354,36 +417,23 @@ class ThreadUploadIngestionService:
|
||||
staged_fd, convert_source_fd = convert_source_fd, None
|
||||
await await_drained(run_file_io(_copy_fd_to_path, staged_fd, conversion_source))
|
||||
md_staged = await uploads.convert_file_to_markdown(conversion_source, output_path=md_staging)
|
||||
except Exception:
|
||||
if not md_staged:
|
||||
self._seen_filenames.discard(unique_md_name)
|
||||
else:
|
||||
unique_md_name, md_path = await self._link_commit_with_retry(uploads, Path(md_staged), unique_md_name)
|
||||
md_virtual_path = uploads.upload_virtual_path(md_path.name)
|
||||
if self._sync_to_sandbox:
|
||||
self._sync_targets.append((md_path, md_virtual_path))
|
||||
file_info["markdown_file"] = md_path.name
|
||||
self._companion_pairs.append((file_path, md_path))
|
||||
file_info["markdown_path"] = str(self._uploads_dir / md_path.name)
|
||||
file_info["markdown_virtual_path"] = md_virtual_path
|
||||
file_info["markdown_artifact_url"] = uploads.upload_artifact_url(self._thread_id, md_path.name)
|
||||
except BaseException:
|
||||
self._seen_filenames.discard(unique_md_name)
|
||||
await run_file_io(md_staging.unlink, True)
|
||||
raise
|
||||
finally:
|
||||
_close_fd(convert_source_fd)
|
||||
if deferred_staged_path is not None:
|
||||
# The descriptor that pinned the staged inode is released
|
||||
# by now, so the deferred staged-name removal can land.
|
||||
await run_file_io(uploads._remove_staged_file, deferred_staged_path)
|
||||
if private_dir is not None:
|
||||
await run_file_io(shutil.rmtree, private_dir, True)
|
||||
if not md_staged:
|
||||
# Conversion failed and wrote nothing (or a partial staged
|
||||
# file, removed here): release the claim; holding it would
|
||||
# rename a later same-stem upload against a name nothing
|
||||
# occupies.
|
||||
self._seen_filenames.discard(unique_md_name)
|
||||
await run_file_io(md_staging.unlink, True)
|
||||
else:
|
||||
unique_md_name, md_path = await self._link_commit_with_retry(uploads, Path(md_staged), unique_md_name)
|
||||
self._written_paths.append(md_path)
|
||||
md_virtual_path = uploads.upload_virtual_path(md_path.name)
|
||||
if self._sync_to_sandbox:
|
||||
self._sync_targets.append((md_path, md_virtual_path))
|
||||
file_info["markdown_file"] = md_path.name
|
||||
self._companion_pairs.append((file_path, md_path))
|
||||
file_info["markdown_path"] = str(self._uploads_dir / md_path.name)
|
||||
file_info["markdown_virtual_path"] = md_virtual_path
|
||||
file_info["markdown_artifact_url"] = uploads.upload_artifact_url(self._thread_id, md_path.name)
|
||||
await await_drained(run_file_io(cleanup_conversion))
|
||||
return file_info
|
||||
|
||||
async def finalize(self) -> None:
|
||||
@ -394,6 +444,9 @@ class ThreadUploadIngestionService:
|
||||
stay in place; the caller does not clean them up, matching ordinary
|
||||
upload behavior for sync-phase errors).
|
||||
"""
|
||||
await await_drained(self._finalize())
|
||||
|
||||
async def _finalize(self) -> None:
|
||||
uploads = _uploads()
|
||||
# Uploaded files are created with 0o600 permissions (owner read/write
|
||||
# only). In Docker sandbox deployments the gateway writes as root but
|
||||
@ -413,12 +466,19 @@ class ThreadUploadIngestionService:
|
||||
|
||||
async def cleanup_written(self) -> None:
|
||||
"""Remove every file this session wrote (ordinary rejected-request cleanup)."""
|
||||
await await_drained(self._cleanup_written())
|
||||
|
||||
async def _cleanup_written(self) -> None:
|
||||
if not self._written_paths:
|
||||
return
|
||||
uploads = _uploads()
|
||||
await run_file_io(uploads._cleanup_uploaded_paths, self._written_paths)
|
||||
for original, _ in self._companion_pairs:
|
||||
await run_file_io(unregister_companion, original)
|
||||
try:
|
||||
await run_file_io(unregister_companion, original)
|
||||
except Exception:
|
||||
with suppress(Exception):
|
||||
logger.warning("Failed to remove upload companion record for %s", original, exc_info=True)
|
||||
self._companion_pairs = []
|
||||
self._written_paths = []
|
||||
|
||||
@ -427,10 +487,11 @@ class ThreadUploadIngestionService:
|
||||
if self._sandbox_lease is None:
|
||||
return
|
||||
try:
|
||||
await self._sandbox_lease.release()
|
||||
await await_drained(self._sandbox_lease.release())
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Failed to release sandbox request lease after upload sync: %s",
|
||||
self._sandbox_lease.sandbox_id,
|
||||
exc_info=True,
|
||||
)
|
||||
with suppress(Exception):
|
||||
logger.warning(
|
||||
"Failed to release sandbox request lease after upload sync: %s",
|
||||
self._sandbox_lease.sandbox_id,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
@ -24,6 +24,11 @@ POST /api/threads/{thread_id}/uploads
|
||||
|
||||
网关会在应用层限制上传规模,默认最多 10 个文件、单文件 50 MiB、单次请求总计 100 MiB。可通过 `config.yaml` 的 `uploads.max_files`、`uploads.max_file_size`、`uploads.max_total_size` 调整;前端会读取同一组限制并在选择文件时提示,超过限制时后端返回 `413 Payload Too Large`。
|
||||
|
||||
文件接收或转换期间取消请求,会清理本次请求已创建的文件,不删除其他请求的文件。
|
||||
等待下一段上传数据时可直接取消;已经启动的文件操作或转换需要先完成并收回资源,
|
||||
因此取消可能等待这些操作结束。开始最终权限设置和沙箱同步后,主机文件按现有
|
||||
同步失败规则保留,重试前可先检查上传列表;同步完成或失败后才会释放请求的沙箱租约。
|
||||
|
||||
文件名匹配 `.upload-*.part`(例如 `.upload-notes.part`)时,网关会返回 `400 Bad Request`,提示改名后重新上传。这是系统保留的临时文件命名规则;文件名按去掉目录后的 basename 判断,HTTP 检查在 Linux 上也同时识别 `/` 和 `\` 两种路径分隔符,例如 `folder\.upload-notes.part`。网关会先检查整批文件,再开始写入聊天的上传目录或获取沙箱,因此保留名称排在批次末尾也不会留下部分上传。聊天界面收到该错误后会提示用户,并停止本次消息发送。`.upload-notes.txt`、`notes.part` 和 `.env` 仍可上传。
|
||||
|
||||
检查也会拒绝 Windows 去掉末尾点号或空格、忽略大小写后会变成保留名称的别名,例如 `.upload-notes.part.`、`.upload-notes.part `、`.UPLOAD-NOTES.PART` 和 `.Upload-NoTeS.Part .`。这项检查在所有主机上生效,HTTP 上传会在整批文件写入前返回 `400`,提示改名。
|
||||
|
||||
@ -7,9 +7,11 @@ umask; replacements preserve existing permission bits before publication. Never
|
||||
read the process umask by changing it in this threaded runtime. Behavioral tests
|
||||
for both tool imports live in `tests/test_browser_outputs_atomic.py`.
|
||||
`cleanup_failed_atomic_write` shares best-effort cleanup across this publisher,
|
||||
`skills/projection.py`'s private manifest writer, and project document staging.
|
||||
`skills/projection.py`'s private manifest writer, project document staging, and
|
||||
Gateway upload staging.
|
||||
Use it on publication failure or after deciding to discard an owned
|
||||
temporary; pass a raw descriptor only until `fdopen` succeeds. Shelf discards
|
||||
temporary; pass a raw descriptor only until `fdopen` succeeds, or an explicitly
|
||||
transferred duplicate descriptor owned by the cleanup caller. Shelf discards
|
||||
must preserve both a rejected project's 404 and a dedup hit's 200 response.
|
||||
Callers retain their own serialization, permissions and publication policy.
|
||||
Failure cleanup must preserve the original publication exception, including
|
||||
|
||||
@ -17,6 +17,7 @@ def cleanup_failed_atomic_write(temporary: Path, fd: int | None = None) -> None:
|
||||
|
||||
Pass *fd* only while the raw descriptor is still owned by the caller;
|
||||
successful ``fdopen`` transfers that ownership to the file object.
|
||||
An explicitly transferred duplicate descriptor is also caller-owned.
|
||||
Also use this to discard staging after a publication decision. Filesystem
|
||||
and logging-handler errors remain independent best-effort steps.
|
||||
"""
|
||||
|
||||
519
backend/tests/test_upload_ingestion_cancellation.py
Normal file
519
backend/tests/test_upload_ingestion_cancellation.py
Normal file
@ -0,0 +1,519 @@
|
||||
"""Cancellation at the real upload consumer must not abandon filesystem owners."""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import shutil
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from _router_auth_helpers import call_unwrapped, make_authed_test_app
|
||||
from httpx import ASGITransport, AsyncClient
|
||||
|
||||
from app.gateway import upload_ingestion
|
||||
from app.gateway.routers import uploads
|
||||
from deerflow.utils.file_io import run_file_io
|
||||
|
||||
|
||||
class _Upload:
|
||||
def __init__(self, name, payload):
|
||||
self.filename = name
|
||||
self._chunks = [payload, b""]
|
||||
|
||||
async def read(self, size):
|
||||
return self._chunks.pop(0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stage", ["prepare", "write", "commit", "directory", "companion", "sync"])
|
||||
@pytest.mark.parametrize("cancelled", [False, True])
|
||||
def test_upload_worker_ownership_survives_request_cancellation(tmp_path, monkeypatch, stage, cancelled):
|
||||
"""Pause native work at each ownership handoff; retain uncancelled controls."""
|
||||
directory = tmp_path / "uploads"
|
||||
directory.mkdir()
|
||||
existing = directory / "keep.txt"
|
||||
existing.write_bytes(b"another request")
|
||||
remote = tmp_path / "remote"
|
||||
remote.mkdir()
|
||||
entered, release, settled = threading.Event(), threading.Event(), threading.Event()
|
||||
temporaries, conversion_dirs, worker_errors = [], [], []
|
||||
released_before_worker = []
|
||||
convert = stage in {"directory", "companion"}
|
||||
filename = "report.pdf" if convert else "report.txt"
|
||||
provider = SimpleNamespace(uses_thread_data_mounts=stage != "sync")
|
||||
monkeypatch.setattr(uploads, "ensure_uploads_dir", lambda *args, **kwargs: directory)
|
||||
monkeypatch.setattr(uploads, "get_sandbox_provider", lambda: provider)
|
||||
monkeypatch.setattr(uploads, "_auto_convert_documents_enabled", lambda _: convert)
|
||||
|
||||
def sync_file(name, content):
|
||||
(remote / Path(name).name).write_bytes(content)
|
||||
|
||||
async def release_lease():
|
||||
released_before_worker.append(entered.is_set() and not settled.is_set())
|
||||
|
||||
lease = SimpleNamespace(sandbox=SimpleNamespace(update_file=sync_file), sandbox_id="synthetic", denied=False, release=release_lease)
|
||||
monkeypatch.setattr(uploads, "try_acquire_sandbox_for_request", AsyncMock(return_value=lease))
|
||||
|
||||
async def convert_file(source, output_path):
|
||||
assert source.read_bytes() == b"owned bytes"
|
||||
await run_file_io(output_path.write_text, "converted bytes", encoding="utf-8")
|
||||
return output_path
|
||||
|
||||
monkeypatch.setattr(uploads, "convert_file_to_markdown", convert_file)
|
||||
prepare = uploads._prepare_upload_destination
|
||||
|
||||
def remember_temp(*args, **kwargs):
|
||||
temporary = prepare(*args, **kwargs)
|
||||
temporaries.append(temporary)
|
||||
return temporary
|
||||
|
||||
monkeypatch.setattr(uploads, "_prepare_upload_destination", remember_temp)
|
||||
|
||||
def pause(call, *, after=False):
|
||||
def worker(*args, **kwargs):
|
||||
try:
|
||||
value = call(*args, **kwargs) if after else None
|
||||
entered.set()
|
||||
assert release.wait(10), "test did not release the native worker"
|
||||
return value if after else call(*args, **kwargs)
|
||||
except BaseException as error:
|
||||
worker_errors.append(error)
|
||||
raise
|
||||
finally:
|
||||
settled.set()
|
||||
|
||||
return worker
|
||||
|
||||
if stage == "prepare":
|
||||
monkeypatch.setattr(uploads, "_prepare_upload_destination", pause(remember_temp, after=True))
|
||||
elif stage == "write":
|
||||
monkeypatch.setattr(uploads, "_write_upload_chunk", pause(uploads._write_upload_chunk))
|
||||
elif stage == "commit":
|
||||
# The native hard link already exists, but its returned owner has not
|
||||
# reached the coroutine that records it for rejected-request cleanup.
|
||||
monkeypatch.setattr(uploads, "_commit_upload_temp_no_overwrite", pause(uploads._commit_upload_temp_no_overwrite, after=True))
|
||||
elif stage == "directory":
|
||||
mkdtemp = upload_ingestion.tempfile.mkdtemp
|
||||
|
||||
def make_directory(*args, **kwargs):
|
||||
path = Path(mkdtemp(*args, **kwargs, dir=tmp_path))
|
||||
conversion_dirs.append(path)
|
||||
return str(path)
|
||||
|
||||
monkeypatch.setattr(upload_ingestion.tempfile, "mkdtemp", pause(make_directory, after=True))
|
||||
elif stage == "companion":
|
||||
link = uploads._link_staged_no_overwrite
|
||||
|
||||
def publish(*args, **kwargs):
|
||||
if args[2].endswith(".md"):
|
||||
return pause(link, after=True)(*args, **kwargs)
|
||||
return link(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(uploads, "_link_staged_no_overwrite", publish)
|
||||
else:
|
||||
monkeypatch.setattr(uploads, "_sync_upload_to_sandbox", pause(uploads._sync_upload_to_sandbox))
|
||||
|
||||
async def scenario():
|
||||
task = asyncio.create_task(call_unwrapped(uploads.upload_files, "thread-owner", request=MagicMock(), files=[_Upload(filename, b"owned bytes")], config=SimpleNamespace()))
|
||||
try:
|
||||
assert await asyncio.to_thread(entered.wait, 10), "owning upload did not reach the worker"
|
||||
if cancelled:
|
||||
task.cancel()
|
||||
# Give both the owning coroutine and any shield/drain task
|
||||
# their ready turns, then deliver a second host cancellation.
|
||||
await asyncio.sleep(0)
|
||||
await asyncio.sleep(0)
|
||||
task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
release.set()
|
||||
if cancelled:
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
else:
|
||||
result = await task
|
||||
assert result.success
|
||||
assert result.files[0].filename == filename
|
||||
assert await asyncio.to_thread(settled.wait, 10)
|
||||
assert existing.read_bytes() == b"another request"
|
||||
assert all(temporary.handle.closed for temporary in temporaries), "lost staged handle"
|
||||
assert all(not temporary.temp_path.exists() for temporary in temporaries), "lost staged name"
|
||||
assert all(not path.exists() for path in conversion_dirs), "lost private conversion directory"
|
||||
assert worker_errors == [], "cleanup raced a still-running native worker"
|
||||
assert not any(directory.glob(".upload-*.part")), "hidden staged companion was abandoned"
|
||||
if cancelled and stage != "sync":
|
||||
assert list(directory.iterdir()) == [existing], "cancelled ingestion left published files"
|
||||
else:
|
||||
assert (directory / filename).read_bytes() == b"owned bytes"
|
||||
if convert:
|
||||
assert (directory / "report.md").read_text() == "converted bytes"
|
||||
if stage == "sync":
|
||||
assert released_before_worker == [False], "sandbox lease released before sync worker ended"
|
||||
assert (remote / filename).read_bytes() == b"owned bytes"
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
assert await asyncio.to_thread(settled.wait, 10)
|
||||
# The red baseline intentionally leaks resources. Reclaim only
|
||||
# the test's captures after all consumer-boundary assertions.
|
||||
for temporary in temporaries:
|
||||
if not temporary.handle.closed:
|
||||
temporary.handle.close()
|
||||
temporary.temp_path.unlink(missing_ok=True)
|
||||
for path in conversion_dirs:
|
||||
shutil.rmtree(path, ignore_errors=True)
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def local_uploads(tmp_path, monkeypatch):
|
||||
directory = tmp_path / "uploads"
|
||||
directory.mkdir()
|
||||
provider = SimpleNamespace(uses_thread_data_mounts=True)
|
||||
monkeypatch.setattr(uploads, "ensure_uploads_dir", lambda *args, **kwargs: directory)
|
||||
monkeypatch.setattr(uploads, "get_sandbox_provider", lambda: provider)
|
||||
monkeypatch.setattr(uploads, "_auto_convert_documents_enabled", lambda _: False)
|
||||
return directory, provider
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stage", ["prepare", "commit"])
|
||||
def test_multipart_http_cancellation_reclaims_native_files(local_uploads, monkeypatch, stage):
|
||||
directory, _ = local_uploads
|
||||
app = make_authed_test_app()
|
||||
app.include_router(uploads.router)
|
||||
app.dependency_overrides[uploads.get_config] = lambda: SimpleNamespace()
|
||||
entered, release = threading.Event(), threading.Event()
|
||||
captured = []
|
||||
name = "_prepare_upload_destination" if stage == "prepare" else "_commit_upload_temp_no_overwrite"
|
||||
native = getattr(uploads, name)
|
||||
|
||||
def pause_after_native(*args, **kwargs):
|
||||
result = native(*args, **kwargs)
|
||||
captured.append(result)
|
||||
entered.set()
|
||||
assert release.wait(10)
|
||||
return result
|
||||
|
||||
monkeypatch.setattr(uploads, name, pause_after_native)
|
||||
|
||||
async def scenario():
|
||||
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
|
||||
task = asyncio.create_task(client.post("/api/threads/http-owner/uploads", files={"files": ("payload.txt", b"http bytes")}))
|
||||
try:
|
||||
assert await asyncio.to_thread(entered.wait, 10)
|
||||
task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
release.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
assert list(directory.iterdir()) == []
|
||||
if stage == "prepare":
|
||||
assert captured[0].handle.closed
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_waiting_for_network_chunks_is_cancellable_and_retry_is_clean(local_uploads):
|
||||
directory, _ = local_uploads
|
||||
survivor = directory / "keep.txt"
|
||||
survivor.write_bytes(b"someone else's file")
|
||||
|
||||
async def scenario():
|
||||
waiting, network = asyncio.Event(), asyncio.Event()
|
||||
|
||||
class PendingUpload:
|
||||
filename = "second.txt"
|
||||
|
||||
async def read(self, size):
|
||||
waiting.set()
|
||||
await network.wait()
|
||||
return b""
|
||||
|
||||
task = asyncio.create_task(call_unwrapped(uploads.upload_files, "network-owner", request=MagicMock(), files=[_Upload("first.txt", b"first bytes"), PendingUpload()], config=SimpleNamespace()))
|
||||
try:
|
||||
await asyncio.wait_for(waiting.wait(), 10)
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(task, 2)
|
||||
assert not network.is_set(), "cancellation must not wait for more request bytes"
|
||||
assert list(directory.iterdir()) == [survivor]
|
||||
retry = await call_unwrapped(uploads.upload_files, "network-owner", request=MagicMock(), files=[_Upload("first.txt", b"retry bytes")], config=SimpleNamespace())
|
||||
assert retry.files[0].filename == "first.txt"
|
||||
assert (directory / "first.txt").read_bytes() == b"retry bytes"
|
||||
assert survivor.read_bytes() == b"someone else's file"
|
||||
finally:
|
||||
network.set()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_cancelled_write_failure_preserves_cancellation_and_cleans_staging(local_uploads, monkeypatch):
|
||||
directory, _ = local_uploads
|
||||
entered, release, settled = threading.Event(), threading.Event(), threading.Event()
|
||||
|
||||
def fail_write(*args):
|
||||
entered.set()
|
||||
try:
|
||||
assert release.wait(10)
|
||||
raise OSError("synthetic disk write failure")
|
||||
finally:
|
||||
settled.set()
|
||||
|
||||
monkeypatch.setattr(uploads, "_write_upload_chunk", fail_write)
|
||||
|
||||
async def scenario():
|
||||
task = asyncio.create_task(call_unwrapped(uploads.upload_files, "failed-owner", request=MagicMock(), files=[_Upload("fail.txt", b"bytes")], config=SimpleNamespace()))
|
||||
try:
|
||||
assert await asyncio.to_thread(entered.wait, 10)
|
||||
task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
task.cancel()
|
||||
release.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
assert settled.is_set()
|
||||
assert list(directory.iterdir()) == []
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("failure", ["cancel", "missing_client"])
|
||||
def test_context_entry_releases_a_late_or_unusable_lease(local_uploads, monkeypatch, failure):
|
||||
_, provider = local_uploads
|
||||
provider.uses_thread_data_mounts = False
|
||||
|
||||
async def scenario():
|
||||
entered, release = asyncio.Event(), asyncio.Event()
|
||||
holder = SimpleNamespace(sandbox=object() if failure == "cancel" else None, sandbox_id="synthetic", denied=False, release=AsyncMock())
|
||||
|
||||
async def acquire(*args, **kwargs):
|
||||
entered.set()
|
||||
await release.wait()
|
||||
return holder
|
||||
|
||||
monkeypatch.setattr(uploads, "try_acquire_sandbox_for_request", acquire)
|
||||
|
||||
async def enter():
|
||||
async with upload_ingestion.ThreadUploadIngestionService(request=None, thread_id="lease-owner", user_id="u", app_config=SimpleNamespace()):
|
||||
pytest.fail("failed/cancelled entry must not yield")
|
||||
|
||||
task = asyncio.create_task(enter())
|
||||
try:
|
||||
await entered.wait()
|
||||
if failure == "cancel":
|
||||
task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
task.cancel()
|
||||
release.set()
|
||||
if failure == "cancel":
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
else:
|
||||
with pytest.raises(uploads.HTTPException) as error:
|
||||
await task
|
||||
assert error.value.status_code == 500
|
||||
holder.release.assert_awaited_once()
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_request_close_drains_release_under_repeated_cancellation():
|
||||
async def scenario():
|
||||
entered, release, closed = asyncio.Event(), asyncio.Event(), asyncio.Event()
|
||||
|
||||
async def release_lease():
|
||||
entered.set()
|
||||
await release.wait()
|
||||
closed.set()
|
||||
|
||||
service = upload_ingestion.ThreadUploadIngestionService(request=None, thread_id="close-owner", user_id="u", app_config=SimpleNamespace())
|
||||
service._sandbox_lease = SimpleNamespace(release=release_lease)
|
||||
task = asyncio.create_task(service.aclose())
|
||||
try:
|
||||
await entered.wait()
|
||||
task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
task.cancel()
|
||||
release.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
assert closed.is_set(), "request holder escaped its release"
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("failure", ["close", "staging_unlink", "published_unlink", "companion_unlink"])
|
||||
@pytest.mark.parametrize("broken_logging", [False, True])
|
||||
@pytest.mark.parametrize("outcome", ["cancel", "cancel_repeated", "too_large"])
|
||||
def test_rejected_request_preserves_outcome_when_cleanup_fails(local_uploads, monkeypatch, failure, broken_logging, outcome):
|
||||
"""Exercise cleanup failures through the real multi-file upload consumer."""
|
||||
directory, _ = local_uploads
|
||||
survivor = directory / "keep.txt"
|
||||
survivor.write_bytes(b"another request")
|
||||
monkeypatch.setattr(uploads, "_auto_convert_documents_enabled", lambda _: True)
|
||||
monkeypatch.setattr(uploads, "_get_upload_limits", lambda _: SimpleNamespace(max_files=10, max_file_size=20, max_total_size=100))
|
||||
captured, record_attempts, diagnostics = [], [], []
|
||||
cleanup_entered, cleanup_release = threading.Event(), threading.Event()
|
||||
prepare = uploads._prepare_upload_destination
|
||||
native_unlink, native_path_unlink = uploads.os.unlink, Path.unlink
|
||||
first_record = hashlib.sha256(b"first.pdf").hexdigest() + ".json"
|
||||
second_record = hashlib.sha256(b"second.pdf").hexdigest() + ".json"
|
||||
|
||||
def fail_cleanup(message):
|
||||
if outcome == "cancel_repeated":
|
||||
cleanup_entered.set()
|
||||
assert cleanup_release.wait(10), "test did not release the cleanup worker"
|
||||
raise OSError(message)
|
||||
|
||||
class FailingClose:
|
||||
def __init__(self, handle):
|
||||
self.handle = handle
|
||||
|
||||
def __getattr__(self, name):
|
||||
return getattr(self.handle, name)
|
||||
|
||||
def close(self):
|
||||
self.handle.close()
|
||||
fail_cleanup("synthetic staged close/flush failure")
|
||||
|
||||
def capture_pending(directory, name):
|
||||
staged = prepare(directory, name)
|
||||
if name == "pending.txt":
|
||||
captured.append(staged)
|
||||
if failure == "close":
|
||||
staged.handle = FailingClose(staged.handle)
|
||||
return staged
|
||||
|
||||
def unlink(path, *args, **kwargs):
|
||||
if failure == "staging_unlink" and captured and Path(path) == captured[0].temp_path:
|
||||
fail_cleanup("synthetic staging unlink failure")
|
||||
if failure == "published_unlink" and Path(path) == directory / "second.pdf":
|
||||
fail_cleanup("synthetic published unlink failure")
|
||||
return native_unlink(path, *args, **kwargs)
|
||||
|
||||
def unlink_record(path, *args, **kwargs):
|
||||
if path.parent == directory.parent / "upload-companions":
|
||||
record_attempts.append(path.name)
|
||||
if failure == "companion_unlink" and path.name == first_record:
|
||||
fail_cleanup("synthetic companion unlink failure")
|
||||
return native_path_unlink(path, *args, **kwargs)
|
||||
|
||||
def warn(message, *args, **kwargs):
|
||||
diagnostics.append(message)
|
||||
if broken_logging:
|
||||
raise RuntimeError("synthetic logging handler failure")
|
||||
|
||||
async def convert(source, output_path):
|
||||
await run_file_io(output_path.write_bytes, b"converted bytes")
|
||||
return output_path
|
||||
|
||||
monkeypatch.setattr(uploads, "_prepare_upload_destination", capture_pending)
|
||||
monkeypatch.setattr(uploads, "convert_file_to_markdown", convert)
|
||||
monkeypatch.setattr(uploads.os, "unlink", unlink)
|
||||
monkeypatch.setattr(Path, "unlink", unlink_record)
|
||||
monkeypatch.setattr(uploads.logger, "warning", warn)
|
||||
monkeypatch.setattr(upload_ingestion.logger, "warning", warn)
|
||||
error_log = MagicMock()
|
||||
monkeypatch.setattr(uploads.logger, "error", error_log)
|
||||
|
||||
async def scenario():
|
||||
waiting, network = asyncio.Event(), asyncio.Event()
|
||||
|
||||
class PendingUpload:
|
||||
filename = "pending.txt"
|
||||
|
||||
async def read(self, size):
|
||||
waiting.set()
|
||||
if outcome == "too_large":
|
||||
return b"x" * 30
|
||||
await network.wait()
|
||||
return b""
|
||||
|
||||
files = [_Upload("first.pdf", b"first"), _Upload("second.pdf", b"second"), PendingUpload()]
|
||||
task = asyncio.create_task(call_unwrapped(uploads.upload_files, "cleanup-owner", request=MagicMock(), files=files, config=SimpleNamespace()))
|
||||
try:
|
||||
await asyncio.wait_for(waiting.wait(), 10)
|
||||
if outcome != "too_large":
|
||||
task.cancel("client disconnected")
|
||||
if outcome == "cancel_repeated":
|
||||
assert await asyncio.to_thread(cleanup_entered.wait, 10)
|
||||
task.cancel("host shutdown")
|
||||
await asyncio.sleep(0)
|
||||
task.cancel("repeated host shutdown")
|
||||
cleanup_release.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
assert task.cancelled()
|
||||
else:
|
||||
with pytest.raises(uploads.HTTPException) as error:
|
||||
await task
|
||||
assert error.value.status_code == 413
|
||||
error_log.assert_not_called()
|
||||
assert diagnostics, "cleanup failure was not reported"
|
||||
assert record_attempts == [first_record, second_record], "one failed cleanup skipped the remaining companion"
|
||||
assert captured[0].handle.closed
|
||||
expected = {survivor}
|
||||
if failure == "staging_unlink":
|
||||
expected.add(captured[0].temp_path)
|
||||
if failure == "published_unlink":
|
||||
expected.add(directory / "second.pdf")
|
||||
assert set(directory.iterdir()) == expected
|
||||
assert survivor.read_bytes() == b"another request"
|
||||
finally:
|
||||
cleanup_release.set()
|
||||
network.set()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
for staged in captured:
|
||||
if not staged.handle.closed:
|
||||
staged.handle.close()
|
||||
if staged.temp_path.exists():
|
||||
native_unlink(staged.temp_path)
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_failed_lease_release_logging_preserves_request_cancellation(local_uploads, monkeypatch):
|
||||
directory, provider = local_uploads
|
||||
provider.uses_thread_data_mounts = False
|
||||
lease = SimpleNamespace(sandbox=object(), sandbox_id="synthetic", denied=False, release=AsyncMock(side_effect=OSError("synthetic lease release failure")))
|
||||
monkeypatch.setattr(uploads, "try_acquire_sandbox_for_request", AsyncMock(return_value=lease))
|
||||
monkeypatch.setattr(upload_ingestion.logger, "warning", MagicMock(side_effect=RuntimeError("synthetic logging handler failure")))
|
||||
|
||||
async def scenario():
|
||||
waiting, network = asyncio.Event(), asyncio.Event()
|
||||
|
||||
class PendingUpload:
|
||||
filename = "pending.txt"
|
||||
|
||||
async def read(self, size):
|
||||
waiting.set()
|
||||
await network.wait()
|
||||
return b""
|
||||
|
||||
task = asyncio.create_task(call_unwrapped(uploads.upload_files, "lease-owner", request=MagicMock(), files=[PendingUpload()], config=SimpleNamespace()))
|
||||
try:
|
||||
await asyncio.wait_for(waiting.wait(), 10)
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
lease.release.assert_awaited_once()
|
||||
assert list(directory.iterdir()) == []
|
||||
finally:
|
||||
network.set()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
|
||||
asyncio.run(scenario())
|
||||
Loading…
x
Reference in New Issue
Block a user