mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-15 00:19:14 +00:00
fix: harden upload lifecycle rollback
This commit is contained in:
hetaoBackend
parent
a540ff46a9
commit
6416c2787a
@ -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)
|
||||
|
||||
@ -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}")
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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"
|
||||
|
||||
@ -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"
|
||||
|
||||
@ -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)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user