From 4d0e8aa70af914ed9b407e6518f71855dda29fd4 Mon Sep 17 00:00:00 2001 From: 237899745 <237899745@users.noreply.git.workyai.cn> Date: Sun, 26 Jul 2026 06:46:29 +0800 Subject: [PATCH] fix: finalize failed batch enqueue atomically --- .gitea/workflows/ci.yml | 2 +- scripts/run_external_state_tests.sh | 2 +- src/api/tasks.rs | 396 +++++++++++++++++++++++++--- src/worker/mod.rs | 4 +- 4 files changed, 361 insertions(+), 43 deletions(-) diff --git a/.gitea/workflows/ci.yml b/.gitea/workflows/ci.yml index 8edd3e1..17db367 100644 --- a/.gitea/workflows/ci.yml +++ b/.gitea/workflows/ci.yml @@ -33,7 +33,7 @@ jobs: IMAGEFORGE_TEST_REDIS_URL: redis://redis:6379/ JWT_SECRET: imageforge-ci-jwt-secret API_KEY_PEPPER: imageforge-ci-api-key-pepper - EXPECTED_EXTERNAL_TESTS: '9' + EXPECTED_EXTERNAL_TESTS: '10' steps: - name: Checkout uses: actions/checkout@v4 diff --git a/scripts/run_external_state_tests.sh b/scripts/run_external_state_tests.sh index dd96322..f08dc87 100644 --- a/scripts/run_external_state_tests.sh +++ b/scripts/run_external_state_tests.sh @@ -1,7 +1,7 @@ #!/usr/bin/env bash set -Eeuo pipefail -EXPECTED_EXTERNAL_TESTS="${EXPECTED_EXTERNAL_TESTS:-9}" +EXPECTED_EXTERNAL_TESTS="${EXPECTED_EXTERNAL_TESTS:-10}" WORK_DIR="$(mktemp -d)" POSTGRES_CONTAINER="" REDIS_CONTAINER="" diff --git a/src/api/tasks.rs b/src/api/tasks.rs index 62e67df..3c535a9 100644 --- a/src/api/tasks.rs +++ b/src/api/tasks.rs @@ -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 = (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 { + 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)> = 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; } } diff --git a/src/worker/mod.rs b/src/worker/mod.rs index 367e620..634904a 100644 --- a/src/worker/mod.rs +++ b/src/worker/mod.rs @@ -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,