fix: 修复分片上传并发竞态和过期清理问题
This commit is contained in:
@@ -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",
|
||||
|
||||
Reference in New Issue
Block a user