fix: 修复分片上传并发竞态和过期清理问题
CI/CD Pipeline / Validate Code Quality And Tests (push) Waiting to run
CI/CD Pipeline / Validate Code Quality And Tests (pull_request) Has been cancelled

This commit is contained in:
CI Test
2026-06-26 20:55:11 +08:00
parent e44ebc535e
commit 81828af62b
+47 -7
View File
@@ -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",