Spaces:
Running
Running
| from __future__ import annotations | |
| import asyncio | |
| import os | |
| import shutil | |
| import time | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| from uuid import UUID | |
| from app.core.config import Settings | |
| from app.core.exceptions import NotFoundError, ProcessingError | |
| from app.core.logger import get_logger | |
| logger = get_logger(__name__) | |
| class RequestWorkspace: | |
| request_id: str | |
| root: Path | |
| uploads: Path | |
| outputs: Path | |
| logs: Path | |
| class CleanupService: | |
| """Owns per-request directories and expires completed or abandoned work.""" | |
| def __init__(self, settings: Settings) -> None: | |
| self.settings = settings | |
| self._active: set[str] = set() | |
| self._lock = asyncio.Lock() | |
| async def create_workspace(self, request_id: str) -> RequestWorkspace: | |
| self._validate_request_id(request_id) | |
| root = self.settings.temp_dir / request_id | |
| workspace = RequestWorkspace( | |
| request_id=request_id, | |
| root=root, | |
| uploads=root / "uploads", | |
| outputs=root / "outputs", | |
| logs=root / "logs", | |
| ) | |
| for directory in (workspace.uploads, workspace.outputs, workspace.logs): | |
| directory.mkdir(parents=True, exist_ok=True) | |
| async with self._lock: | |
| self._active.add(request_id) | |
| return workspace | |
| async def complete(self, request_id: str) -> None: | |
| async with self._lock: | |
| self._active.discard(request_id) | |
| for base in (self.settings.temp_dir, self.settings.output_dir): | |
| path = base / request_id | |
| if path.exists(): | |
| await asyncio.to_thread(os.utime, path, None) | |
| async def publish(self, request_id: str, source: Path, filename: str) -> Path: | |
| self._validate_request_id(request_id) | |
| safe_name = Path(filename).name | |
| if not safe_name or safe_name in {".", ".."}: | |
| raise ProcessingError("The generated output filename is invalid") | |
| destination_dir = self.settings.output_dir / request_id | |
| destination_dir.mkdir(parents=True, exist_ok=True) | |
| destination = destination_dir / safe_name | |
| try: | |
| await asyncio.to_thread(os.replace, source, destination) | |
| except OSError: | |
| await asyncio.to_thread(shutil.move, str(source), str(destination)) | |
| return destination | |
| async def publish_new(self, request_id: str, source: Path, filename: str) -> Path | None: | |
| """Publish a generated output without replacing an existing file. | |
| This is used by durable generation-job reconciliation, where a retry | |
| must never overwrite a canonical output created by an earlier worker | |
| attempt. Existing media operations continue to use ``publish`` and | |
| retain their established replacement semantics. | |
| """ | |
| self._validate_request_id(request_id) | |
| safe_name = Path(filename).name | |
| if not safe_name or safe_name in {".", ".."}: | |
| raise ProcessingError("The generated output filename is invalid") | |
| destination_dir = self.settings.output_dir / request_id | |
| destination_dir.mkdir(parents=True, exist_ok=True) | |
| destination = destination_dir / safe_name | |
| created = await asyncio.to_thread(self._publish_new_sync, source, destination) | |
| return destination if created else None | |
| def _publish_new_sync(source: Path, destination: Path) -> bool: | |
| """Atomically claim a destination, with a cross-device fallback.""" | |
| try: | |
| os.link(source, destination) | |
| except FileExistsError: | |
| return False | |
| except OSError: | |
| # ``temp_dir`` and ``output_dir`` can be different mounts. An | |
| # exclusive create still prevents overwrite in that arrangement. | |
| try: | |
| with source.open("rb") as input_stream, destination.open("xb") as output_stream: | |
| shutil.copyfileobj(input_stream, output_stream, length=1024 * 1024) | |
| except FileExistsError: | |
| return False | |
| except OSError: | |
| # Never leave a partial file reachable from the output | |
| # directory when cross-device publication fails. | |
| try: | |
| destination.unlink(missing_ok=True) | |
| except OSError: | |
| pass | |
| raise | |
| try: | |
| source.unlink() | |
| except FileNotFoundError: | |
| pass | |
| return True | |
| def resolve_download(self, request_id: str, filename: str) -> Path: | |
| self._validate_request_id(request_id) | |
| if filename != Path(filename).name: | |
| raise NotFoundError("Output file not found") | |
| root = (self.settings.output_dir / request_id).resolve() | |
| candidate = (root / filename).resolve() | |
| if candidate.parent != root or not candidate.is_file(): | |
| raise NotFoundError("Output file not found") | |
| return candidate | |
| async def cleanup_expired(self) -> int: | |
| cutoff = time.time() - self.settings.cleanup_minutes * 60 | |
| async with self._lock: | |
| active = self._active.copy() | |
| removed = 0 | |
| for base in (self.settings.temp_dir, self.settings.output_dir): | |
| if not base.exists(): | |
| continue | |
| for path in list(base.iterdir()): | |
| if not path.is_dir() or path.name in active: | |
| continue | |
| try: | |
| if path.stat().st_mtime < cutoff: | |
| await asyncio.to_thread(shutil.rmtree, path) | |
| removed += 1 | |
| logger.info("expired workspace removed", extra={"path": str(path)}) | |
| except FileNotFoundError: | |
| continue | |
| except OSError as exc: | |
| logger.warning( | |
| "workspace cleanup failed", | |
| extra={"path": str(path), "error": str(exc)}, | |
| ) | |
| return removed | |
| async def remove_request(self, request_id: str) -> None: | |
| self._validate_request_id(request_id) | |
| async with self._lock: | |
| self._active.discard(request_id) | |
| for base in (self.settings.temp_dir, self.settings.output_dir): | |
| path = base / request_id | |
| if path.is_dir(): | |
| await asyncio.to_thread(shutil.rmtree, path) | |
| async def remove_temporary_request(self, request_id: str) -> None: | |
| """Remove only bounded staging data while retaining published output.""" | |
| self._validate_request_id(request_id) | |
| async with self._lock: | |
| self._active.discard(request_id) | |
| path = self.settings.temp_dir / request_id | |
| if path.is_dir(): | |
| await asyncio.to_thread(shutil.rmtree, path) | |
| def _validate_request_id(request_id: str) -> None: | |
| try: | |
| UUID(request_id) | |
| except (ValueError, AttributeError) as exc: | |
| raise NotFoundError("Output file not found") from exc | |