From 6416c2787a38703afa5f70e6634c8b4c870827a4 Mon Sep 17 00:00:00 2001 From: hetaoBackend Date: Thu, 6 Aug 2026 20:55:31 +0800 Subject: [PATCH] fix: harden upload lifecycle rollback --- backend/app/gateway/routers/uploads.py | 88 +++++--- .../harness/deerflow/sandbox/sandbox.py | 10 + .../harness/deerflow/uploads/async_helpers.py | 56 ++++-- .../harness/deerflow/uploads/conversion.py | 60 +++--- .../harness/deerflow/uploads/layout.py | 18 ++ .../harness/deerflow/uploads/lease.py | 189 ++++++++++++++++-- .../harness/deerflow/uploads/manager.py | 71 +++++-- backend/tests/test_upload_conversion.py | 117 +++++++++++ backend/tests/test_uploads_manager.py | 127 ++++++++++++ backend/tests/test_uploads_router.py | 134 ++++++++++++- 10 files changed, 768 insertions(+), 102 deletions(-) diff --git a/backend/app/gateway/routers/uploads.py b/backend/app/gateway/routers/uploads.py index 6d7a69710..a0c40ea46 100644 --- a/backend/app/gateway/routers/uploads.py +++ b/backend/app/gateway/routers/uploads.py @@ -15,6 +15,7 @@ from deerflow.config.app_config import AppConfig from deerflow.config.paths import get_paths from deerflow.runtime.user_context import get_effective_user_id from deerflow.sandbox.sandbox_provider import SandboxProvider, get_sandbox_provider +from deerflow.uploads.async_helpers import run_upload_lease_io, wait_for_task_completion from deerflow.uploads.conversion import convert_uploaded_file_to_markdown from deerflow.uploads.layout import artifact_url_for_virtual_path, conversion_virtual_path from deerflow.uploads.manager import ( @@ -180,6 +181,29 @@ def _cleanup_published_uploads( logger.warning("Failed to roll back published upload after rejected request: %s", publication.path, exc_info=True) +def _cleanup_synced_sandbox_paths(sandbox, virtual_paths: list[str]) -> None: + """Best-effort removal of the exact remote copies completed by this request.""" + if sandbox is None or not virtual_paths: + return + for virtual_path in reversed(virtual_paths): + try: + sandbox.remove_file(virtual_path) + except Exception: + logger.warning("Failed to remove synchronized sandbox upload path: %s", virtual_path, exc_info=True) + + +def _rollback_upload_request( + sandbox, + synced_sandbox_paths: list[str], + publications: list[PublishedUpload], + generated_paths: list[os.PathLike[str] | str], +) -> None: + try: + _cleanup_synced_sandbox_paths(sandbox, synced_sandbox_paths) + finally: + _cleanup_published_uploads(publications, generated_paths) + + def _release_publications(publications: list[PublishedUpload]) -> None: for publication in reversed(publications): try: @@ -197,12 +221,7 @@ def _rollback_and_release_publication(publication: PublishedUpload) -> None: async def _run_file_io_cancellation_safe(function, *args): task = asyncio.create_task(run_file_io(function, *args)) - cancelled = False - while not task.done(): - try: - await asyncio.shield(task) - except asyncio.CancelledError: - cancelled = True + cancelled = await wait_for_task_completion(task) result = task.result() if cancelled: raise asyncio.CancelledError @@ -210,24 +229,14 @@ async def _run_file_io_cancellation_safe(function, *args): async def _publish_staged_upload_cancellation_safe(staged: StagedUpload, filename: str) -> PublishedUpload: - publish_task = asyncio.create_task(run_file_io(publish_staged_upload_leased, staged, filename)) + publish_task = asyncio.create_task(run_upload_lease_io(publish_staged_upload_leased, staged, filename)) try: return await asyncio.shield(publish_task) except asyncio.CancelledError: - while not publish_task.done(): - try: - await asyncio.shield(publish_task) - except asyncio.CancelledError: - continue - except Exception: - break + await wait_for_task_completion(publish_task) if not publish_task.cancelled() and publish_task.exception() is None: cleanup_task = asyncio.create_task(run_file_io(_rollback_and_release_publication, publish_task.result())) - while not cleanup_task.done(): - try: - await asyncio.shield(cleanup_task) - except asyncio.CancelledError: - continue + await wait_for_task_completion(cleanup_task) try: cleanup_task.result() except Exception: @@ -240,9 +249,15 @@ def _make_uploaded_paths_sandbox_readable(paths: list[os.PathLike[str] | str]) - _make_file_sandbox_readable(file_path) -def _sync_upload_to_sandbox(sandbox, file_path: os.PathLike[str] | str, virtual_path: str) -> None: +def _sync_upload_to_sandbox( + sandbox, + file_path: os.PathLike[str] | str, + virtual_path: str, + synced_sandbox_paths: list[str], +) -> None: _make_file_sandbox_writable(file_path) sandbox.update_file(virtual_path, Path(file_path).read_bytes()) + synced_sandbox_paths.append(virtual_path) def _list_uploaded_files_for_thread(thread_id: str, user_id: str) -> dict: @@ -333,6 +348,7 @@ async def upload_files( publications: list[PublishedUpload] = [] generated_paths: list[Path] = [] sandbox_sync_targets = [] + synced_sandbox_paths: list[str] = [] skipped_files = [] total_size = 0 sandbox_provider = get_sandbox_provider() @@ -414,7 +430,13 @@ async def upload_files( if sync_to_sandbox: for file_path, virtual_path in sandbox_sync_targets: - await _run_file_io_cancellation_safe(_sync_upload_to_sandbox, sandbox, file_path, virtual_path) + await _run_file_io_cancellation_safe( + _sync_upload_to_sandbox, + sandbox, + file_path, + virtual_path, + synced_sandbox_paths, + ) message = f"Successfully uploaded {len(uploaded_files)} file(s)" if skipped_files: @@ -427,14 +449,32 @@ async def upload_files( skipped_files=skipped_files, ) except HTTPException: - await _run_file_io_cancellation_safe(_cleanup_published_uploads, publications, generated_paths) + await _run_file_io_cancellation_safe( + _rollback_upload_request, + sandbox, + synced_sandbox_paths, + publications, + generated_paths, + ) raise except Exception as exc: logger.error("Failed to upload %s: %s", current_filename, exc) - await _run_file_io_cancellation_safe(_cleanup_published_uploads, publications, generated_paths) + await _run_file_io_cancellation_safe( + _rollback_upload_request, + sandbox, + synced_sandbox_paths, + publications, + generated_paths, + ) raise HTTPException(status_code=500, detail=f"Failed to upload {current_filename}: {str(exc)}") from exc except BaseException: - await _run_file_io_cancellation_safe(_cleanup_published_uploads, publications, generated_paths) + await _run_file_io_cancellation_safe( + _rollback_upload_request, + sandbox, + synced_sandbox_paths, + publications, + generated_paths, + ) raise finally: await _run_file_io_cancellation_safe(_release_publications, publications) diff --git a/backend/packages/harness/deerflow/sandbox/sandbox.py b/backend/packages/harness/deerflow/sandbox/sandbox.py index 9ca966d57..552938649 100644 --- a/backend/packages/harness/deerflow/sandbox/sandbox.py +++ b/backend/packages/harness/deerflow/sandbox/sandbox.py @@ -1,4 +1,5 @@ import re +import shlex from abc import ABC, abstractmethod from deerflow.sandbox.search import GrepMatch @@ -180,3 +181,12 @@ class Sandbox(ABC): content: The binary content to write to the file. """ pass + + def remove_file(self, path: str) -> None: + """Remove one exact sandbox file using the provider's virtual-path mapping.""" + resolver = getattr(self, "_resolve_path", None) + resolved = resolver(path) if callable(resolver) else path + marker = "__DEERFLOW_REMOVE_FILE_OK__" + output = self.execute_command(f"rm -f -- {shlex.quote(resolved)} && printf '%s' {marker}") + if marker not in str(output): + raise OSError(f"Sandbox did not confirm removal of {path}") diff --git a/backend/packages/harness/deerflow/uploads/async_helpers.py b/backend/packages/harness/deerflow/uploads/async_helpers.py index dc911959b..e256f5c5d 100644 --- a/backend/packages/harness/deerflow/uploads/async_helpers.py +++ b/backend/packages/harness/deerflow/uploads/async_helpers.py @@ -1,7 +1,13 @@ """Cancellation-safe async adapters for blocking upload publication APIs.""" import asyncio +import atexit +import contextvars +import functools import logging +import os +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor from pathlib import Path from deerflow.uploads.manager import ( @@ -12,6 +18,33 @@ from deerflow.uploads.manager import ( logger = logging.getLogger(__name__) +_UPLOAD_LEASE_EXECUTOR = ThreadPoolExecutor( + max_workers=min(32, (os.cpu_count() or 1) + 4), + thread_name_prefix="upload-lease-wait", +) +atexit.register(_UPLOAD_LEASE_EXECUTOR.shutdown, wait=False, cancel_futures=True) + + +async def run_upload_lease_io[**P, T](func: Callable[P, T], /, *args: P.args, **kwargs: P.kwargs) -> T: + """Run work that may wait for an upload lease outside general I/O pools.""" + loop = asyncio.get_running_loop() + context = contextvars.copy_context() + call = functools.partial(func, *args, **kwargs) + return await loop.run_in_executor(_UPLOAD_LEASE_EXECUTOR, context.run, call) + + +async def wait_for_task_completion(task: asyncio.Task) -> bool: + """Drain *task* despite cancellation and report whether cancellation arrived.""" + cancelled = False + while not task.done(): + try: + await asyncio.shield(task) + except asyncio.CancelledError: + cancelled = True + except BaseException: + break + return cancelled + def _rollback_and_release(publication: PublishedUpload) -> None: try: @@ -20,16 +53,6 @@ def _rollback_and_release(publication: PublishedUpload) -> None: publication.release() -async def _drain_task(task: asyncio.Task) -> None: - while not task.done(): - try: - await asyncio.shield(task) - except asyncio.CancelledError: - continue - except BaseException: - break - - async def publish_upload_bytes_leased_async( base_dir: Path, preferred_filename: str, @@ -37,19 +60,19 @@ async def publish_upload_bytes_leased_async( ) -> PublishedUpload: """Publish bytes off-thread without leaking a lease when cancelled.""" publish_task = asyncio.create_task( - asyncio.to_thread(publish_upload_bytes_leased, base_dir, preferred_filename, data), + run_upload_lease_io(publish_upload_bytes_leased, base_dir, preferred_filename, data), name=f"publish-upload:{preferred_filename}", ) try: return await asyncio.shield(publish_task) except asyncio.CancelledError: - await _drain_task(publish_task) + await wait_for_task_completion(publish_task) if not publish_task.cancelled() and publish_task.exception() is None: cleanup_task = asyncio.create_task( asyncio.to_thread(_rollback_and_release, publish_task.result()), name=f"rollback-cancelled-upload:{preferred_filename}", ) - await _drain_task(cleanup_task) + await wait_for_task_completion(cleanup_task) try: cleanup_task.result() except Exception: @@ -63,12 +86,7 @@ async def release_published_upload_async(publication: PublishedUpload) -> None: asyncio.to_thread(publication.release), name=f"release-upload:{publication.path.name}", ) - cancelled = False - while not release_task.done(): - try: - await asyncio.shield(release_task) - except asyncio.CancelledError: - cancelled = True + cancelled = await wait_for_task_completion(release_task) release_task.result() if cancelled: raise asyncio.CancelledError diff --git a/backend/packages/harness/deerflow/uploads/conversion.py b/backend/packages/harness/deerflow/uploads/conversion.py index b49d2995a..2aaf877c5 100644 --- a/backend/packages/harness/deerflow/uploads/conversion.py +++ b/backend/packages/harness/deerflow/uploads/conversion.py @@ -7,6 +7,7 @@ import stat from dataclasses import dataclass from pathlib import Path +from deerflow.uploads.async_helpers import run_upload_lease_io, wait_for_task_completion from deerflow.uploads.errors import UnsafeUploadPathError from deerflow.uploads.layout import ( UnsafeConversionPathError, @@ -118,40 +119,28 @@ async def _prepare_conversion_cancellation_safe( publication: PublishedUpload | None, ) -> _PreparedConversion: prepare_task = asyncio.create_task( - asyncio.to_thread(_prepare_conversion, upload_path, publication), + run_upload_lease_io(_prepare_conversion, upload_path, publication), name=f"prepare-upload-conversion:{upload_path.name}", ) - try: - return await asyncio.shield(prepare_task) - except asyncio.CancelledError: - while not prepare_task.done(): - try: - await asyncio.shield(prepare_task) - except asyncio.CancelledError: - continue - except Exception: - break + cancelled = await wait_for_task_completion(prepare_task) + if cancelled: if not prepare_task.cancelled() and prepare_task.exception() is None: cleanup_task = asyncio.create_task( asyncio.to_thread(_discard_prepared_conversion, prepare_task.result()), name=f"discard-upload-conversion:{upload_path.name}", ) - while not cleanup_task.done(): - try: - await asyncio.shield(cleanup_task) - except asyncio.CancelledError: - continue - raise + await wait_for_task_completion(cleanup_task) + cleanup_task.result() + raise asyncio.CancelledError + return prepare_task.result() async def _run_cleanup_off_thread(function, *args) -> None: cleanup_task = asyncio.create_task(asyncio.to_thread(function, *args)) - while not cleanup_task.done(): - try: - await asyncio.shield(cleanup_task) - except asyncio.CancelledError: - continue + cancelled = await wait_for_task_completion(cleanup_task) cleanup_task.result() + if cancelled: + raise asyncio.CancelledError async def convert_uploaded_file_to_markdown( @@ -163,16 +152,35 @@ async def convert_uploaded_file_to_markdown( prepared = await _prepare_conversion_cancellation_safe(upload_path, publication) stage_consumed = False try: - result = await convert_file_to_markdown(upload_path, output_path=prepared.staged.path) + conversion_task = asyncio.create_task( + convert_file_to_markdown(upload_path, output_path=prepared.staged.path), + name=f"convert-upload:{upload_path.name}", + ) + conversion_cancelled = await wait_for_task_completion(conversion_task) + if conversion_cancelled: + raise asyncio.CancelledError + result = conversion_task.result() if result is None: await _run_cleanup_off_thread(_abort_stage_without_masking, prepared.staged) stage_consumed = True return None - converted = await asyncio.to_thread(_publish_prepared_conversion, prepared, Path(result)) + publish_task = asyncio.create_task( + asyncio.to_thread(_publish_prepared_conversion, prepared, Path(result)), + name=f"publish-upload-conversion:{upload_path.name}", + ) + publish_cancelled = await wait_for_task_completion(publish_task) + if publish_cancelled: + if not publish_task.cancelled() and publish_task.exception() is None: + stage_consumed = True + raise asyncio.CancelledError + converted = publish_task.result() stage_consumed = True return converted finally: if not stage_consumed: - await _run_cleanup_off_thread(_abort_stage_without_masking, prepared.staged) - if prepared.release_publication: + if prepared.release_publication: + await _run_cleanup_off_thread(_discard_prepared_conversion, prepared) + else: + await _run_cleanup_off_thread(_abort_stage_without_masking, prepared.staged) + elif prepared.release_publication: await _run_cleanup_off_thread(_release_publication_without_masking, prepared.publication) diff --git a/backend/packages/harness/deerflow/uploads/layout.py b/backend/packages/harness/deerflow/uploads/layout.py index 0f008d092..52d20a259 100644 --- a/backend/packages/harness/deerflow/uploads/layout.py +++ b/backend/packages/harness/deerflow/uploads/layout.py @@ -10,6 +10,7 @@ from deerflow.config.paths import VIRTUAL_PATH_PREFIX UPLOAD_CONVERSIONS_DIRNAME = ".upload-conversions" UPLOAD_LOCKS_DIRNAME = ".locks" +UPLOAD_STAGE_LOCKS_DIRNAME = "stages" class UnsafeConversionPathError(ValueError): @@ -91,6 +92,23 @@ def ensure_upload_lock_dir(uploads_dir: Path) -> Path: return lock_dir +def ensure_upload_stage_lock_dir(uploads_dir: Path) -> Path: + """Create and validate the liveness-lock directory for upload stages.""" + lock_dir = ensure_upload_lock_dir(uploads_dir) + stage_lock_dir = lock_dir / UPLOAD_STAGE_LOCKS_DIRNAME + try: + stage_lock_dir.mkdir(mode=0o700) + except FileExistsError: + pass + try: + stage_lock_stat = os.lstat(stage_lock_dir) + except FileNotFoundError as exc: + raise UnsafeConversionPathError("Upload stage lock directory disappeared") from exc + if stat.S_ISLNK(stage_lock_stat.st_mode) or not stat.S_ISDIR(stage_lock_stat.st_mode): + raise UnsafeConversionPathError("Unsafe upload stage lock directory") + return stage_lock_dir + + def existing_conversion_path_for_upload(upload_path: Path) -> Path | None: """Return an existing safe generated file owned by ``upload_path``.""" if validate_conversion_dir(upload_path.parent) is None: diff --git a/backend/packages/harness/deerflow/uploads/lease.py b/backend/packages/harness/deerflow/uploads/lease.py index 109e44ae8..80430aaea 100644 --- a/backend/packages/harness/deerflow/uploads/lease.py +++ b/backend/packages/harness/deerflow/uploads/lease.py @@ -1,15 +1,18 @@ """Cross-process leases and inode identities for published uploads.""" +import errno import hashlib +import logging import os import stat import threading +import time from dataclasses import dataclass, field from pathlib import Path from typing import BinaryIO from deerflow.uploads.errors import UnsafeUploadPathError -from deerflow.uploads.layout import ensure_upload_lock_dir +from deerflow.uploads.layout import UPLOAD_CONVERSIONS_DIRNAME, ensure_upload_lock_dir, ensure_upload_stage_lock_dir try: import fcntl @@ -17,8 +20,49 @@ except ImportError: # pragma: no cover - Windows only fcntl = None # type: ignore[assignment] import msvcrt +logger = logging.getLogger(__name__) -_LOCK_STRIPES = tuple(threading.Lock() for _ in range(64)) + +_THREAD_LOCKS_GUARD = threading.Lock() + + +@dataclass(slots=True) +class _ThreadLockEntry: + lock: threading.Lock = field(default_factory=threading.Lock) + references: int = 0 + + +_THREAD_LOCKS: dict[tuple[int, int, str], _ThreadLockEntry] = {} + + +def _acquire_thread_lock(uploads_dir: Path, filename: str) -> tuple[tuple[int, int, str], _ThreadLockEntry]: + directory_stat = os.lstat(uploads_dir) + if stat.S_ISLNK(directory_stat.st_mode) or not stat.S_ISDIR(directory_stat.st_mode): + raise UnsafeUploadPathError("Unsafe upload lease directory") + key = (directory_stat.st_dev, directory_stat.st_ino, filename) + with _THREAD_LOCKS_GUARD: + entry = _THREAD_LOCKS.get(key) + if entry is None: + entry = _ThreadLockEntry() + _THREAD_LOCKS[key] = entry + entry.references += 1 + try: + entry.lock.acquire() + except BaseException: + with _THREAD_LOCKS_GUARD: + entry.references -= 1 + if entry.references == 0 and _THREAD_LOCKS.get(key) is entry: + del _THREAD_LOCKS[key] + raise + return key, entry + + +def _release_thread_lock(key: tuple[int, int, str], entry: _ThreadLockEntry) -> None: + entry.lock.release() + with _THREAD_LOCKS_GUARD: + entry.references -= 1 + if entry.references == 0 and _THREAD_LOCKS.get(key) is entry: + del _THREAD_LOCKS[key] @dataclass(frozen=True, slots=True) @@ -72,8 +116,15 @@ def _lock_file(lock_file: BinaryIO) -> None: if fcntl is not None: fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX) return - lock_file.seek(0) - msvcrt.locking(lock_file.fileno(), msvcrt.LK_LOCK, 1) + while True: + lock_file.seek(0) + try: + msvcrt.locking(lock_file.fileno(), msvcrt.LK_NBLCK, 1) + return + except OSError as exc: + if exc.errno not in {errno.EACCES, errno.EAGAIN, errno.EDEADLK}: + raise + time.sleep(0.05) def _unlock_file(lock_file: BinaryIO) -> None: @@ -84,6 +135,116 @@ def _unlock_file(lock_file: BinaryIO) -> None: msvcrt.locking(lock_file.fileno(), msvcrt.LK_UNLCK, 1) +def _try_lock_file(lock_file: BinaryIO) -> bool: + if fcntl is not None: + try: + fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + except OSError as exc: + if exc.errno in {errno.EACCES, errno.EAGAIN}: + return False + raise + return True + lock_file.seek(0) + try: + msvcrt.locking(lock_file.fileno(), msvcrt.LK_NBLCK, 1) + except OSError as exc: + if exc.errno in {errno.EACCES, errno.EAGAIN, errno.EDEADLK}: + return False + raise + return True + + +def _stage_uploads_dir(stage_dir: Path) -> Path: + if stage_dir.name == UPLOAD_CONVERSIONS_DIRNAME: + return stage_dir.parent / "uploads" + return stage_dir + + +def _stage_lock_path(stage_dir: Path, stage_filename: str) -> Path: + digest = hashlib.sha256(f"{stage_dir.name}\0{stage_filename}".encode()).hexdigest() + return ensure_upload_stage_lock_dir(_stage_uploads_dir(stage_dir)) / f"{digest}.lock" + + +@dataclass(slots=True) +class UploadStageLease: + """Cross-process liveness lease for one hidden upload staging file.""" + + stage_dir: Path + stage_filename: str + lock_path: Path + _lock_file: BinaryIO + _identity: UploadIdentity + _active: bool = True + _state_lock: threading.Lock = field(default_factory=threading.Lock) + + @classmethod + def acquire(cls, stage_dir: Path, stage_filename: str) -> "UploadStageLease": + lock_path = _stage_lock_path(Path(stage_dir), stage_filename) + lock_file = _open_lock_file(lock_path) + try: + _lock_file(lock_file) + return cls( + stage_dir=Path(stage_dir), + stage_filename=stage_filename, + lock_path=lock_path, + _lock_file=lock_file, + _identity=UploadIdentity.from_path(lock_path), + ) + except BaseException: + lock_file.close() + raise + + @classmethod + def try_acquire(cls, stage_dir: Path, stage_filename: str) -> "UploadStageLease | None": + lock_path = _stage_lock_path(Path(stage_dir), stage_filename) + lock_file = _open_lock_file(lock_path) + try: + if not _try_lock_file(lock_file): + lock_file.close() + return None + return cls( + stage_dir=Path(stage_dir), + stage_filename=stage_filename, + lock_path=lock_path, + _lock_file=lock_file, + _identity=UploadIdentity.from_path(lock_path), + ) + except BaseException: + lock_file.close() + raise + + @property + def is_active(self) -> bool: + with self._state_lock: + return self._active + + def _remove_matching_lock_file(self) -> None: + try: + if self._identity.matches(self.lock_path): + self.lock_path.unlink(missing_ok=True) + except BaseException: + logger.warning("Failed to remove upload stage liveness file: %s", self.lock_path, exc_info=True) + + def release(self) -> None: + with self._state_lock: + if not self._active: + return + if fcntl is not None: + self._remove_matching_lock_file() + try: + _unlock_file(self._lock_file) + except BaseException: + logger.warning("Failed to unlock upload stage liveness file: %s", self.lock_path, exc_info=True) + try: + self._lock_file.close() + except BaseException: + logger.warning("Failed to close upload stage liveness file: %s", self.lock_path, exc_info=True) + finally: + self._active = False + if fcntl is None: + self._remove_matching_lock_file() + + @dataclass(slots=True) class UploadNameLease: """Exclusive thread-and-process lease for one actual upload filename.""" @@ -92,7 +253,8 @@ class UploadNameLease: filename: str lock_path: Path _lock_file: BinaryIO - _stripe: threading.Lock + _thread_lock_key: tuple[int, int, str] + _thread_lock_entry: _ThreadLockEntry _active: bool = True _state_lock: threading.Lock = field(default_factory=threading.Lock) @@ -104,25 +266,26 @@ class UploadNameLease: if len(filename.encode("utf-8")) > 255: raise UnsafeUploadPathError("Upload lease filename is too long") + uploads_dir = Path(uploads_dir) digest = hashlib.sha256(filename.encode("utf-8")).hexdigest() - stripe = _LOCK_STRIPES[int(digest[:2], 16) % len(_LOCK_STRIPES)] - stripe.acquire() + thread_lock_key, thread_lock_entry = _acquire_thread_lock(uploads_dir, filename) lock_file: BinaryIO | None = None try: - lock_path = ensure_upload_lock_dir(Path(uploads_dir)) / f"{digest}.lock" + lock_path = ensure_upload_lock_dir(uploads_dir) / f"{digest}.lock" lock_file = _open_lock_file(lock_path) _lock_file(lock_file) return cls( - uploads_dir=Path(uploads_dir), + uploads_dir=uploads_dir, filename=filename, lock_path=lock_path, _lock_file=lock_file, - _stripe=stripe, + _thread_lock_key=thread_lock_key, + _thread_lock_entry=thread_lock_entry, ) except BaseException: if lock_file is not None: lock_file.close() - stripe.release() + _release_thread_lock(thread_lock_key, thread_lock_entry) raise @property @@ -132,7 +295,7 @@ class UploadNameLease: return self._active def release(self) -> None: - """Release the OS lock and process stripe; repeated calls are harmless.""" + """Release the OS and exact-name thread locks; repeated calls are harmless.""" with self._state_lock: if not self._active: return @@ -148,7 +311,7 @@ class UploadNameLease: error = exc finally: self._active = False - self._stripe.release() + _release_thread_lock(self._thread_lock_key, self._thread_lock_entry) if error is not None: raise error diff --git a/backend/packages/harness/deerflow/uploads/manager.py b/backend/packages/harness/deerflow/uploads/manager.py index 2ab6f2a00..d2e98308d 100644 --- a/backend/packages/harness/deerflow/uploads/manager.py +++ b/backend/packages/harness/deerflow/uploads/manager.py @@ -7,9 +7,9 @@ Both Gateway and Client delegate to these functions. import errno import logging import os +import secrets import shutil import stat -import tempfile from collections.abc import Iterator from dataclasses import dataclass from pathlib import Path @@ -25,7 +25,7 @@ from deerflow.uploads.layout import ( existing_conversion_path_for_upload, upload_virtual_path, ) -from deerflow.uploads.lease import UploadIdentity, UploadNameLease +from deerflow.uploads.lease import UploadIdentity, UploadNameLease, UploadStageLease from deerflow.utils.thread_id import validate_thread_id logger = logging.getLogger(__name__) @@ -41,6 +41,7 @@ class StagedUpload: base_dir: Path path: Path handle: BinaryIO + lease: UploadStageLease @dataclass(slots=True) @@ -147,24 +148,35 @@ def _validate_upload_directory(base_dir: Path) -> Path: def create_upload_staging_file(base_dir: Path) -> StagedUpload: """Create a hidden same-directory staging file for a complete payload.""" base_dir = _validate_upload_directory(Path(base_dir)) - fd, temp_path_str = tempfile.mkstemp( - prefix=UPLOAD_STAGING_PREFIX, - suffix=UPLOAD_STAGING_SUFFIX, - dir=base_dir, - ) - temp_path = Path(temp_path_str) - try: - handle = os.fdopen(fd, "wb") - except Exception: - os.close(fd) - temp_path.unlink(missing_ok=True) - raise - return StagedUpload(base_dir=base_dir, path=temp_path, handle=handle) + while True: + temp_path = base_dir / f"{UPLOAD_STAGING_PREFIX}{secrets.token_hex(16)}{UPLOAD_STAGING_SUFFIX}" + lease = UploadStageLease.acquire(base_dir, temp_path.name) + flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL + if hasattr(os, "O_BINARY"): + flags |= os.O_BINARY + try: + fd = os.open(temp_path, flags, 0o600) + except FileExistsError: + lease.release() + continue + except BaseException: + lease.release() + raise + try: + handle = os.fdopen(fd, "wb") + except Exception: + os.close(fd) + temp_path.unlink(missing_ok=True) + lease.release() + raise + return StagedUpload(base_dir=base_dir, path=temp_path, handle=handle, lease=lease) def abort_staged_upload(staged: StagedUpload) -> None: """Close and remove a staging file, tolerating repeated cleanup.""" close_error: BaseException | None = None + unlink_error: BaseException | None = None + lease_error: BaseException | None = None try: if not staged.handle.closed: staged.handle.close() @@ -172,17 +184,31 @@ def abort_staged_upload(staged: StagedUpload) -> None: close_error = exc try: staged.path.unlink(missing_ok=True) - except BaseException as unlink_error: - if close_error is not None: - raise close_error from unlink_error - raise + except BaseException as exc: + unlink_error = exc + try: + staged.lease.release() + except BaseException as exc: + lease_error = exc if close_error is not None: + if unlink_error is not None: + raise close_error from unlink_error + if lease_error is not None: + raise close_error from lease_error raise close_error + if unlink_error is not None: + if lease_error is not None: + raise unlink_error from lease_error + raise unlink_error + if lease_error is not None: + raise lease_error def _validate_staged_upload(staged: StagedUpload) -> None: """Reject a staging path that was replaced or moved outside its directory.""" _validate_upload_directory(staged.base_dir) + if not staged.lease.is_active or staged.lease.stage_dir != staged.base_dir or staged.lease.stage_filename != staged.path.name: + raise UnsafeUploadPathError("Upload staging liveness lease is missing") if staged.path.parent.resolve() != staged.base_dir.resolve(): raise UnsafeUploadPathError("Upload staging path escaped its directory") try: @@ -272,6 +298,7 @@ def publish_staged_upload_leased(staged: StagedUpload, preferred_filename: str) except OSError as exc: _rollback_link_without_masking(candidate, staged_identity) raise AtomicUploadPublishError("Failed to remove upload staging link after publication") from exc + staged.lease.release() try: candidate_stat = os.lstat(candidate) @@ -370,6 +397,7 @@ def replace_system_owned_staged_file(staged: StagedUpload, filename: str) -> Pat _validate_staged_upload(staged) target = staged.base_dir / normalize_filename(filename) os.replace(staged.path, target) + staged.lease.release() return target @@ -429,6 +457,9 @@ def cleanup_stale_upload_staging_files(base_dir: Path | str | None = None) -> in for entry in entries: if not is_upload_staging_file(entry.name) or not entry.is_file(follow_symlinks=False): continue + stage_lease = UploadStageLease.try_acquire(uploads_dir, entry.name) + if stage_lease is None: + continue try: os.unlink(entry.path) removed += 1 @@ -436,6 +467,8 @@ def cleanup_stale_upload_staging_files(base_dir: Path | str | None = None) -> in pass except OSError: logger.warning("Failed to remove stale upload staging file: %s", entry.path, exc_info=True) + finally: + stage_lease.release() except FileNotFoundError: continue except OSError: diff --git a/backend/tests/test_upload_conversion.py b/backend/tests/test_upload_conversion.py index 694a4f333..847a3cc34 100644 --- a/backend/tests/test_upload_conversion.py +++ b/backend/tests/test_upload_conversion.py @@ -2,6 +2,7 @@ import asyncio import hashlib +import threading from pathlib import Path from unittest.mock import AsyncMock, patch @@ -17,6 +18,122 @@ from deerflow.uploads.layout import ( from deerflow.uploads.manager import delete_file_safe, publish_upload_bytes, publish_upload_bytes_leased +@pytest.mark.asyncio +async def test_cancellation_waits_for_converter_before_cleanup_and_lease_release(tmp_path): + uploads = tmp_path / "user-data" / "uploads" + uploads.mkdir(parents=True) + upload = publish_upload_bytes(uploads, "report.pdf", b"OLD") + converter_started = threading.Event() + allow_converter = threading.Event() + + def paused_sync_convert(_source, output_path): + converter_started.set() + assert allow_converter.wait(5) + output_path.write_text("converted", encoding="utf-8") + return output_path + + with patch("deerflow.utils.file_conversion._convert_file_to_markdown_sync", side_effect=paused_sync_convert): + conversion = asyncio.create_task(convert_uploaded_file_to_markdown(upload)) + assert await asyncio.to_thread(converter_started.wait, 5) + conversion.cancel() + deletion = asyncio.create_task(asyncio.to_thread(delete_file_safe, uploads, upload.name)) + await asyncio.sleep(0.05) + assert not conversion.done() + assert not deletion.done() + allow_converter.set() + with pytest.raises(asyncio.CancelledError): + await conversion + await deletion + + conversion_dir = uploads.parent / ".upload-conversions" + assert not list(conversion_dir.glob(".upload-*.part")) + assert not conversion_path_for_upload(upload).exists() + + +@pytest.mark.asyncio +async def test_cancellation_waits_for_conversion_publication_worker(tmp_path): + import deerflow.uploads.conversion as conversion_module + + uploads = tmp_path / "user-data" / "uploads" + uploads.mkdir(parents=True) + upload = publish_upload_bytes(uploads, "report.pdf", b"OLD") + publish_started = threading.Event() + allow_publish = threading.Event() + real_publish = conversion_module._publish_prepared_conversion + + async def fake_convert(_source, output_path=None): + output_path.write_text("converted", encoding="utf-8") + return output_path + + def paused_publish(prepared, result): + publish_started.set() + assert allow_publish.wait(5) + return real_publish(prepared, result) + + with ( + patch("deerflow.uploads.conversion.convert_file_to_markdown", side_effect=fake_convert), + patch("deerflow.uploads.conversion._publish_prepared_conversion", side_effect=paused_publish), + ): + conversion = asyncio.create_task(convert_uploaded_file_to_markdown(upload)) + assert await asyncio.to_thread(publish_started.wait, 5) + conversion.cancel() + deletion = asyncio.create_task(asyncio.to_thread(delete_file_safe, uploads, upload.name)) + await asyncio.sleep(0.05) + assert not conversion.done() + assert not deletion.done() + allow_publish.set() + with pytest.raises(asyncio.CancelledError): + await conversion + await deletion + + assert not conversion_path_for_upload(upload).exists() + + +@pytest.mark.asyncio +async def test_repeated_cancellation_still_releases_standalone_conversion_lease(tmp_path): + import deerflow.uploads.conversion as conversion_module + + uploads = tmp_path / "user-data" / "uploads" + uploads.mkdir(parents=True) + upload = publish_upload_bytes(uploads, "report.pdf", b"OLD") + converter_started = asyncio.Event() + allow_converter = asyncio.Event() + cleanup_started = threading.Event() + allow_cleanup = threading.Event() + real_abort = conversion_module._abort_stage_without_masking + + async def paused_convert(_source, output_path=None): + converter_started.set() + await allow_converter.wait() + output_path.write_text("converted", encoding="utf-8") + return output_path + + def paused_abort(staged): + cleanup_started.set() + assert allow_cleanup.wait(5) + real_abort(staged) + + with ( + patch("deerflow.uploads.conversion.convert_file_to_markdown", side_effect=paused_convert), + patch("deerflow.uploads.conversion._abort_stage_without_masking", side_effect=paused_abort), + ): + conversion = asyncio.create_task(convert_uploaded_file_to_markdown(upload)) + await converter_started.wait() + conversion.cancel() + allow_converter.set() + assert await asyncio.to_thread(cleanup_started.wait, 5) + conversion.cancel() + deletion = asyncio.create_task(asyncio.to_thread(delete_file_safe, uploads, upload.name)) + await asyncio.sleep(0.05) + assert not deletion.done() + allow_cleanup.set() + with pytest.raises(asyncio.CancelledError): + await conversion + await asyncio.wait_for(deletion, timeout=2) + + assert not conversion_path_for_upload(upload).exists() + + @pytest.mark.asyncio async def test_delete_and_reupload_cannot_receive_old_conversion(tmp_path): uploads = tmp_path / "user-data" / "uploads" diff --git a/backend/tests/test_uploads_manager.py b/backend/tests/test_uploads_manager.py index 8d7867f17..f8a43d7d5 100644 --- a/backend/tests/test_uploads_manager.py +++ b/backend/tests/test_uploads_manager.py @@ -17,6 +17,7 @@ from deerflow.uploads.layout import ( conversion_path_for_upload, conversion_virtual_path, ) +from deerflow.uploads.lease import UploadNameLease from deerflow.uploads.manager import ( AtomicUploadPublishError, PathTraversalError, @@ -54,6 +55,32 @@ def _delete_upload_in_process( finished.set() +def _hold_staged_upload_in_process( + uploads_dir: str, + started: Any, + release: Any, + staged_paths: Any, + errors: Any, +) -> None: + staged = None + try: + staged = create_upload_staging_file(Path(uploads_dir)) + staged.handle.write(b"in progress") + staged.handle.flush() + staged_paths.put(str(staged.path)) + started.set() + if not release.wait(5): + raise TimeoutError("parent did not release the staged upload") + except BaseException as exc: # pragma: no cover - surfaced in the parent + errors.put(repr(exc)) + finally: + if staged is not None: + try: + abort_staged_upload(staged) + except BaseException as exc: # pragma: no cover - surfaced in the parent + errors.put(repr(exc)) + + # --------------------------------------------------------------------------- # normalize_filename # --------------------------------------------------------------------------- @@ -146,6 +173,52 @@ class TestValidatePathTraversal: class TestUploadPublication: + def test_unrelated_names_that_shared_a_legacy_stripe_do_not_block(self, tmp_path): + first = UploadNameLease.acquire(tmp_path, "f0.txt") + second = None + pool = ThreadPoolExecutor(max_workers=1) + try: + future = pool.submit(UploadNameLease.acquire, tmp_path, "f15.txt") + second = future.result(timeout=1) + finally: + first.release() + if second is not None: + second.release() + pool.shutdown() + + def test_windows_lock_retries_until_the_holder_releases(self, monkeypatch): + import deerflow.uploads.lease as lease_module + + attempts = 0 + + class FakeMsvcrt: + LK_NBLCK = 1 + + @staticmethod + def locking(_fd, mode, _length): + nonlocal attempts + assert mode == FakeMsvcrt.LK_NBLCK + attempts += 1 + if attempts < 3: + raise OSError(errno.EACCES, "locked") + + class FakeLockFile: + @staticmethod + def fileno(): + return 7 + + @staticmethod + def seek(_offset): + return None + + monkeypatch.setattr(lease_module, "fcntl", None) + monkeypatch.setattr(lease_module, "msvcrt", FakeMsvcrt, raising=False) + monkeypatch.setattr(lease_module.time, "sleep", lambda _seconds: None) + + lease_module._lock_file(FakeLockFile()) + + assert attempts == 3 + def test_reserved_staging_name_is_rejected_before_stage_creation(self, tmp_path): with patch("deerflow.uploads.manager.create_upload_staging_file") as create_stage: with pytest.raises(ValueError, match="reserved"): @@ -283,6 +356,16 @@ class TestUploadPublication: assert not staged.path.exists() + def test_abort_releases_stage_lease_when_unlink_raises(self, tmp_path): + staged = create_upload_staging_file(tmp_path) + + with patch.object(Path, "unlink", autospec=True, side_effect=OSError("unlink failed")): + with pytest.raises(OSError, match="unlink failed"): + abort_staged_upload(staged) + + assert not staged.lease.is_active + staged.path.unlink() + def test_compatibility_wrapper_writes_new_file(self, tmp_path): dest = write_upload_file_no_symlink(tmp_path, "notes.txt", b"hello") @@ -460,6 +543,50 @@ class TestListFilesInDir: class TestCleanupStaleUploadStagingFiles: + def test_skips_stage_held_by_another_process(self, tmp_path): + uploads = tmp_path / "threads" / "thread-live" / "user-data" / "uploads" + uploads.mkdir(parents=True) + context = multiprocessing.get_context("spawn") + started = context.Event() + release = context.Event() + staged_paths = context.Queue() + errors = context.Queue() + worker = context.Process( + target=_hold_staged_upload_in_process, + args=(str(uploads), started, release, staged_paths, errors), + ) + worker.start() + try: + assert started.wait(5) + staged_path = Path(staged_paths.get(timeout=1)) + + assert cleanup_stale_upload_staging_files(tmp_path) == 0 + assert staged_path.exists() + finally: + release.set() + worker.join(timeout=5) + if worker.is_alive(): + worker.terminate() + worker.join(timeout=5) + + assert worker.exitcode == 0 + with pytest.raises(Empty): + errors.get_nowait() + + def test_skips_live_stage_and_removes_it_after_lease_is_abandoned(self, tmp_path): + uploads = tmp_path / "threads" / "thread-live" / "user-data" / "uploads" + uploads.mkdir(parents=True) + staged = create_upload_staging_file(uploads) + staged.handle.write(b"in progress") + + assert cleanup_stale_upload_staging_files(tmp_path) == 0 + assert staged.path.exists() + + staged.handle.close() + staged.lease.release() + assert cleanup_stale_upload_staging_files(tmp_path) == 1 + assert not staged.path.exists() + def test_removes_only_stale_staging_files_from_all_upload_layouts(self, tmp_path): legacy_uploads = tmp_path / "threads" / "thread-legacy" / "user-data" / "uploads" user_uploads = tmp_path / "users" / "owner-1" / "threads" / "thread-owned" / "user-data" / "uploads" diff --git a/backend/tests/test_uploads_router.py b/backend/tests/test_uploads_router.py index e82e0b75a..ec11a293b 100644 --- a/backend/tests/test_uploads_router.py +++ b/backend/tests/test_uploads_router.py @@ -2,6 +2,7 @@ import asyncio import os import stat import threading +from concurrent.futures import ThreadPoolExecutor from io import BytesIO from pathlib import Path from types import SimpleNamespace @@ -14,8 +15,10 @@ from fastapi.testclient import TestClient from app.gateway.deps import get_config from app.gateway.routers import uploads +from deerflow.sandbox.sandbox import Sandbox from deerflow.uploads.layout import conversion_path_for_upload -from deerflow.uploads.manager import delete_file_safe, publish_upload_bytes_leased +from deerflow.uploads.lease import UploadNameLease +from deerflow.uploads.manager import create_upload_staging_file, delete_file_safe, publish_upload_bytes_leased class ChunkedUpload: @@ -65,6 +68,25 @@ def _fake_owned_conversion(content_by_source: dict[str, str] | None = None): return fake_convert +def test_sandbox_remove_file_uses_provider_virtual_path_resolution(): + commands: list[str] = [] + + class FakeSandbox: + @staticmethod + def _resolve_path(path: str) -> str: + assert path == "/mnt/user-data/uploads/report final.pdf" + return "/home/sandbox/uploads/report final.pdf" + + @staticmethod + def execute_command(command: str) -> str: + commands.append(command) + return "__DEERFLOW_REMOVE_FILE_OK__" + + Sandbox.remove_file(FakeSandbox(), "/mnt/user-data/uploads/report final.pdf") + + assert commands == ["rm -f -- '/home/sandbox/uploads/report final.pdf' && printf '%s' __DEERFLOW_REMOVE_FILE_OK__"] + + def test_cleanup_uses_publication_identity_not_reused_path(tmp_path): publication = publish_upload_bytes_leased(tmp_path, "report.pdf", b"old") try: @@ -155,6 +177,116 @@ def test_sandbox_sync_failure_rolls_back_published_generation(tmp_path): assert not (thread_uploads_dir / "notes.txt").exists() +def test_partial_sandbox_sync_failure_removes_only_completed_remote_paths(tmp_path): + thread_uploads_dir = tmp_path / "uploads" + thread_uploads_dir.mkdir(parents=True) + provider = MagicMock() + provider.uses_thread_data_mounts = False + provider.acquire_async = AsyncMock(return_value="remote-1") + sandbox = MagicMock() + synced_paths: list[str] = [] + + def update_file(virtual_path: str, _data: bytes) -> None: + if virtual_path.endswith("second.txt"): + raise RuntimeError("second sync failed") + synced_paths.append(virtual_path) + + sandbox.update_file.side_effect = update_file + provider.get.return_value = sandbox + + with ( + patch.object(uploads, "ensure_uploads_dir", return_value=thread_uploads_dir), + patch.object(uploads, "get_sandbox_provider", return_value=provider), + ): + with pytest.raises(HTTPException) as exc_info: + asyncio.run( + call_unwrapped( + uploads.upload_files, + "thread-remote", + request=MagicMock(), + files=[ + UploadFile(filename="first.txt", file=BytesIO(b"first")), + UploadFile(filename="second.txt", file=BytesIO(b"second")), + ], + config=SimpleNamespace(), + ) + ) + + assert exc_info.value.status_code == 500 + assert synced_paths == ["/mnt/user-data/uploads/first.txt"] + sandbox.remove_file.assert_called_once_with("/mnt/user-data/uploads/first.txt") + assert not (thread_uploads_dir / "first.txt").exists() + assert not (thread_uploads_dir / "second.txt").exists() + + +@pytest.mark.asyncio +async def test_cancellation_after_remote_sync_still_removes_the_completed_copy(tmp_path): + thread_uploads_dir = tmp_path / "uploads" + thread_uploads_dir.mkdir(parents=True) + provider = MagicMock() + provider.uses_thread_data_mounts = False + provider.acquire_async = AsyncMock(return_value="remote-1") + sandbox = MagicMock() + sync_started = threading.Event() + allow_sync = threading.Event() + + def update_file(_virtual_path: str, _data: bytes) -> None: + sync_started.set() + assert allow_sync.wait(5) + + sandbox.update_file.side_effect = update_file + provider.get.return_value = sandbox + + with ( + patch.object(uploads, "ensure_uploads_dir", return_value=thread_uploads_dir), + patch.object(uploads, "get_sandbox_provider", return_value=provider), + ): + upload_task = asyncio.create_task( + call_unwrapped( + uploads.upload_files, + "thread-remote", + request=MagicMock(), + files=[UploadFile(filename="notes.txt", file=BytesIO(b"payload"))], + config=SimpleNamespace(), + ) + ) + assert await asyncio.to_thread(sync_started.wait, 5) + upload_task.cancel() + allow_sync.set() + with pytest.raises(asyncio.CancelledError): + await upload_task + + sandbox.remove_file.assert_called_once_with("/mnt/user-data/uploads/notes.txt") + assert not (thread_uploads_dir / "notes.txt").exists() + + +@pytest.mark.asyncio +async def test_waiting_publication_cannot_starve_lease_release(tmp_path, monkeypatch): + import deerflow.utils.file_io as file_io_module + + first = UploadNameLease.acquire(tmp_path, "report.pdf") + staged = create_upload_staging_file(tmp_path) + staged.handle.write(b"second") + single_worker = ThreadPoolExecutor(max_workers=1) + monkeypatch.setattr(file_io_module, "_FILE_IO_EXECUTOR", single_worker) + + publication_task = asyncio.create_task(uploads._publish_staged_upload_cancellation_safe(staged, "report.pdf")) + await asyncio.sleep(0.05) + release_task = asyncio.create_task(uploads._run_file_io_cancellation_safe(first.release)) + second = None + try: + await asyncio.wait_for(asyncio.shield(release_task), timeout=0.2) + finally: + if first.is_active: + first.release() + await release_task + second = await asyncio.wait_for(publication_task, timeout=2) + second.release() + single_worker.shutdown(wait=True) + + assert second.path.name == "report.pdf" + + def test_upload_files_writes_thread_storage_and_skips_local_sandbox_sync(tmp_path): thread_uploads_dir = tmp_path / "uploads" thread_uploads_dir.mkdir(parents=True)