diff --git a/apps/api/app/api/routes/chunked_upload.py b/apps/api/app/api/routes/chunked_upload.py index ff9a2c8d2..45664d786 100644 --- a/apps/api/app/api/routes/chunked_upload.py +++ b/apps/api/app/api/routes/chunked_upload.py @@ -3,6 +3,7 @@ Chunked upload routes for large file uploads (up to 2GB). Supports chunked upload, resume, and automatic cleanup of expired uploads. """ +import fcntl import json import logging import os @@ -64,6 +65,33 @@ def _get_upload_meta_path(upload_id: str) -> Path: return CHUNK_STORAGE_ROOT / f"{upload_id}.meta.json" +def _atomic_check_and_record(upload_id: str, chunk_index: int) -> bool: + """ + Atomically check if chunk is uploaded and record if not. + Uses file locking to prevent race conditions. + + Returns: + True if chunk was newly recorded, False if already exists + """ + meta_path = _get_upload_meta_path(upload_id) + CHUNK_STORAGE_ROOT.mkdir(parents=True, exist_ok=True) + + with open(meta_path, "r+", encoding="utf-8") as f: + fcntl.flock(f.fileno(), fcntl.LOCK_EX) + try: + meta = json.load(f) + if chunk_index in meta["uploaded_chunks"]: + return False + meta["uploaded_chunks"].append(chunk_index) + meta["status"] = "uploading" + f.seek(0) + json.dump(meta, f, ensure_ascii=False, indent=2) + f.truncate() + return True + finally: + fcntl.flock(f.fileno(), fcntl.LOCK_UN) + + def _require_workspace_member( workspace_id: str, authenticated_user: AuthenticatedUser, @@ -144,7 +172,8 @@ def _cleanup_expired_uploads() -> int: if expires_at.tzinfo is None: expires_at = expires_at.replace(tzinfo=timezone.utc) - if expires_at < now: + # Only cleanup uploads that are not actively being uploaded + if expires_at < now and meta.get("status") != "uploading": upload_id = meta["upload_id"] chunk_dir = _get_chunk_dir(upload_id) if chunk_dir.exists(): @@ -270,8 +299,8 @@ async def upload_chunk( # Verify permission _require_workspace_member(meta["workspace_id"], authenticated_user, workspace_member_repository) - # Check if already uploaded - if chunk_index in meta["uploaded_chunks"]: + # Atomic check and record to prevent race conditions + if not _atomic_check_and_record(upload_id, chunk_index): return {"message": "Chunk already uploaded", "chunk_index": chunk_index} # Read chunk data @@ -283,6 +312,19 @@ async def upload_chunk( expected_size = meta["file_size"] - (chunk_index * DEFAULT_CHUNK_SIZE) if len(chunk_data) != expected_size: + # Rollback the recorded chunk + meta_path = _get_upload_meta_path(upload_id) + with open(meta_path, "r+", encoding="utf-8") as f: + fcntl.flock(f.fileno(), fcntl.LOCK_EX) + try: + meta = json.load(f) + if chunk_index in meta["uploaded_chunks"]: + meta["uploaded_chunks"].remove(chunk_index) + f.seek(0) + json.dump(meta, f, ensure_ascii=False, indent=2) + f.truncate() + finally: + fcntl.flock(f.fileno(), fcntl.LOCK_UN) raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"Chunk size mismatch. Expected {expected_size}, got {len(chunk_data)}", @@ -293,10 +335,8 @@ async def upload_chunk( with open(chunk_path, "wb") as f: f.write(chunk_data) - # Update metadata - meta["uploaded_chunks"].append(chunk_index) - meta["status"] = "uploading" - _save_upload_meta(upload_id, meta) + # Reload metadata for response + meta = _load_upload_meta(upload_id) return { "message": "Chunk uploaded successfully",