237899745/ystp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: fence concurrent ZIP archive builds

Browse Source
This commit is contained in:
237899745
2026-07-26 05:47:19 +08:00
parent 037e83e92f
commit 923ba495c4
16 changed files with 777 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
.env.example
View File

@@ -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

5
docker/.env.production.example
View File

@@ -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

3
docker/docker-compose.prod.yml
View File

@@ -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}

2
docs/architecture.md
View File

@@ -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 缓存用户会话

2
docs/s3-storage-plan.md
View File

@@ -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/<retention>/.../attempts/` 前缀，发布失败会立即删除，进程崩溃遗留项仍由同一前缀生命周期兜底。S3 生命周期多保留 2 天，只负责处理数据库故障、进程崩溃或上传后未能落库的孤儿对象，不能作为精确会员权限判断。未完成的分片上传 1 天后由生命周期中止。
## 5. 119 首期容量

12
migrations/021_zip_build_leases.sql Normal file
View File

@@ -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;

1
src/api/billing.rs
View File

@@ -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,

787
src/api/downloads.rs
View File

@@ -241,8 +241,27 @@ struct TaskZipFileRow {
storage_key: Option<String>,
original_name: String,
output_format: String,
compressed_size: Option<i64>,
}
#[derive(Debug, FromRow)]
struct ZipBuildStateRow {
zip_storage_backend: Option<String>,
zip_storage_endpoint_id: Option<Uuid>,
zip_storage_key: Option<String>,
expires_at: DateTime<Utc>,
}
#[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<AppState>,
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<storage::ObjectLocator, AppError> {
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<storage::StoredObject, AppError> = async {
let mut used_names: HashMap<String, usize> = 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<u64, AppError> {
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<ZipBuildClaim, AppError> {
let token = Uuid::new_v4();
let acquired: Option<i64> = 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<storage::ObjectLocator, AppError> {
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<storage::StoredObject, AppError> {
tokio::fs::create_dir_all(temp_dir).await.map_err(|err| {
AppError::new(ErrorCode::StorageUnavailable, "创建 ZIP 临时目录失败").with_source(err)
})?;
let mut used_names: HashMap<String, usize> = 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<storage::ObjectLocator, AppError> {
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<Option<storage::ObjectLocator>, AppError> {
let row: Option<(Option<String>, Option<Uuid>, Option<String>)> = 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<Uuid> {
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");
}
}

1
src/api/user.rs
View File

@@ -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,

1
src/api/webhooks.rs
View File

@@ -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,

15
src/config.rs
View File

@@ -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,

4
src/main.rs
View File

@@ -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(),
};

1
src/services/quota.rs
View File

@@ -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,

6
src/services/storage.rs
View File

@@ -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"));
}

1
src/state.rs
View File

@@ -10,6 +10,7 @@ pub struct AppState {
pub redis: redis::aio::ConnectionManager,
pub mailer: std::sync::Arc<Mailer>,
pub image_processing_semaphore: std::sync::Arc<tokio::sync::Semaphore>,
pub zip_build_semaphore: std::sync::Arc<tokio::sync::Semaphore>,
pub runtime_policy_cache: RuntimePolicyCache,
pub storage_cache: StorageCache,
}

1
src/worker/mod.rs
View File

@@ -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,