diff --git a/.env.example b/.env.example index 541bad7..5980fb4 100644 --- a/.env.example +++ b/.env.example @@ -23,6 +23,11 @@ WORKER_CONCURRENCY=4 # 单进程图片处理并发上限（API 与 Worker 均生效，默认等于 CPU 线程数） IMAGE_PROCESSING_CONCURRENCY=4 +# ZIP 使用任务租约做 single-flight；总大小按解压前源文件字节计算。 +ZIP_BUILD_CONCURRENCY=2 +ZIP_MAX_ENTRIES=200 +ZIP_MAX_UNCOMPRESSED_BYTES=2147483648 + # 仅当后端只能由可信反向代理访问时启用，否则客户端可伪造来源 IP TRUST_PROXY_HEADERS=false diff --git a/docker/.env.production.example b/docker/.env.production.example index 590eccf..9b437cd 100644 --- a/docker/.env.production.example +++ b/docker/.env.production.example @@ -27,6 +27,11 @@ WORKER_TASK_CONCURRENCY=4 WORKER_CONCURRENCY=2 IMAGE_PROCESSING_CONCURRENCY=4 +# Two concurrent 2 GiB ZIP builds require roughly 8 GiB of temporary disk. +ZIP_BUILD_CONCURRENCY=2 +ZIP_MAX_ENTRIES=200 +ZIP_MAX_UNCOMPRESSED_BYTES=2147483648 + # Resource ceilings tuned for an 8-core / 16 GB application host. POSTGRES_MEMORY_LIMIT=2g REDIS_MEMORY_LIMIT=1g diff --git a/docker/docker-compose.prod.yml b/docker/docker-compose.prod.yml index d0c4606..ef44028 100644 --- a/docker/docker-compose.prod.yml +++ b/docker/docker-compose.prod.yml @@ -13,6 +13,9 @@ x-imageforge-environment: &imageforge-environment WORKER_TASK_CONCURRENCY: ${WORKER_TASK_CONCURRENCY:-4} WORKER_CONCURRENCY: ${WORKER_CONCURRENCY:-2} IMAGE_PROCESSING_CONCURRENCY: ${IMAGE_PROCESSING_CONCURRENCY:-2} + ZIP_BUILD_CONCURRENCY: ${ZIP_BUILD_CONCURRENCY:-2} + ZIP_MAX_ENTRIES: ${ZIP_MAX_ENTRIES:-200} + ZIP_MAX_UNCOMPRESSED_BYTES: ${ZIP_MAX_UNCOMPRESSED_BYTES:-2147483648} ALLOW_ANONYMOUS_UPLOAD: ${ALLOW_ANONYMOUS_UPLOAD:-true} ANON_MAX_FILE_SIZE_MB: ${ANON_MAX_FILE_SIZE_MB:-5} ANON_MAX_FILES_PER_BATCH: ${ANON_MAX_FILES_PER_BATCH:-5} diff --git a/docs/architecture.md b/docs/architecture.md index c951b2e..313bd99 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -278,6 +278,7 @@ dotenvy = "0.15" - 图片压缩使用 `spawn_blocking` 避免阻塞异步线程 - `WORKER_TASK_CONCURRENCY` 控制任务级并发，避免大批量任务独占 Worker - `WORKER_CONCURRENCY` 控制单任务内文件并发，`IMAGE_PROCESSING_CONCURRENCY` 作为进程级 CPU 闸门 +- `ZIP_BUILD_CONCURRENCY` 是 API 进程级 ZIP 闸门；数据库租约保证同一任务跨实例只构建一次，`ZIP_MAX_ENTRIES` 和 `ZIP_MAX_UNCOMPRESSED_BYTES` 在下载源对象前拒绝超预算任务 ```rust // 在独立线程池中执行 CPU 密集型压缩 @@ -290,6 +291,7 @@ let result = tokio::task::spawn_blocking(move || { - 流式处理大文件 - 限制并发压缩任务数 - 压缩完成后立即清理临时文件 +- ZIP attempt 使用独立临时目录和对象键；发布 CAS 失败时立即删除，两项默认并发且每项 2 GiB 上限时应至少预留约 8 GiB 临时磁盘余量 ### 3. 缓存策略 - Redis 缓存用户会话 diff --git a/docs/s3-storage-plan.md b/docs/s3-storage-plan.md index ecafcb8..992f21b 100644 --- a/docs/s3-storage-plan.md +++ b/docs/s3-storage-plan.md @@ -66,7 +66,7 @@ flowchart LR | 低级会员 Pro | 7 天 | `results/7d/`、`archives/7d/` | 9 天 | | 高级会员 Business | 15 天 | `results/15d/`、`archives/15d/` | 17 天 | -Worker 每 5 分钟按 `expires_at` 精确删除对象，删除成功后才删除数据库任务。S3 生命周期多保留 2 天，只负责处理数据库故障、进程崩溃或上传后未能落库的孤儿对象，不能作为精确会员权限判断。未完成的分片上传 1 天后由生命周期中止。 +Worker 每 5 分钟按 `expires_at` 精确删除对象，删除成功后才删除数据库任务。ZIP 构建 attempt 位于对应 `archives//.../attempts/` 前缀，发布失败会立即删除，进程崩溃遗留项仍由同一前缀生命周期兜底。S3 生命周期多保留 2 天，只负责处理数据库故障、进程崩溃或上传后未能落库的孤儿对象，不能作为精确会员权限判断。未完成的分片上传 1 天后由生命周期中止。 ## 5. 119 首期容量 diff --git a/migrations/021_zip_build_leases.sql b/migrations/021_zip_build_leases.sql new file mode 100644 index 0000000..bfcdad6 --- /dev/null +++ b/migrations/021_zip_build_leases.sql @@ -0,0 +1,12 @@ +ALTER TABLE tasks + ADD COLUMN zip_build_token UUID, + ADD COLUMN zip_build_lease_until TIMESTAMPTZ, + ADD COLUMN zip_build_attempt BIGINT NOT NULL DEFAULT 0; + +ALTER TABLE tasks + ADD CONSTRAINT tasks_zip_build_lease_pair_check + CHECK ((zip_build_token IS NULL) = (zip_build_lease_until IS NULL)); + +CREATE INDEX idx_tasks_zip_build_lease + ON tasks(zip_build_lease_until) + WHERE zip_storage_key IS NULL AND zip_build_token IS NOT NULL; diff --git a/src/api/billing.rs b/src/api/billing.rs index 600e325..ee460ba 100644 --- a/src/api/billing.rs +++ b/src/api/billing.rs @@ -1091,6 +1091,7 @@ mod tests { AppState { mailer: Arc::new(Mailer::new(&config).expect("create disabled test mailer")), image_processing_semaphore: Arc::new(Semaphore::new(2)), + zip_build_semaphore: Arc::new(Semaphore::new(2)), runtime_policy_cache: crate::services::settings::RuntimePolicyCache::new(), storage_cache: crate::services::storage::StorageCache::new(), config, diff --git a/src/api/downloads.rs b/src/api/downloads.rs index cbcaccb..e0e76c5 100644 --- a/src/api/downloads.rs +++ b/src/api/downloads.rs @@ -241,8 +241,27 @@ struct TaskZipFileRow { storage_key: Option, original_name: String, output_format: String, + compressed_size: Option, } +#[derive(Debug, FromRow)] +struct ZipBuildStateRow { + zip_storage_backend: Option, + zip_storage_endpoint_id: Option, + zip_storage_key: Option, + expires_at: DateTime, +} + +#[derive(Debug)] +enum ZipBuildClaim { + Acquired { token: Uuid }, + Cached(storage::ObjectLocator), + Busy, +} + +const ZIP_BUILD_LEASE_SECONDS: i64 = 15 * 60; +const ZIP_BUILD_WAIT_SECONDS: u64 = 30; + async fn download_task_zip( State(state): State, jar: axum_extra::extract::cookie::CookieJar, @@ -315,11 +334,31 @@ async fn download_task_zip( .await; } + let object = resolve_task_zip(&state, task_id, task.retention_hours as i64).await?; + + respond_object( + &state, + jar, + &object, + &format!("task_{task_id}.zip"), + "application/zip", + ) + .await +} + +async fn resolve_task_zip( + state: &AppState, + task_id: Uuid, + retention_hours: i64, +) -> Result { + if let Some(object) = load_published_zip(state, task_id).await? { + return Ok(object); + } let rows = sqlx::query_as::<_, TaskZipFileRow>( r#" SELECT storage_backend, storage_endpoint_id, COALESCE(storage_key, storage_path) AS storage_key, - original_name, output_format + original_name, output_format, compressed_size FROM task_files WHERE task_id = $1 AND status = 'completed' ORDER BY created_at ASC @@ -333,93 +372,386 @@ async fn download_task_zip( if rows.is_empty() { return Err(AppError::new(ErrorCode::NotFound, "没有可打包的文件")); } + validate_zip_budget(state, &rows)?; - let temp_dir = PathBuf::from(format!( - "{}/tmp/zips/{task_id}-{}", - state.config.storage_path, - Uuid::new_v4() - )); - tokio::fs::create_dir_all(&temp_dir).await.map_err(|err| { - AppError::new(ErrorCode::StorageUnavailable, "创建 ZIP 临时目录失败").with_source(err) - })?; - let zip_path = temp_dir.join(format!("task_{task_id}.zip")); - - let build_result: Result = async { - let mut used_names: HashMap = HashMap::new(); - let mut entries: Vec<(String, String)> = Vec::new(); - for (index, row) in rows.into_iter().enumerate() { - let Some(key) = row.storage_key else { - continue; - }; - let path = temp_dir.join(format!("entry-{index}")); - storage::download_to_file( - &state, - &storage::ObjectLocator { - backend: row.storage_backend, - endpoint_id: row.storage_endpoint_id, - key, - }, - &path, - ) - .await?; - let name = - build_zip_entry_name(&row.original_name, &row.output_format, &mut used_names); - entries.push((name, path.to_string_lossy().to_string())); + let deadline = + tokio::time::Instant::now() + std::time::Duration::from_secs(ZIP_BUILD_WAIT_SECONDS); + loop { + match claim_zip_build(state, task_id).await? { + ZipBuildClaim::Cached(object) => return Ok(object), + ZipBuildClaim::Acquired { token } => { + return build_claimed_zip(state, task_id, retention_hours, token, &rows).await; + } + ZipBuildClaim::Busy if tokio::time::Instant::now() < deadline => { + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } + ZipBuildClaim::Busy => { + return Err(AppError::new( + ErrorCode::StorageUnavailable, + "ZIP 正在生成，请稍后重试", + )); + } } - if entries.is_empty() { - return Err(AppError::new(ErrorCode::NotFound, "没有可打包的文件")); - } - - let zip_path_cloned = zip_path.clone(); - let task_id_str = task_id.to_string(); - tokio::task::spawn_blocking(move || { - generate_zip_file(&zip_path_cloned, &task_id_str, &entries) - }) - .await - .map_err(|err| AppError::new(ErrorCode::Internal, "生成 ZIP 失败").with_source(err))? - .map_err(|err| AppError::new(ErrorCode::Internal, "生成 ZIP 失败").with_source(err))?; - - let object_key = storage::archive_key(task.retention_hours as i64, task_id); - storage::store_file(&state, &object_key, &zip_path, "application/zip").await } - .await; +} - let _ = tokio::fs::remove_dir_all(&temp_dir).await; - let stored = build_result?; +fn validate_zip_budget(state: &AppState, rows: &[TaskZipFileRow]) -> Result { + if rows.len() > state.config.zip_max_entries as usize { + return Err(AppError::new( + ErrorCode::FileTooLarge, + format!("ZIP 文件数量超过 {} 个上限", state.config.zip_max_entries), + )); + } + rows.iter().try_fold(0_u64, |total, row| { + if row.storage_key.is_none() { + return Err(AppError::new( + ErrorCode::StorageUnavailable, + "ZIP 源文件存储信息不完整", + )); + } + let size = row + .compressed_size + .and_then(|value| u64::try_from(value).ok()) + .ok_or_else(|| AppError::new(ErrorCode::StorageUnavailable, "ZIP 源文件大小无效"))?; + let next = total + .checked_add(size) + .ok_or_else(|| AppError::new(ErrorCode::FileTooLarge, "ZIP 源文件总大小超出限制"))?; + if next > state.config.zip_max_uncompressed_bytes { + return Err(AppError::new( + ErrorCode::FileTooLarge, + format!( + "ZIP 源文件总大小超过 {} 字节上限", + state.config.zip_max_uncompressed_bytes + ), + )); + } + Ok(next) + }) +} - sqlx::query( +async fn claim_zip_build(state: &AppState, task_id: Uuid) -> Result { + let token = Uuid::new_v4(); + let acquired: Option = sqlx::query_scalar( r#" UPDATE tasks - SET zip_storage_backend = $2, - zip_storage_endpoint_id = $3, - zip_storage_key = $4, - zip_storage_etag = $5, - zip_size = $6 - WHERE id = $1 AND zip_storage_key IS NULL + SET zip_build_token = $2, + zip_build_lease_until = NOW() + ($3 * INTERVAL '1 second'), + zip_build_attempt = zip_build_attempt + 1 + WHERE id = $1 + AND zip_storage_key IS NULL + AND completed_at IS NOT NULL + AND expires_at > NOW() + AND ( + zip_build_token IS NULL + OR zip_build_lease_until <= NOW() + ) + RETURNING zip_build_attempt "#, ) .bind(task_id) + .bind(token) + .bind(ZIP_BUILD_LEASE_SECONDS) + .fetch_optional(&state.db) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "获取 ZIP 构建租约失败").with_source(err))?; + if acquired.is_some() { + return Ok(ZipBuildClaim::Acquired { token }); + } + + let current = sqlx::query_as::<_, ZipBuildStateRow>( + r#" + SELECT zip_storage_backend, zip_storage_endpoint_id, zip_storage_key, expires_at + FROM tasks + WHERE id = $1 + "#, + ) + .bind(task_id) + .fetch_optional(&state.db) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "查询 ZIP 构建状态失败").with_source(err))? + .ok_or_else(|| AppError::new(ErrorCode::NotFound, "任务不存在"))?; + if current.expires_at <= Utc::now() { + return Err(AppError::new(ErrorCode::NotFound, "任务已过期或不存在")); + } + if let (Some(backend), Some(key)) = (current.zip_storage_backend, current.zip_storage_key) { + return Ok(ZipBuildClaim::Cached(storage::ObjectLocator { + backend, + endpoint_id: current.zip_storage_endpoint_id, + key, + })); + } + // A lease may have been released between the UPDATE and this read. The + // caller retries the atomic claim after a short wait in either case. + Ok(ZipBuildClaim::Busy) +} + +async fn renew_zip_build(state: &AppState, task_id: Uuid, token: Uuid) -> Result<(), AppError> { + let updated = sqlx::query( + r#" + UPDATE tasks + SET zip_build_lease_until = NOW() + ($3 * INTERVAL '1 second') + WHERE id = $1 + AND zip_build_token = $2 + AND zip_storage_key IS NULL + AND expires_at > NOW() + "#, + ) + .bind(task_id) + .bind(token) + .bind(ZIP_BUILD_LEASE_SECONDS) + .execute(&state.db) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "续租 ZIP 构建失败").with_source(err))?; + if updated.rows_affected() != 1 { + return Err(AppError::new( + ErrorCode::StorageUnavailable, + "ZIP 构建租约已失效，请重试", + )); + } + Ok(()) +} + +async fn release_zip_build(state: &AppState, task_id: Uuid, token: Uuid) { + if let Err(err) = sqlx::query( + r#" + UPDATE tasks + SET zip_build_token = NULL, + zip_build_lease_until = NULL + WHERE id = $1 + AND zip_build_token = $2 + AND zip_storage_key IS NULL + "#, + ) + .bind(task_id) + .bind(token) + .execute(&state.db) + .await + { + tracing::warn!(task_id = %task_id, zip_build_token = %token, error = %err, "failed to release ZIP build lease"); + } +} + +async fn build_claimed_zip( + state: &AppState, + task_id: Uuid, + retention_hours: i64, + token: Uuid, + rows: &[TaskZipFileRow], +) -> Result { + let permit = match tokio::time::timeout( + std::time::Duration::from_secs(ZIP_BUILD_WAIT_SECONDS), + state.zip_build_semaphore.clone().acquire_owned(), + ) + .await + { + Ok(Ok(permit)) => permit, + Ok(Err(err)) => { + release_zip_build(state, task_id, token).await; + return Err(AppError::new(ErrorCode::Internal, "ZIP 并发闸门已关闭").with_source(err)); + } + Err(_) => { + release_zip_build(state, task_id, token).await; + return Err(AppError::new( + ErrorCode::StorageUnavailable, + "ZIP 生成繁忙，请稍后重试", + )); + } + }; + + let temp_dir = PathBuf::from(format!( + "{}/tmp/zips/{task_id}-{token}", + state.config.storage_path + )); + let zip_path = temp_dir.join(format!("task_{task_id}.zip")); + let build_result = build_zip_attempt( + state, + task_id, + token, + retention_hours, + rows, + &temp_dir, + &zip_path, + ) + .await; + drop(permit); + let _ = tokio::fs::remove_dir_all(&temp_dir).await; + + let stored = match build_result { + Ok(stored) => stored, + Err(err) => { + release_zip_build(state, task_id, token).await; + return Err(err); + } + }; + publish_zip_attempt(state, task_id, token, stored).await +} + +async fn build_zip_attempt( + state: &AppState, + task_id: Uuid, + token: Uuid, + retention_hours: i64, + rows: &[TaskZipFileRow], + temp_dir: &std::path::Path, + zip_path: &std::path::Path, +) -> Result { + tokio::fs::create_dir_all(temp_dir).await.map_err(|err| { + AppError::new(ErrorCode::StorageUnavailable, "创建 ZIP 临时目录失败").with_source(err) + })?; + let mut used_names: HashMap = HashMap::new(); + let mut entries: Vec<(String, String)> = Vec::with_capacity(rows.len()); + let mut actual_bytes = 0_u64; + for (index, row) in rows.iter().enumerate() { + renew_zip_build(state, task_id, token).await?; + let key = row.storage_key.as_ref().ok_or_else(|| { + AppError::new(ErrorCode::StorageUnavailable, "ZIP 源文件存储信息不完整") + })?; + let path = temp_dir.join(format!("entry-{index}")); + storage::download_to_file( + state, + &storage::ObjectLocator { + backend: row.storage_backend.clone(), + endpoint_id: row.storage_endpoint_id, + key: key.clone(), + }, + &path, + ) + .await?; + let size = tokio::fs::metadata(&path) + .await + .map_err(|err| { + AppError::new(ErrorCode::StorageUnavailable, "读取 ZIP 临时文件失败") + .with_source(err) + })? + .len(); + actual_bytes = actual_bytes + .checked_add(size) + .ok_or_else(|| AppError::new(ErrorCode::FileTooLarge, "ZIP 实际文件总大小超出限制"))?; + if actual_bytes > state.config.zip_max_uncompressed_bytes { + return Err(AppError::new( + ErrorCode::FileTooLarge, + "ZIP 实际文件总大小超出限制", + )); + } + let name = build_zip_entry_name(&row.original_name, &row.output_format, &mut used_names); + entries.push((name, path.to_string_lossy().to_string())); + } + + renew_zip_build(state, task_id, token).await?; + let zip_path_cloned = zip_path.to_path_buf(); + let task_id_str = task_id.to_string(); + tokio::task::spawn_blocking(move || { + generate_zip_file(&zip_path_cloned, &task_id_str, &entries) + }) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "生成 ZIP 失败").with_source(err))? + .map_err(|err| AppError::new(ErrorCode::Internal, "生成 ZIP 失败").with_source(err))?; + + renew_zip_build(state, task_id, token).await?; + let object_key = storage::archive_attempt_key(retention_hours, task_id, token); + storage::store_file(state, &object_key, zip_path, "application/zip").await +} + +async fn publish_zip_attempt( + state: &AppState, + task_id: Uuid, + token: Uuid, + stored: storage::StoredObject, +) -> Result { + let published = sqlx::query( + r#" + UPDATE tasks + SET zip_storage_backend = $3, + zip_storage_endpoint_id = $4, + zip_storage_key = $5, + zip_storage_etag = $6, + zip_size = $7, + zip_build_token = NULL, + zip_build_lease_until = NULL + WHERE id = $1 + AND zip_build_token = $2 + AND zip_storage_key IS NULL + AND expires_at > NOW() + "#, + ) + .bind(task_id) + .bind(token) .bind(&stored.backend) .bind(stored.endpoint_id) .bind(&stored.key) .bind(&stored.etag) .bind(stored.size as i64) .execute(&state.db) - .await - .map_err(|err| AppError::new(ErrorCode::Internal, "记录 ZIP 对象失败").with_source(err))?; + .await; - respond_object( - &state, - jar, - &storage::ObjectLocator { + match published { + Ok(result) if result.rows_affected() == 1 => Ok(storage::ObjectLocator { backend: stored.backend, endpoint_id: stored.endpoint_id, key: stored.key, - }, - &format!("task_{task_id}.zip"), - "application/zip", + }), + Ok(_) => { + delete_unpublished_zip(state, task_id, token, &stored).await; + let current = load_published_zip(state, task_id).await?; + current.ok_or_else(|| { + AppError::new(ErrorCode::StorageUnavailable, "ZIP 发布租约已失效，请重试") + }) + } + Err(err) => { + delete_unpublished_zip(state, task_id, token, &stored).await; + Err(AppError::new(ErrorCode::Internal, "记录 ZIP 对象失败").with_source(err)) + } + } +} + +async fn delete_unpublished_zip( + state: &AppState, + task_id: Uuid, + token: Uuid, + stored: &storage::StoredObject, +) { + let object = storage::ObjectLocator { + backend: stored.backend.clone(), + endpoint_id: stored.endpoint_id, + key: stored.key.clone(), + }; + let mut last_error = None; + for attempt in 1..=3_u64 { + match storage::delete_object(state, &object).await { + Ok(()) => { + last_error = None; + break; + } + Err(err) => { + last_error = Some(err); + tokio::time::sleep(std::time::Duration::from_millis(100 * attempt)).await; + } + } + } + if let Some(err) = last_error { + tracing::error!(task_id = %task_id, zip_build_token = %token, object_key = %stored.key, error = %err, "failed to delete unpublished ZIP attempt after retries"); + } + release_zip_build(state, task_id, token).await; +} + +async fn load_published_zip( + state: &AppState, + task_id: Uuid, +) -> Result, AppError> { + let row: Option<(Option, Option, Option)> = sqlx::query_as( + "SELECT zip_storage_backend, zip_storage_endpoint_id, zip_storage_key FROM tasks WHERE id = $1", ) + .bind(task_id) + .fetch_optional(&state.db) .await + .map_err(|err| AppError::new(ErrorCode::Internal, "查询已发布 ZIP 失败").with_source(err))?; + Ok(row.and_then(|(backend, endpoint_id, key)| { + Some(storage::ObjectLocator { + backend: backend?, + endpoint_id, + key: key?, + }) + })) } fn build_zip_entry_name( @@ -497,6 +829,12 @@ fn generate_zip_file( #[cfg(test)] mod tests { use super::*; + use crate::config::Config; + use crate::services::mail::Mailer; + use crate::services::settings; + use sqlx::postgres::PgPoolOptions; + use std::sync::Arc; + use tokio::sync::{Barrier, Semaphore}; #[test] fn content_disposition_supports_unicode_names() { @@ -525,4 +863,321 @@ mod tests { assert_eq!(output_file_name("photo.png", "webp"), "photo.webp"); assert_eq!(output_file_name("没有扩展名", "jpeg"), "没有扩展名.jpg"); } + + async fn build_zip_test_state( + pool: sqlx::PgPool, + database_url: String, + redis_url: String, + storage_path: String, + ) -> AppState { + let mut config = Config::from_env().expect("load ZIP test config"); + config.database_url = database_url; + config.redis_url = redis_url; + config.storage_path = storage_path; + config.zip_build_concurrency = 2; + config.zip_max_entries = 200; + config.zip_max_uncompressed_bytes = 2 * 1024 * 1024; + config.mail_enabled = false; + config.mail_log_links_when_disabled = false; + let redis = redis::Client::open(config.redis_url.clone()) + .expect("create ZIP test Redis client") + .get_connection_manager() + .await + .expect("connect ZIP test Redis"); + AppState { + mailer: Arc::new(Mailer::new(&config).expect("create disabled ZIP test mailer")), + image_processing_semaphore: Arc::new(Semaphore::new(2)), + zip_build_semaphore: Arc::new(Semaphore::new(2)), + runtime_policy_cache: crate::services::settings::RuntimePolicyCache::new(), + storage_cache: storage::StorageCache::new(), + config, + db: pool, + redis, + } + } + + async fn insert_zip_task( + pool: &sqlx::PgPool, + task_id: Uuid, + marker: &str, + input_path: &std::path::Path, + recorded_size: i64, + ) { + sqlx::query( + r#" + INSERT INTO tasks ( + id, session_id, status, total_files, completed_files, + total_original_size, total_compressed_size, + started_at, completed_at, expires_at, retention_hours + ) VALUES ( + $1, $2, 'completed', 1, 1, + $3, $3, NOW(), NOW(), NOW() + INTERVAL '1 day', 24 + ) + "#, + ) + .bind(task_id) + .bind(format!("zip-session-{marker}")) + .bind(recorded_size) + .execute(pool) + .await + .expect("insert ZIP test task"); + sqlx::query( + r#" + INSERT INTO task_files ( + id, task_id, original_name, original_format, output_format, + original_size, compressed_size, saved_percent, + storage_path, storage_backend, storage_key, + status, completed_at + ) VALUES ( + $1, $2, $3, 'png', 'png', + $4, $4, 0, + $5, 'local', $5, + 'completed', NOW() + ) + "#, + ) + .bind(Uuid::new_v4()) + .bind(task_id) + .bind(format!("{marker}.png")) + .bind(recorded_size) + .bind(input_path.to_string_lossy().to_string()) + .execute(pool) + .await + .expect("insert ZIP test file"); + } + + async fn configure_test_s3(state: &AppState, marker: &str) -> Option { + let endpoint = std::env::var("IMAGEFORGE_TEST_S3_ENDPOINT").ok()?; + let bucket = std::env::var("IMAGEFORGE_TEST_S3_BUCKET").ok()?; + let access_key = std::env::var("IMAGEFORGE_TEST_S3_ACCESS_KEY").ok()?; + let secret_key = std::env::var("IMAGEFORGE_TEST_S3_SECRET_KEY").ok()?; + let endpoint_id = Uuid::new_v4(); + let encrypted_access = + settings::encrypt_secret(state, &access_key).expect("encrypt ZIP test S3 access key"); + let encrypted_secret = + settings::encrypt_secret(state, &secret_key).expect("encrypt ZIP test S3 secret key"); + sqlx::query( + r#" + INSERT INTO storage_endpoints ( + id, name, internal_endpoint, public_endpoint, + bucket, region, + access_key_encrypted, secret_key_encrypted, access_key_hint, + force_path_style, is_active + ) VALUES ( + $1, $2, $3, $3, + $4, 'us-east-1', + $5, $6, 'test', + true, true + ) + "#, + ) + .bind(endpoint_id) + .bind(format!("zip-test-{marker}")) + .bind(endpoint) + .bind(bucket) + .bind(encrypted_access) + .bind(encrypted_secret) + .execute(&state.db) + .await + .expect("insert ZIP test S3 endpoint"); + Some(endpoint_id) + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + #[ignore = "requires isolated IMAGEFORGE_TEST_DATABASE_URL and IMAGEFORGE_TEST_REDIS_URL; optional IMAGEFORGE_TEST_S3_* uses MinIO"] + async fn zip_build_is_single_flight_bounded_and_fenced() { + let database_url = std::env::var("IMAGEFORGE_TEST_DATABASE_URL") + .expect("IMAGEFORGE_TEST_DATABASE_URL must be set"); + assert!( + database_url.to_ascii_lowercase().contains("test"), + "refusing to run destructive integration test outside a test database" + ); + let redis_url = std::env::var("IMAGEFORGE_TEST_REDIS_URL") + .expect("IMAGEFORGE_TEST_REDIS_URL must be set"); + let pool = PgPoolOptions::new() + .max_connections(32) + .connect(&database_url) + .await + .expect("connect ZIP test database"); + sqlx::migrate!().run(&pool).await.expect("run migrations"); + let marker = Uuid::new_v4().simple().to_string(); + let storage_root = std::env::temp_dir().join(format!("imageforge-zip-test-{marker}")); + tokio::fs::create_dir_all(&storage_root) + .await + .expect("create ZIP test storage root"); + let state = build_zip_test_state( + pool.clone(), + database_url, + redis_url, + storage_root.to_string_lossy().to_string(), + ) + .await; + let endpoint_id = configure_test_s3(&state, &marker).await; + + let input_path = storage_root.join("single-flight-input.png"); + tokio::fs::write(&input_path, b"single-flight-payload") + .await + .expect("write ZIP input"); + let task_id = Uuid::new_v4(); + insert_zip_task(&pool, task_id, &marker, &input_path, 21).await; + let barrier = Arc::new(Barrier::new(20)); + let mut joins = Vec::new(); + for _ in 0..20 { + let state = state.clone(); + let barrier = barrier.clone(); + joins.push(tokio::spawn(async move { + barrier.wait().await; + resolve_task_zip(&state, task_id, 24).await + })); + } + let mut locators = Vec::new(); + for join in joins { + locators.push( + join.await + .expect("join concurrent ZIP request") + .expect("resolve concurrent ZIP request"), + ); + } + assert!(locators + .iter() + .all(|locator| locator.key == locators[0].key)); + assert!(locators[0] + .key + .replace('\\', "/") + .contains(&format!("/attempts/{task_id}/"))); + if endpoint_id.is_some() { + assert_eq!(locators[0].backend, "s3"); + } + let attempts: i64 = sqlx::query_scalar("SELECT zip_build_attempt FROM tasks WHERE id = $1") + .bind(task_id) + .fetch_one(&pool) + .await + .expect("query ZIP build attempts"); + assert_eq!(attempts, 1, "concurrent ZIP requests built more than once"); + + let over_budget_task = Uuid::new_v4(); + let missing_path = storage_root.join("must-not-be-downloaded.png"); + insert_zip_task( + &pool, + over_budget_task, + &format!("{marker}-over-budget"), + &missing_path, + state.config.zip_max_uncompressed_bytes as i64 + 1, + ) + .await; + let over_budget = resolve_task_zip(&state, over_budget_task, 24) + .await + .expect_err("over-budget ZIP reached the download phase"); + assert_eq!(over_budget.code, ErrorCode::FileTooLarge); + let over_budget_attempts: i64 = + sqlx::query_scalar("SELECT zip_build_attempt FROM tasks WHERE id = $1") + .bind(over_budget_task) + .fetch_one(&pool) + .await + .expect("query over-budget ZIP attempts"); + assert_eq!(over_budget_attempts, 0); + + let takeover_task = Uuid::new_v4(); + insert_zip_task( + &pool, + takeover_task, + &format!("{marker}-takeover"), + &input_path, + 21, + ) + .await; + let first_token = match claim_zip_build(&state, takeover_task) + .await + .expect("claim simulated failing ZIP builder") + { + ZipBuildClaim::Acquired { token } => token, + other => panic!("unexpected initial ZIP claim: {other:?}"), + }; + let waiter_state = state.clone(); + let waiter = + tokio::spawn(async move { resolve_task_zip(&waiter_state, takeover_task, 24).await }); + tokio::time::sleep(std::time::Duration::from_millis(250)).await; + release_zip_build(&state, takeover_task, first_token).await; + let takeover_locator = waiter + .await + .expect("join ZIP takeover waiter") + .expect("waiter safely took over ZIP build"); + let takeover_attempts: i64 = + sqlx::query_scalar("SELECT zip_build_attempt FROM tasks WHERE id = $1") + .bind(takeover_task) + .fetch_one(&pool) + .await + .expect("query takeover attempts"); + assert_eq!(takeover_attempts, 2); + + let deleted_task = Uuid::new_v4(); + insert_zip_task( + &pool, + deleted_task, + &format!("{marker}-deleted"), + &input_path, + 21, + ) + .await; + let deleted_token = match claim_zip_build(&state, deleted_task) + .await + .expect("claim deleted-task ZIP builder") + { + ZipBuildClaim::Acquired { token } => token, + other => panic!("unexpected deleted-task ZIP claim: {other:?}"), + }; + let unpublished_path = storage_root.join("unpublished.zip"); + tokio::fs::write(&unpublished_path, b"unpublished-zip") + .await + .expect("write unpublished ZIP fixture"); + let unpublished = storage::store_file( + &state, + &storage::archive_attempt_key(24, deleted_task, deleted_token), + &unpublished_path, + "application/zip", + ) + .await + .expect("store unpublished ZIP attempt"); + sqlx::query("DELETE FROM tasks WHERE id = $1") + .bind(deleted_task) + .execute(&pool) + .await + .expect("delete task before ZIP publish"); + let publish_error = + publish_zip_attempt(&state, deleted_task, deleted_token, unpublished.clone()) + .await + .expect_err("published ZIP after task deletion"); + assert_eq!(publish_error.code, ErrorCode::StorageUnavailable); + let orphan_read = storage::read_bytes( + &state, + &storage::ObjectLocator { + backend: unpublished.backend, + endpoint_id: unpublished.endpoint_id, + key: unpublished.key, + }, + ) + .await; + assert!(orphan_read.is_err(), "unpublished ZIP object was orphaned"); + + for locator in [&locators[0], &takeover_locator] { + storage::delete_object(&state, locator) + .await + .expect("delete published ZIP test object"); + } + sqlx::query("DELETE FROM tasks WHERE id = ANY($1)") + .bind(&[task_id, over_budget_task, takeover_task][..]) + .execute(&pool) + .await + .expect("delete ZIP test tasks"); + if let Some(endpoint_id) = endpoint_id { + sqlx::query("DELETE FROM storage_endpoints WHERE id = $1") + .bind(endpoint_id) + .execute(&pool) + .await + .expect("delete ZIP test S3 endpoint"); + } + tokio::fs::remove_dir_all(&storage_root) + .await + .expect("remove ZIP test storage root"); + } } diff --git a/src/api/user.rs b/src/api/user.rs index 000f2e5..b353181 100644 --- a/src/api/user.rs +++ b/src/api/user.rs @@ -1187,6 +1187,7 @@ mod tests { let state = AppState { mailer: Arc::new(Mailer::new(&config).expect("create disabled test mailer")), image_processing_semaphore: Arc::new(Semaphore::new(2)), + zip_build_semaphore: Arc::new(Semaphore::new(2)), runtime_policy_cache: crate::services::settings::RuntimePolicyCache::new(), storage_cache: crate::services::storage::StorageCache::new(), config, diff --git a/src/api/webhooks.rs b/src/api/webhooks.rs index 4e230b6..974c14c 100644 --- a/src/api/webhooks.rs +++ b/src/api/webhooks.rs @@ -1338,6 +1338,7 @@ mod tests { AppState { mailer: Arc::new(Mailer::new(&config).expect("create disabled test mailer")), image_processing_semaphore: Arc::new(Semaphore::new(2)), + zip_build_semaphore: Arc::new(Semaphore::new(2)), runtime_policy_cache: crate::services::settings::RuntimePolicyCache::new(), storage_cache: crate::services::storage::StorageCache::new(), config, diff --git a/src/config.rs b/src/config.rs index 72a15ca..2998f2b 100644 --- a/src/config.rs +++ b/src/config.rs @@ -17,6 +17,9 @@ pub struct Config { pub worker_task_concurrency: u32, pub worker_concurrency: u32, pub image_processing_concurrency: u32, + pub zip_build_concurrency: u32, + pub zip_max_entries: u32, + pub zip_max_uncompressed_bytes: u64, pub jwt_secret: String, pub jwt_expiry_hours: i64, @@ -81,6 +84,15 @@ impl Config { .map(|v| v.get() as u32) .unwrap_or(4) }); + let zip_build_concurrency = env_u32("ZIP_BUILD_CONCURRENCY") + .filter(|value| *value > 0) + .unwrap_or(2); + let zip_max_entries = env_u32("ZIP_MAX_ENTRIES") + .filter(|value| *value > 0) + .unwrap_or(200); + let zip_max_uncompressed_bytes = env_u64("ZIP_MAX_UNCOMPRESSED_BYTES") + .filter(|value| *value > 0) + .unwrap_or(2 * 1024 * 1024 * 1024); let jwt_secret = env_string("JWT_SECRET") .ok_or_else(|| AppError::new(ErrorCode::InvalidRequest, "缺少环境变量 JWT_SECRET"))?; @@ -140,6 +152,9 @@ impl Config { worker_task_concurrency, worker_concurrency, image_processing_concurrency, + zip_build_concurrency, + zip_max_entries, + zip_max_uncompressed_bytes, jwt_secret, jwt_expiry_hours, api_key_pepper, diff --git a/src/main.rs b/src/main.rs index 822a219..4fb0ca4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -36,6 +36,9 @@ async fn main() -> Result<(), AppError> { let image_processing_semaphore = std::sync::Arc::new(tokio::sync::Semaphore::new( config.image_processing_concurrency as usize, )); + let zip_build_semaphore = std::sync::Arc::new(tokio::sync::Semaphore::new( + config.zip_build_concurrency as usize, + )); let state = AppState { config, @@ -43,6 +46,7 @@ async fn main() -> Result<(), AppError> { redis, mailer: std::sync::Arc::new(mailer), image_processing_semaphore, + zip_build_semaphore, runtime_policy_cache: crate::services::settings::RuntimePolicyCache::new(), storage_cache: crate::services::storage::StorageCache::new(), }; diff --git a/src/services/quota.rs b/src/services/quota.rs index 06f35f2..76a09cc 100644 --- a/src/services/quota.rs +++ b/src/services/quota.rs @@ -605,6 +605,7 @@ mod tests { AppState { mailer: Arc::new(Mailer::new(&config).expect("create disabled quota test mailer")), image_processing_semaphore: Arc::new(Semaphore::new(2)), + zip_build_semaphore: Arc::new(Semaphore::new(2)), runtime_policy_cache: crate::services::settings::RuntimePolicyCache::new(), storage_cache: crate::services::storage::StorageCache::new(), config, diff --git a/src/services/storage.rs b/src/services/storage.rs index fa5a667..dd90687 100644 --- a/src/services/storage.rs +++ b/src/services/storage.rs @@ -311,10 +311,10 @@ pub fn result_attempt_key( ) } -pub fn archive_key(retention_hours: i64, task_id: Uuid) -> String { +pub fn archive_attempt_key(retention_hours: i64, task_id: Uuid, token: Uuid) -> String { let now = Utc::now(); format!( - "archives/{}/{:04}/{:02}/{task_id}.zip", + "archives/{}/{:04}/{:02}/attempts/{task_id}/{token}.zip", retention_prefix(retention_hours), now.year(), now.month() @@ -983,7 +983,7 @@ mod tests { let key = result_key(168, task_id, file_id, "webp"); assert!(key.starts_with("results/7d/")); assert!(key.ends_with("/00000000-0000-0000-0000-000000000001.webp")); - assert!(archive_key(360, task_id).starts_with("archives/15d/")); + assert!(archive_attempt_key(360, task_id, Uuid::new_v4()).starts_with("archives/15d/")); let attempt_key = result_attempt_key(24, task_id, file_id, 2, 3, "avif"); assert!(attempt_key.contains("-t2-f3.avif")); } diff --git a/src/state.rs b/src/state.rs index 33ba26c..d0e7fa6 100644 --- a/src/state.rs +++ b/src/state.rs @@ -10,6 +10,7 @@ pub struct AppState { pub redis: redis::aio::ConnectionManager, pub mailer: std::sync::Arc, pub image_processing_semaphore: std::sync::Arc, + pub zip_build_semaphore: std::sync::Arc, pub runtime_policy_cache: RuntimePolicyCache, pub storage_cache: StorageCache, } diff --git a/src/worker/mod.rs b/src/worker/mod.rs index d3c09e1..b292a44 100644 --- a/src/worker/mod.rs +++ b/src/worker/mod.rs @@ -2026,6 +2026,7 @@ mod tests { let state = AppState { mailer: Arc::new(Mailer::new(&config).expect("create disabled test mailer")), image_processing_semaphore: Arc::new(Semaphore::new(2)), + zip_build_semaphore: Arc::new(Semaphore::new(2)), runtime_policy_cache: crate::services::settings::RuntimePolicyCache::new(), storage_cache: storage::StorageCache::new(), config,