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

fix: finalize failed batch enqueue atomically

Browse Source
This commit is contained in:
237899745
2026-07-26 06:46:29 +08:00
parent 380e89b058
commit 4d0e8aa70a
4 changed files with 361 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files

396
src/api/tasks.rs
View File

@@ -21,6 +21,7 @@ use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use sqlx::FromRow;
use std::net::SocketAddr;
use std::path::PathBuf;
use tokio::io::AsyncWriteExt;
use uuid::Uuid;
@@ -115,16 +116,13 @@ async fn create_batch_task(
let (files, opts, request_hash) = match parsed {
Ok(parsed) => parsed,
Err(err) => {
let base_dir = format!("{}/orig/{task_id}", state.config.storage_path);
let _ = tokio::fs::remove_dir_all(base_dir).await;
cleanup_task_input_dir(&state, task_id).await;
return Err(err);
}
};
if files.is_empty() {
cleanup_file_paths(&files).await;
let base_dir = format!("{}/orig/{task_id}", state.config.storage_path);
let _ = tokio::fs::remove_dir_all(base_dir).await;
cleanup_task_input_dir(&state, task_id).await;
return Err(AppError::new(ErrorCode::InvalidRequest, "缺少 files[]"));
}
@@ -141,13 +139,13 @@ async fn create_batch_task(
let begin_result = match begin_result {
Ok(result) => result,
Err(err) => {
cleanup_file_paths(&files).await;
cleanup_task_input_dir(&state, task_id).await;
return Err(err);
}
};
match begin_result {
idempotency::BeginResult::Replay { response_body, .. } => {
cleanup_file_paths(&files).await;
cleanup_task_input_dir(&state, task_id).await;
let resp: BatchCreateResponse =
serde_json::from_value(response_body).map_err(|err| {
AppError::new(ErrorCode::Internal, "幂等结果解析失败").with_source(err)
@@ -161,7 +159,7 @@ async fn create_batch_task(
));
}
idempotency::BeginResult::InProgress => {
cleanup_file_paths(&files).await;
cleanup_task_input_dir(&state, task_id).await;
if let Some((_status, body)) =
idempotency::wait_for_replay(&state, scope, idem_key, &request_hash, 10_000)
.await?
@@ -191,6 +189,9 @@ async fn create_batch_task(
let mut anonymous_reserved_units = 0u32;
let mut anonymous_quota_date = None;
let mut task_persisted = false;
let mut enqueue_failure_finalized = false;
let mut cleanup_inputs_on_error = true;
let create_result: Result<BatchCreateResponse, AppError> = (async {
match &admission.task_owner {
TaskOwner::Anonymous { session_id } => {
@@ -301,14 +302,21 @@ async fn create_batch_task(
tx.commit()
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "提交事务失败").with_source(err))?;
task_persisted = true;
if let Err(err) = enqueue_task(&state, task_id).await {
let _ =
sqlx::query("UPDATE tasks SET status = 'failed', error_message = $2 WHERE id = $1")
.bind(task_id)
.bind("队列提交失败")
.execute(&state.db)
.await;
match finalize_enqueue_failure(&state, task_id, "队列提交失败").await {
Ok(true) => enqueue_failure_finalized = true,
Ok(false) => {
// XADD may have succeeded even if the client saw an error. A worker
// that already claimed the task owns both the input and settlement.
cleanup_inputs_on_error = false;
}
Err(finalize_err) => {
cleanup_inputs_on_error = false;
tracing::error!(task_id = %task_id, error = %finalize_err, "failed to finalize task after queue submission error");
}
}
return Err(err);
}
@@ -347,28 +355,32 @@ async fn create_batch_task(
Err(err) => {
if anonymous_reserved_units > 0 {
if let context::Principal::Anonymous { session_id } = &principal {
let settlement =
quota::settle_anonymous_task_reservation(&state, task_id).await;
match settlement {
Ok(Some(_)) => {}
Ok(None) => {
if let Some(date) = anonymous_quota_date {
if let Err(refund_err) = quota::refund_anonymous_reservation_once(
&state,
task_id,
session_id,
ip,
date,
anonymous_reserved_units,
)
.await
{
tracing::warn!(task_id = %task_id, error = %refund_err, "failed to refund anonymous batch admission");
}
let should_refund_directly = if enqueue_failure_finalized {
match quota::settle_anonymous_task_reservation(&state, task_id).await {
Ok(Some(_)) => false,
Ok(None) => true,
Err(settle_err) => {
tracing::warn!(task_id = %task_id, error = %settle_err, "failed to settle anonymous batch admission");
false
}
}
Err(settle_err) => {
tracing::warn!(task_id = %task_id, error = %settle_err, "failed to settle anonymous batch admission");
} else {
!task_persisted
};
if should_refund_directly {
if let Some(date) = anonymous_quota_date {
if let Err(refund_err) = quota::refund_anonymous_reservation_once(
&state,
task_id,
session_id,
ip,
date,
anonymous_reserved_units,
)
.await
{
tracing::warn!(task_id = %task_id, error = %refund_err, "failed to refund anonymous batch admission");
}
}
}
}
@@ -378,7 +390,9 @@ async fn create_batch_task(
let _ = idempotency::abort(&state, scope, idem_key, &request_hash).await;
}
}
cleanup_file_paths(&files).await;
if cleanup_inputs_on_error {
cleanup_task_input_dir(&state, task_id).await;
}
Err(err)
}
}
@@ -489,6 +503,88 @@ async fn enqueue_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
Ok(())
}
async fn finalize_enqueue_failure(
state: &AppState,
task_id: Uuid,
error_message: &str,
) -> Result<bool, AppError> {
let mut tx = state.db.begin().await.map_err(|err| {
AppError::new(ErrorCode::Internal, "开启队列失败收口事务失败").with_source(err)
})?;
let task: Option<(String, i32)> =
sqlx::query_as("SELECT status::text, total_files FROM tasks WHERE id = $1 FOR UPDATE")
.bind(task_id)
.fetch_optional(&mut *tx)
.await
.map_err(|err| {
AppError::new(ErrorCode::Internal, "锁定队列失败任务失败").with_source(err)
})?;
let Some((status, total_files)) = task else {
tx.rollback().await.ok();
return Ok(false);
};
if status != "pending" {
tx.rollback().await.ok();
return Ok(false);
}
let files = sqlx::query(
r#"
UPDATE task_files
SET status = 'failed',
error_message = $2,
completed_at = NOW(),
input_path = NULL,
storage_path = NULL,
lease_owner = NULL,
lease_until = NULL
WHERE task_id = $1
AND status = 'pending'
"#,
)
.bind(task_id)
.bind(error_message)
.execute(&mut *tx)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "终结未入队文件失败").with_source(err))?;
if files.rows_affected() != u64::try_from(total_files.max(0)).unwrap_or(0) {
tx.rollback().await.ok();
return Err(AppError::new(
ErrorCode::Internal,
"未入队任务的文件状态不一致",
));
}
let task = sqlx::query(
r#"
UPDATE tasks
SET status = 'failed',
error_message = $2,
completed_at = NOW(),
completed_files = 0,
failed_files = total_files,
lease_owner = NULL,
lease_until = NULL
WHERE id = $1
AND status = 'pending'
"#,
)
.bind(task_id)
.bind(error_message)
.execute(&mut *tx)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "终结未入队任务失败").with_source(err))?;
if task.rows_affected() != 1 {
tx.rollback().await.ok();
return Ok(false);
}
tx.commit().await.map_err(|err| {
AppError::new(ErrorCode::Internal, "提交队列失败收口事务失败").with_source(err)
})?;
Ok(true)
}
async fn parse_batch_request(
state: &AppState,
task_id: Uuid,
@@ -1173,10 +1269,232 @@ async fn cleanup_file_paths(files: &[BatchFileInput]) {
.first()
.and_then(|file| std::path::Path::new(&file.storage_path).parent())
.map(std::path::Path::to_path_buf);
for f in files {
let _ = tokio::fs::remove_file(&f.storage_path).await;
}
if let Some(parent) = parent {
let _ = tokio::fs::remove_dir(parent).await;
let _ = tokio::fs::remove_dir_all(parent).await;
}
}
fn task_input_dir(state: &AppState, task_id: Uuid) -> PathBuf {
PathBuf::from(&state.config.storage_path)
.join("orig")
.join(task_id.to_string())
}
async fn cleanup_task_input_dir(state: &AppState, task_id: Uuid) {
let path = task_input_dir(state, task_id);
match tokio::fs::remove_dir_all(&path).await {
Ok(()) => {}
Err(err) if err.kind() == std::io::ErrorKind::NotFound => {}
Err(err) => {
tracing::warn!(task_id = %task_id, path = %path.display(), error = %err, "failed to clean task input directory");
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::Config;
use crate::services::mail::Mailer;
use crate::services::settings::RuntimePolicyCache;
use crate::services::storage::StorageCache;
use crate::worker::TaskProcessOutcome;
use sqlx::postgres::PgPoolOptions;
use std::sync::Arc;
use tokio::sync::Semaphore;
async fn build_batch_test_state(
pool: sqlx::PgPool,
database_url: String,
redis_url: String,
storage_path: String,
) -> AppState {
let mut config = Config::from_env().expect("load batch test config");
config.database_url = database_url;
config.redis_url = redis_url;
config.storage_path = storage_path;
config.mail_enabled = false;
config.mail_log_links_when_disabled = false;
let redis = redis::Client::open(config.redis_url.clone())
.expect("create batch test Redis client")
.get_connection_manager()
.await
.expect("connect batch test Redis");
AppState {
mailer: Arc::new(Mailer::new(&config).expect("create disabled batch test mailer")),
image_processing_semaphore: Arc::new(Semaphore::new(2)),
zip_build_semaphore: Arc::new(Semaphore::new(1)),
runtime_policy_cache: RuntimePolicyCache::new(),
storage_cache: StorageCache::new(),
config,
db: pool,
redis,
}
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[ignore = "requires isolated IMAGEFORGE_TEST_DATABASE_URL and IMAGEFORGE_TEST_REDIS_URL"]
async fn enqueue_failure_finalizes_files_and_removes_exact_input_directory() {
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(8)
.connect(&database_url)
.await
.expect("connect batch 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-batch-test-{marker}"));
let state = build_batch_test_state(
pool.clone(),
database_url,
redis_url,
storage_root.to_string_lossy().to_string(),
)
.await;
let task_id = Uuid::new_v4();
let input_dir = task_input_dir(&state, task_id);
tokio::fs::create_dir_all(&input_dir)
.await
.expect("create batch input directory");
let first_path = input_dir.join("first.png");
let second_path = input_dir.join("second.png");
tokio::fs::write(&first_path, b"first")
.await
.expect("write first input");
tokio::fs::write(&second_path, b"second")
.await
.expect("write second input");
tokio::fs::write(input_dir.join("interrupted.upload"), b"partial")
.await
.expect("write interrupted upload fixture");
sqlx::query(
r#"
INSERT INTO tasks (
id, session_id, client_ip, status, total_files,
total_original_size, expires_at, retention_hours
) VALUES ($1, $2, '127.0.0.1'::inet, 'pending', 2, 11, NOW() + INTERVAL '1 day', 24)
"#,
)
.bind(task_id)
.bind(format!("batch-session-{marker}"))
.execute(&pool)
.await
.expect("insert pending batch task");
for (name, path, size) in [
("first.png", &first_path, 5_i64),
("second.png", &second_path, 6_i64),
] {
sqlx::query(
r#"
INSERT INTO task_files (
id, task_id, original_name, original_format, output_format,
original_size, input_path, status
) VALUES ($1, $2, $3, 'png', 'png', $4, $5, 'pending')
"#,
)
.bind(Uuid::new_v4())
.bind(task_id)
.bind(name)
.bind(size)
.bind(path.to_string_lossy().to_string())
.execute(&pool)
.await
.expect("insert pending batch file");
}
let mut redis = state.redis.clone();
let _: i64 = redis::cmd("DEL")
.arg("stream:compress_jobs")
.query_async(&mut redis)
.await
.expect("clear compression stream");
let _: () = redis::cmd("SET")
.arg("stream:compress_jobs")
.arg("wrong-type-fixture")
.query_async(&mut redis)
.await
.expect("install WRONGTYPE fixture");
let enqueue_error = enqueue_task(&state, task_id)
.await
.expect_err("XADD unexpectedly accepted a string key");
assert_eq!(enqueue_error.code, ErrorCode::Internal);
assert!(
finalize_enqueue_failure(&state, task_id, "队列提交失败")
.await
.expect("finalize enqueue failure"),
"pending task was not finalized"
);
cleanup_task_input_dir(&state, task_id).await;
let task: (String, bool, i32, i32) = sqlx::query_as(
r#"
SELECT status::text, completed_at IS NOT NULL, completed_files, failed_files
FROM tasks WHERE id = $1
"#,
)
.bind(task_id)
.fetch_one(&pool)
.await
.expect("query finalized task");
assert_eq!(task, ("failed".to_string(), true, 0, 2));
let files: Vec<(String, bool, bool, Option<String>)> = sqlx::query_as(
r#"
SELECT status::text, completed_at IS NOT NULL, input_path IS NULL, error_message
FROM task_files WHERE task_id = $1 ORDER BY original_name
"#,
)
.bind(task_id)
.fetch_all(&pool)
.await
.expect("query finalized task files");
assert_eq!(files.len(), 2);
assert!(files.iter().all(|row| {
row.0 == "failed" && row.1 && row.2 && row.3.as_deref() == Some("队列提交失败")
}));
assert!(
!tokio::fs::try_exists(&input_dir)
.await
.expect("check input directory"),
"task input directory or partial upload survived cleanup"
);
let before = task.clone();
let outcome = crate::worker::process_task(&state, task_id, Uuid::new_v4())
.await
.expect("reprocess terminal task");
assert_eq!(outcome, TaskProcessOutcome::Done);
let after: (String, bool, i32, i32) = sqlx::query_as(
r#"
SELECT status::text, completed_at IS NOT NULL, completed_files, failed_files
FROM tasks WHERE id = $1
"#,
)
.bind(task_id)
.fetch_one(&pool)
.await
.expect("query task after reprocess");
assert_eq!(after, before, "reprocessing changed a terminal task");
let _: i64 = redis::cmd("DEL")
.arg("stream:compress_jobs")
.query_async(&mut redis)
.await
.expect("remove WRONGTYPE fixture");
sqlx::query("DELETE FROM tasks WHERE id = $1")
.bind(task_id)
.execute(&pool)
.await
.expect("delete batch test task");
let _ = tokio::fs::remove_dir_all(&storage_root).await;
}
}

4
src/worker/mod.rs
View File

@@ -703,12 +703,12 @@ struct TaskContext {
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum TaskProcessOutcome {
pub(crate) enum TaskProcessOutcome {
Done,
LeaseBusy,
}
async fn process_task(
pub(crate) async fn process_task(
state: &AppState,
task_id: Uuid,
worker_id: Uuid,