From 408e09cda8bc8ced6bb2e50fdee634cade84ac11 Mon Sep 17 00:00:00 2001 From: 237899745 <237899745@users.noreply.git.workyai.cn> Date: Sun, 26 Jul 2026 10:54:13 +0800 Subject: [PATCH] fix(storage): make queue and object ownership durable --- migrations/024_task_queue_outbox.sql | 28 + migrations/025_storage_object_lifecycle.sql | 101 ++ .../026_idempotency_operation_leases.sql | 11 + src/api/compress.rs | 444 +++++-- src/api/downloads.rs | 116 +- src/api/metrics.rs | 49 +- src/api/mod.rs | 7 + src/api/tasks.rs | 628 +++++----- src/services/idempotency.rs | 368 +++++- src/services/mod.rs | 2 + src/services/object_lifecycle.rs | 1039 +++++++++++++++++ src/services/storage.rs | 57 +- src/services/task_queue.rs | 367 ++++++ src/worker/mod.rs | 246 ++-- 14 files changed, 2855 insertions(+), 608 deletions(-) create mode 100644 migrations/024_task_queue_outbox.sql create mode 100644 migrations/025_storage_object_lifecycle.sql create mode 100644 migrations/026_idempotency_operation_leases.sql create mode 100644 src/services/object_lifecycle.rs create mode 100644 src/services/task_queue.rs diff --git a/migrations/024_task_queue_outbox.sql b/migrations/024_task_queue_outbox.sql new file mode 100644 index 0000000..c6f6db3 --- /dev/null +++ b/migrations/024_task_queue_outbox.sql @@ -0,0 +1,28 @@ +CREATE TABLE task_queue_outbox ( + task_id UUID PRIMARY KEY REFERENCES tasks(id) ON DELETE CASCADE, + status VARCHAR(20) NOT NULL DEFAULT 'pending', + attempts INTEGER NOT NULL DEFAULT 0, + next_attempt_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + lease_owner UUID, + lease_until TIMESTAMPTZ, + last_error TEXT, + delivered_at TIMESTAMPTZ, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + CONSTRAINT task_queue_outbox_status_check + CHECK (status IN ('pending', 'delivering', 'delivered', 'dead')), + CONSTRAINT task_queue_outbox_attempts_check + CHECK (attempts >= 0), + CONSTRAINT task_queue_outbox_lease_pair_check + CHECK ((lease_owner IS NULL) = (lease_until IS NULL)) +); + +CREATE INDEX task_queue_outbox_ready + ON task_queue_outbox(next_attempt_at, created_at) + WHERE status IN ('pending', 'delivering'); + +INSERT INTO task_queue_outbox (task_id) +SELECT id +FROM tasks +WHERE status = 'pending' +ON CONFLICT (task_id) DO NOTHING; diff --git a/migrations/025_storage_object_lifecycle.sql b/migrations/025_storage_object_lifecycle.sql new file mode 100644 index 0000000..f281552 --- /dev/null +++ b/migrations/025_storage_object_lifecycle.sql @@ -0,0 +1,101 @@ +ALTER TABLE tasks + ADD COLUMN deletion_started_at TIMESTAMPTZ, + ADD COLUMN deletion_reason VARCHAR(32); + +CREATE INDEX tasks_deletion_pending + ON tasks(deletion_started_at) + WHERE deletion_started_at IS NOT NULL; + +CREATE TABLE storage_objects ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + task_id UUID NOT NULL, + task_file_id UUID, + object_kind VARCHAR(32) NOT NULL, + state VARCHAR(20) NOT NULL DEFAULT 'staging', + backend VARCHAR(16) NOT NULL, + storage_endpoint_id UUID REFERENCES storage_endpoints(id) ON DELETE RESTRICT, + object_key TEXT NOT NULL, + storage_etag TEXT, + size_bytes BIGINT, + lease_owner UUID, + lease_until TIMESTAMPTZ, + delete_attempts INTEGER NOT NULL DEFAULT 0, + next_attempt_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + last_error TEXT, + published_at TIMESTAMPTZ, + deleted_at TIMESTAMPTZ, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + CONSTRAINT storage_objects_kind_check + CHECK (object_kind IN ('result', 'zip_attempt', 'input', 'input_dir', 'legacy_zip')), + CONSTRAINT storage_objects_state_check + CHECK (state IN ('staging', 'published', 'delete_pending', 'deleted')), + CONSTRAINT storage_objects_backend_check + CHECK (backend IN ('s3', 'local', 'local_dir')), + CONSTRAINT storage_objects_attempts_check + CHECK (delete_attempts >= 0), + CONSTRAINT storage_objects_size_check + CHECK (size_bytes IS NULL OR size_bytes >= 0), + CONSTRAINT storage_objects_endpoint_check + CHECK ( + (backend = 's3' AND storage_endpoint_id IS NOT NULL) + OR (backend IN ('local', 'local_dir') AND storage_endpoint_id IS NULL) + ), + CONSTRAINT storage_objects_lease_pair_check + CHECK ((lease_owner IS NULL) = (lease_until IS NULL)) +); + +CREATE UNIQUE INDEX storage_objects_locator_unique + ON storage_objects( + backend, + COALESCE(storage_endpoint_id, '00000000-0000-0000-0000-000000000000'::uuid), + object_key + ); + +CREATE INDEX storage_objects_cleanup_ready + ON storage_objects(next_attempt_at, created_at) + WHERE state IN ('staging', 'delete_pending'); + +CREATE INDEX storage_objects_task + ON storage_objects(task_id, state); + +INSERT INTO storage_objects ( + task_id, task_file_id, object_kind, state, + backend, storage_endpoint_id, object_key, storage_etag, + size_bytes, published_at +) +SELECT + file.task_id, + file.id, + 'result', + 'published', + file.storage_backend, + file.storage_endpoint_id, + COALESCE(file.storage_key, file.storage_path), + file.storage_etag, + file.compressed_size, + COALESCE(file.completed_at, file.created_at) +FROM task_files AS file +WHERE file.status = 'completed' + AND COALESCE(file.storage_key, file.storage_path) IS NOT NULL +ON CONFLICT DO NOTHING; + +INSERT INTO storage_objects ( + task_id, object_kind, state, + backend, storage_endpoint_id, object_key, storage_etag, + size_bytes, published_at +) +SELECT + task.id, + 'zip_attempt', + 'published', + task.zip_storage_backend, + task.zip_storage_endpoint_id, + task.zip_storage_key, + task.zip_storage_etag, + task.zip_size, + COALESCE(task.completed_at, task.created_at) +FROM tasks AS task +WHERE task.zip_storage_backend IS NOT NULL + AND task.zip_storage_key IS NOT NULL +ON CONFLICT DO NOTHING; diff --git a/migrations/026_idempotency_operation_leases.sql b/migrations/026_idempotency_operation_leases.sql new file mode 100644 index 0000000..1fa21ae --- /dev/null +++ b/migrations/026_idempotency_operation_leases.sql @@ -0,0 +1,11 @@ +ALTER TABLE idempotency_keys + ADD COLUMN lease_owner UUID, + ADD COLUMN lease_until TIMESTAMPTZ; + +ALTER TABLE idempotency_keys + ADD CONSTRAINT idempotency_keys_lease_pair_check + CHECK ((lease_owner IS NULL) = (lease_until IS NULL)); + +CREATE INDEX idempotency_keys_stale_operations + ON idempotency_keys(lease_until) + WHERE response_status = 0; diff --git a/src/api/compress.rs b/src/api/compress.rs index 66be550..c542976 100644 --- a/src/api/compress.rs +++ b/src/api/compress.rs @@ -8,6 +8,7 @@ use crate::services::compress; use crate::services::compress::{CompressionLevel, ImageFmt}; use crate::services::filename; use crate::services::idempotency; +use crate::services::object_lifecycle; use crate::services::quota; use crate::services::storage; use crate::state::AppState; @@ -21,6 +22,7 @@ use chrono::{DateTime, Duration, Utc}; use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; use sqlx::FromRow; +use std::future::Future; use std::net::{IpAddr, SocketAddr}; use uuid::Uuid; @@ -30,6 +32,14 @@ pub fn router() -> Router { .route("/compress/direct", post(compress_direct)) } +fn spawn_detached_operation(future: F) -> tokio::task::JoinHandle +where + F: Future + Send + 'static, + T: Send + 'static, +{ + tokio::spawn(future) +} + #[derive(Debug, Serialize, Deserialize)] struct BillingView { units_charged: i32, @@ -214,6 +224,7 @@ async fn compress_json( let quota_ctx = admission.quota_ctx; let mut idem_acquired = false; + let mut idem_owner = None; if let (Some(scope), Some(idem_key), Some(request_hash)) = ( idempotency_scope, idempotency_key.as_deref(), @@ -262,15 +273,35 @@ async fn compress_json( "请求正在处理中，请稍后重试", )); } - idempotency::BeginResult::Acquired => { + idempotency::BeginResult::Acquired { owner } => { idem_acquired = true; + idem_owner = Some(owner); } } } let task_id = Uuid::new_v4(); - let mut anonymous_reservation_date = None; - let op: Result = (async { + let operation_state = state.clone(); + let operation_principal = principal.clone(); + let operation_quota_ctx = quota_ctx.clone(); + let operation_idempotency_key = idempotency_key.clone(); + let operation_request_hash = request_hash.clone(); + let operation_idem_owner = idem_owner; + let operation = spawn_detached_operation(async move { + let state = operation_state; + let principal = operation_principal; + let quota_ctx = operation_quota_ctx; + let idempotency_key = operation_idempotency_key; + let request_hash = operation_request_hash; + let _idempotency_heartbeat = start_idempotency_heartbeat( + &state, + idempotency_scope, + idempotency_key.as_deref(), + request_hash.as_deref(), + operation_idem_owner, + ); + let mut anonymous_reservation_date = None; + let op: Result = (async { match "a_ctx { QuotaContext::User(billing) => ensure_quota_available(&state, billing, 1).await?, QuotaContext::ApiKey(billing, _) => ensure_quota_available(&state, billing, 1).await?, @@ -318,11 +349,41 @@ async fn compress_json( let retention_hours = retention.num_hours(); let object_key = storage::result_key(retention_hours, task_id, file_id, format_out.extension()); - let stored = - storage::store_bytes(&state, &object_key, compressed, format_out.content_type()) - .await?; + let tracked = object_lifecycle::store_tracked_bytes( + &state, + task_id, + Some(file_id), + "result", + &object_key, + compressed, + format_out.content_type(), + ) + .await?; let expires_at = Utc::now() + retention; + let response = CompressResponse { + task_id, + file_id, + format_in: format_in.as_str().to_string(), + format_out: format_out.as_str().to_string(), + original_size, + compressed_size, + saved_bytes, + saved_percent, + download_url: format!("/downloads/{file_id}"), + expires_at, + billing: BillingView { + units_charged: metered_units(charge_units), + }, + }; + let idem_completion = build_idempotency_completion( + idem_acquired, + operation_idem_owner, + idempotency_scope, + idempotency_key.as_deref(), + request_hash.as_deref(), + &response, + )?; if let Err(err) = record_task_and_metering( &state, @@ -330,7 +391,7 @@ async fn compress_json( ip, task_id, file_id, - &stored, + &tracked, &req.file_name, req.max_width, req.max_height, @@ -345,19 +406,31 @@ async fn compress_json( retention_hours, "a_ctx, charge_units, + idem_completion.as_ref(), ) .await { - let _ = storage::delete_object( - &state, - &storage::ObjectLocator { - backend: stored.backend.clone(), - endpoint_id: stored.endpoint_id, - key: stored.key.clone(), - }, - ) - .await; - return Err(err); + match sync_result_was_committed(&state, task_id, file_id, &tracked).await { + Ok(true) => { + tracing::warn!(task_id = %task_id, file_id = %file_id, error = %err, "sync result commit response was lost; recovered committed publication"); + } + Ok(false) => { + if let Err(cleanup_err) = object_lifecycle::schedule_tracked_delete( + &state, + &tracked, + Some(&err), + ) + .await + { + tracing::error!(task_id = %task_id, storage_object_id = %tracked.lifecycle_id, error = %cleanup_err, "failed to persist rejected sync result cleanup"); + } + return Err(err); + } + Err(probe_err) => { + tracing::error!(task_id = %task_id, storage_object_id = %tracked.lifecycle_id, error = %probe_err, original_error = %err, "sync result commit state is unknown; staging lease will reconcile object"); + return Err(err); + } + } } if anonymous_reservation_date.is_some() { @@ -371,51 +444,23 @@ async fn compress_json( } } - Ok(CompressResponse { - task_id, - file_id, - format_in: format_in.as_str().to_string(), - format_out: format_out.as_str().to_string(), - original_size, - compressed_size, - saved_bytes, - saved_percent, - download_url: format!("/downloads/{file_id}"), - expires_at, - billing: BillingView { - units_charged: metered_units(charge_units), - }, + Ok(response) }) - }) - .await; + .await; + (op, anonymous_reservation_date) + }); + let (op, anonymous_reservation_date) = operation.await.map_err(|err| { + AppError::new(ErrorCode::Internal, "同步压缩后台任务异常退出").with_source(err) + })?; match op { - Ok(resp) => { - if let (Some(scope), Some(idem_key), Some(request_hash)) = ( - idempotency_scope, - idempotency_key.as_deref(), - request_hash.as_deref(), - ) { - if idem_acquired { - let _ = idempotency::complete( - &state, - scope, - idem_key, - request_hash, - 200, - serde_json::to_value(&resp).unwrap_or(serde_json::Value::Null), - ) - .await; - } - } - Ok(( - jar, - Json(Envelope { - success: true, - data: resp, - }), - )) - } + Ok(resp) => Ok(( + jar, + Json(Envelope { + success: true, + data: resp, + }), + )), Err(err) => { if anonymous_reservation_date.is_some() { if let Err(refund_err) = @@ -430,7 +475,10 @@ async fn compress_json( request_hash.as_deref(), ) { if idem_acquired { - let _ = idempotency::abort(&state, scope, idem_key, request_hash).await; + if let Some(owner) = idem_owner { + let _ = + idempotency::abort(&state, scope, idem_key, request_hash, owner).await; + } } } Err(err) @@ -525,6 +573,7 @@ async fn compress_direct( let quota_ctx = admission.quota_ctx; let mut idem_acquired = false; + let mut idem_owner = None; if let (Some(scope), Some(idem_key), Some(request_hash)) = ( idempotency_scope, idempotency_key.as_deref(), @@ -566,13 +615,33 @@ async fn compress_direct( "请求正在处理中，请稍后重试", )); } - idempotency::BeginResult::Acquired => { + idempotency::BeginResult::Acquired { owner } => { idem_acquired = true; + idem_owner = Some(owner); } } } - let op: Result<(axum::response::Response, DirectIdempotencyData), AppError> = (async { + let operation_state = state.clone(); + let operation_principal = principal.clone(); + let operation_quota_ctx = quota_ctx.clone(); + let operation_idempotency_key = idempotency_key.clone(); + let operation_request_hash = request_hash.clone(); + let operation_idem_owner = idem_owner; + let operation = spawn_detached_operation(async move { + let state = operation_state; + let principal = operation_principal; + let quota_ctx = operation_quota_ctx; + let idempotency_key = operation_idempotency_key; + let request_hash = operation_request_hash; + let _idempotency_heartbeat = start_idempotency_heartbeat( + &state, + idempotency_scope, + idempotency_key.as_deref(), + request_hash.as_deref(), + operation_idem_owner, + ); + let op: Result<(axum::response::Response, DirectIdempotencyData), AppError> = (async { match "a_ctx { QuotaContext::User(billing) => ensure_quota_available(&state, billing, 1).await?, QuotaContext::ApiKey(billing, _) => ensure_quota_available(&state, billing, 1).await?, @@ -617,8 +686,11 @@ async fn compress_direct( let retention_hours = retention.num_hours(); let object_key = storage::result_key(retention_hours, task_id, file_id, format_out.extension()); - let stored = storage::store_bytes( + let tracked = object_lifecycle::store_tracked_bytes( &state, + task_id, + Some(file_id), + "result", &object_key, compressed.clone(), format_out.content_type(), @@ -626,6 +698,23 @@ async fn compress_direct( .await?; let expires_at = Utc::now() + retention; + let idem_data = DirectIdempotencyData { + file_id, + format_out: format_out.as_str().to_string(), + original_size, + compressed_size, + saved_bytes, + saved_percent, + units_charged: metered_units(charge_units), + }; + let idem_completion = build_idempotency_completion( + idem_acquired, + operation_idem_owner, + idempotency_scope, + idempotency_key.as_deref(), + request_hash.as_deref(), + &idem_data, + )?; if let Err(err) = record_task_and_metering( &state, @@ -633,7 +722,7 @@ async fn compress_direct( ip, task_id, file_id, - &stored, + &tracked, &req.file_name, req.max_width, req.max_height, @@ -648,56 +737,45 @@ async fn compress_direct( retention_hours, "a_ctx, charge_units, + idem_completion.as_ref(), ) .await { - let _ = storage::delete_object( - &state, - &storage::ObjectLocator { - backend: stored.backend.clone(), - endpoint_id: stored.endpoint_id, - key: stored.key.clone(), - }, - ) - .await; - return Err(err); - } - - let idem_data = DirectIdempotencyData { - file_id, - format_out: format_out.as_str().to_string(), - original_size, - compressed_size, - saved_bytes, - saved_percent, - units_charged: metered_units(charge_units), - }; - let response = direct_response(compressed, format_out, &idem_data); - Ok((response, idem_data)) - }) - .await; - - match op { - Ok((response, idem_data)) => { - if let (Some(scope), Some(idem_key), Some(request_hash)) = ( - idempotency_scope, - idempotency_key.as_deref(), - request_hash.as_deref(), - ) { - if idem_acquired { - let _ = idempotency::complete( + match sync_result_was_committed(&state, task_id, file_id, &tracked).await { + Ok(true) => { + tracing::warn!(task_id = %task_id, file_id = %file_id, error = %err, "direct result commit response was lost; recovered committed publication"); + } + Ok(false) => { + if let Err(cleanup_err) = object_lifecycle::schedule_tracked_delete( &state, - scope, - idem_key, - request_hash, - 200, - serde_json::to_value(&idem_data).unwrap_or(serde_json::Value::Null), + &tracked, + Some(&err), ) - .await; + .await + { + tracing::error!(task_id = %task_id, storage_object_id = %tracked.lifecycle_id, error = %cleanup_err, "failed to persist rejected direct result cleanup"); + } + return Err(err); + } + Err(probe_err) => { + tracing::error!(task_id = %task_id, storage_object_id = %tracked.lifecycle_id, error = %probe_err, original_error = %err, "direct result commit state is unknown; staging lease will reconcile object"); + return Err(err); } } - Ok((jar, response)) } + + let response = direct_response(compressed, format_out, &idem_data); + Ok((response, idem_data)) + }) + .await; + op + }); + let op = operation.await.map_err(|err| { + AppError::new(ErrorCode::Internal, "直接压缩后台任务异常退出").with_source(err) + })?; + + match op { + Ok((response, _idem_data)) => Ok((jar, response)), Err(err) => { if let (Some(scope), Some(idem_key), Some(request_hash)) = ( idempotency_scope, @@ -705,7 +783,10 @@ async fn compress_direct( request_hash.as_deref(), ) { if idem_acquired { - let _ = idempotency::abort(&state, scope, idem_key, request_hash).await; + if let Some(owner) = idem_owner { + let _ = + idempotency::abort(&state, scope, idem_key, request_hash, owner).await; + } } } Err(err) @@ -713,6 +794,45 @@ async fn compress_direct( } } +async fn sync_result_was_committed( + state: &AppState, + task_id: Uuid, + file_id: Uuid, + tracked: &object_lifecycle::TrackedStoredObject, +) -> Result { + sqlx::query_scalar( + r#" + SELECT EXISTS( + SELECT 1 + FROM tasks AS task + JOIN task_files AS file ON file.task_id = task.id + JOIN storage_objects AS object ON object.id = $3 + WHERE task.id = $1 + AND task.status = 'completed' + AND file.id = $2 + AND file.status = 'completed' + AND file.storage_backend = $4 + AND file.storage_endpoint_id IS NOT DISTINCT FROM $5 + AND COALESCE(file.storage_key, file.storage_path) = $6 + AND object.state = 'published' + AND object.task_id = task.id + AND object.task_file_id = file.id + ) + "#, + ) + .bind(task_id) + .bind(file_id) + .bind(tracked.lifecycle_id) + .bind(&tracked.stored.backend) + .bind(tracked.stored.endpoint_id) + .bind(&tracked.stored.key) + .fetch_one(&state.db) + .await + .map_err(|err| { + AppError::new(ErrorCode::StorageUnavailable, "核验同步压缩提交结果失败").with_source(err) + }) +} + #[derive(Debug, FromRow)] struct DirectReplayRow { storage_backend: String, @@ -742,6 +862,7 @@ async fn load_direct_replay_bytes( FROM task_files f JOIN tasks t ON t.id = f.task_id WHERE f.id = $1 AND t.user_id = $2 + AND t.deletion_started_at IS NULL "#, ) .bind(file_id) @@ -762,6 +883,7 @@ async fn load_direct_replay_bytes( FROM task_files f JOIN tasks t ON t.id = f.task_id WHERE f.id = $1 AND t.api_key_id = $2 + AND t.deletion_started_at IS NULL "#, ) .bind(file_id) @@ -946,6 +1068,63 @@ enum QuotaContext { ApiKey(BillingContext, Uuid), } +struct IdempotencyCompletion { + scope: idempotency::Scope, + owner: Uuid, + key: String, + request_hash: String, + response_body: serde_json::Value, +} + +fn build_idempotency_completion( + acquired: bool, + owner: Option, + scope: Option, + key: Option<&str>, + request_hash: Option<&str>, + response: &T, +) -> Result, AppError> { + if !acquired { + return Ok(None); + } + let (owner, scope, key, request_hash) = match (owner, scope, key, request_hash) { + (Some(owner), Some(scope), Some(key), Some(request_hash)) => { + (owner, scope, key, request_hash) + } + _ => return Err(AppError::new(ErrorCode::Internal, "幂等请求上下文不完整")), + }; + Ok(Some(IdempotencyCompletion { + scope, + owner, + key: key.to_string(), + request_hash: request_hash.to_string(), + response_body: serde_json::to_value(response).map_err(|err| { + AppError::new(ErrorCode::Internal, "序列化幂等响应失败").with_source(err) + })?, + })) +} + +fn start_idempotency_heartbeat( + state: &AppState, + scope: Option, + key: Option<&str>, + request_hash: Option<&str>, + owner: Option, +) -> Option { + match (scope, key, request_hash, owner) { + (Some(scope), Some(key), Some(request_hash), Some(owner)) => { + Some(idempotency::start_lease_heartbeat( + state.clone(), + scope, + key.to_string(), + request_hash.to_string(), + owner, + )) + } + _ => None, + } +} + struct SingleAdmission { retention: Duration, quota_ctx: QuotaContext, @@ -1044,7 +1223,7 @@ async fn record_task_and_metering( client_ip: IpAddr, task_id: Uuid, file_id: Uuid, - stored: &storage::StoredObject, + tracked: &object_lifecycle::TrackedStoredObject, original_name: &str, max_width: Option, max_height: Option, @@ -1059,7 +1238,9 @@ async fn record_task_and_metering( retention_hours: i64, quota_ctx: &QuotaContext, charge_units: bool, + idempotency_completion: Option<&IdempotencyCompletion>, ) -> Result<(), AppError> { + let stored = &tracked.stored; let (user_id, session_id, api_key_id, source) = match principal { context::Principal::Anonymous { session_id } => { (None, Some(session_id.clone()), None, "web") @@ -1191,6 +1372,20 @@ async fn record_task_and_metering( } } + object_lifecycle::publish_in_tx(&mut tx, tracked).await?; + if let Some(idempotency_completion) = idempotency_completion { + idempotency::complete_in_tx( + &mut tx, + idempotency_completion.scope, + &idempotency_completion.key, + &idempotency_completion.request_hash, + idempotency_completion.owner, + 200, + idempotency_completion.response_body.clone(), + ) + .await?; + } + tx.commit() .await .map_err(|err| AppError::new(ErrorCode::Internal, "提交事务失败").with_source(err))?; @@ -1244,6 +1439,8 @@ async fn charge_one_unit( #[cfg(test)] mod tests { use super::*; + use std::sync::atomic::{AtomicBool, Ordering}; + use std::sync::Arc; #[test] fn direct_response_has_consistent_compression_headers() { @@ -1296,4 +1493,35 @@ mod tests { assert_eq!(metered_units(charged), expected_units); } } + + #[tokio::test] + async fn detached_sync_operation_survives_waiter_abort() { + let started = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Notify::new()); + let completed = Arc::new(AtomicBool::new(false)); + let request_started = started.clone(); + let request_release = release.clone(); + let request_completed = completed.clone(); + let request = tokio::spawn(async move { + spawn_detached_operation(async move { + request_started.notify_one(); + request_release.notified().await; + request_completed.store(true, Ordering::SeqCst); + }) + .await + .expect("detached operation panicked"); + }); + + started.notified().await; + request.abort(); + request.await.expect_err("request waiter was not aborted"); + release.notify_one(); + tokio::time::timeout(std::time::Duration::from_secs(1), async { + while !completed.load(Ordering::SeqCst) { + tokio::task::yield_now().await; + } + }) + .await + .expect("detached operation was canceled with its request waiter"); + } } diff --git a/src/api/downloads.rs b/src/api/downloads.rs index ad44a4f..10fc47f 100644 --- a/src/api/downloads.rs +++ b/src/api/downloads.rs @@ -1,5 +1,6 @@ use crate::api::context; use crate::error::{AppError, ErrorCode}; +use crate::services::object_lifecycle; use crate::services::storage; use crate::state::AppState; @@ -63,7 +64,7 @@ async fn download_file( t.expires_at FROM task_files f JOIN tasks t ON t.id = f.task_id - WHERE f.id = $1 + WHERE f.id = $1 AND t.deletion_started_at IS NULL "#, ) .bind(file_id) @@ -289,7 +290,7 @@ async fn download_task_zip( zip_storage_endpoint_id, zip_storage_key FROM tasks - WHERE id = $1 + WHERE id = $1 AND deletion_started_at IS NULL "#, ) .bind(task_id) @@ -451,6 +452,7 @@ async fn claim_zip_build(state: &AppState, task_id: Uuid) -> Result NOW() AND ( @@ -474,7 +476,7 @@ async fn claim_zip_build(state: &AppState, task_id: Uuid) -> Result NOW() + AND deletion_started_at IS NULL "#, ) .bind(task_id) @@ -711,7 +714,7 @@ async fn build_zip_attempt( rows: &[TaskZipFileRow], temp_dir: &std::path::Path, zip_path: &std::path::Path, -) -> Result { +) -> Result { tokio::fs::create_dir_all(temp_dir).await.map_err(|err| { AppError::new(ErrorCode::StorageUnavailable, "创建 ZIP 临时目录失败").with_source(err) })?; @@ -766,15 +769,27 @@ async fn build_zip_attempt( 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 + object_lifecycle::store_tracked_file( + state, + task_id, + "zip_attempt", + &object_key, + zip_path, + "application/zip", + ) + .await } async fn publish_zip_attempt( state: &AppState, task_id: Uuid, token: Uuid, - stored: storage::StoredObject, + tracked: object_lifecycle::TrackedStoredObject, ) -> Result { + let stored = &tracked.stored; + let mut tx = state.db.begin().await.map_err(|err| { + AppError::new(ErrorCode::Internal, "开启 ZIP 发布事务失败").with_source(err) + })?; let published = sqlx::query( r#" UPDATE tasks @@ -789,6 +804,7 @@ async fn publish_zip_attempt( AND zip_build_token = $2 AND zip_storage_key IS NULL AND expires_at > NOW() + AND deletion_started_at IS NULL "#, ) .bind(task_id) @@ -798,24 +814,36 @@ async fn publish_zip_attempt( .bind(&stored.key) .bind(&stored.etag) .bind(stored.size as i64) - .execute(&state.db) + .execute(&mut *tx) .await; match published { - Ok(result) if result.rows_affected() == 1 => Ok(storage::ObjectLocator { - backend: stored.backend, - endpoint_id: stored.endpoint_id, - key: stored.key, - }), + Ok(result) if result.rows_affected() == 1 => { + if let Err(err) = object_lifecycle::publish_in_tx(&mut tx, &tracked).await { + tx.rollback().await.ok(); + delete_unpublished_zip(state, task_id, token, &tracked).await; + return Err(err); + } + tx.commit().await.map_err(|err| { + AppError::new(ErrorCode::Internal, "提交 ZIP 发布事务失败").with_source(err) + })?; + Ok(storage::ObjectLocator { + backend: tracked.stored.backend, + endpoint_id: tracked.stored.endpoint_id, + key: tracked.stored.key, + }) + } Ok(_) => { - delete_unpublished_zip(state, task_id, token, &stored).await; + tx.rollback().await.ok(); + delete_unpublished_zip(state, task_id, token, &tracked).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; + tx.rollback().await.ok(); + delete_unpublished_zip(state, task_id, token, &tracked).await; Err(AppError::new(ErrorCode::Internal, "记录 ZIP 对象失败").with_source(err)) } } @@ -825,28 +853,13 @@ async fn delete_unpublished_zip( state: &AppState, task_id: Uuid, token: Uuid, - stored: &storage::StoredObject, + tracked: &object_lifecycle::TrackedStoredObject, ) { - 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"); + if let Err(err) = object_lifecycle::schedule_tracked_delete(state, tracked, None).await { + tracing::error!(task_id = %task_id, zip_build_token = %token, object_key = %tracked.stored.key, error = %err, "failed to persist unpublished ZIP deletion"); + } else if let Err(err) = object_lifecycle::cleanup_ready_objects(state, 1, Some(task_id)).await + { + tracing::warn!(task_id = %task_id, zip_build_token = %token, object_key = %tracked.stored.key, error = %err, "unpublished ZIP deletion deferred"); } release_zip_build(state, task_id, token).await; } @@ -856,7 +869,7 @@ async fn load_published_zip( 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", + "SELECT zip_storage_backend, zip_storage_endpoint_id, zip_storage_key FROM tasks WHERE id = $1 AND deletion_started_at IS NULL", ) .bind(task_id) .fetch_optional(&state.db) @@ -1247,8 +1260,10 @@ mod tests { tokio::fs::write(&unpublished_path, b"unpublished-zip") .await .expect("write unpublished ZIP fixture"); - let unpublished = storage::store_file( + let unpublished = object_lifecycle::store_tracked_file( &state, + deleted_task, + "zip_attempt", &storage::archive_attempt_key(24, deleted_task, deleted_token), &unpublished_path, "application/zip", @@ -1268,13 +1283,20 @@ mod tests { let orphan_read = storage::read_bytes( &state, &storage::ObjectLocator { - backend: unpublished.backend, - endpoint_id: unpublished.endpoint_id, - key: unpublished.key, + backend: unpublished.stored.backend, + endpoint_id: unpublished.stored.endpoint_id, + key: unpublished.stored.key, }, ) .await; assert!(orphan_read.is_err(), "unpublished ZIP object was orphaned"); + let unpublished_state: String = + sqlx::query_scalar("SELECT state FROM storage_objects WHERE id = $1") + .bind(unpublished.lifecycle_id) + .fetch_one(&pool) + .await + .expect("query unpublished ZIP lifecycle state"); + assert_eq!(unpublished_state, "deleted"); let cancelled_task = Uuid::new_v4(); insert_zip_task( @@ -1415,6 +1437,20 @@ mod tests { .execute(&pool) .await .expect("delete ZIP test tasks"); + sqlx::query("DELETE FROM storage_objects WHERE task_id = ANY($1)") + .bind( + &[ + task_id, + over_budget_task, + takeover_task, + deleted_task, + cancelled_task, + heartbeat_task, + ][..], + ) + .execute(&pool) + .await + .expect("delete ZIP lifecycle test rows"); if let Some(endpoint_id) = endpoint_id { sqlx::query("DELETE FROM storage_endpoints WHERE id = $1") .bind(endpoint_id) diff --git a/src/api/metrics.rs b/src/api/metrics.rs index d583777..5c28bc0 100644 --- a/src/api/metrics.rs +++ b/src/api/metrics.rs @@ -17,18 +17,33 @@ const SCRAPE_TIMEOUT: Duration = Duration::from_secs(2); pub async fn metrics(State(state): State) -> impl IntoResponse { let database = tokio::time::timeout( SCRAPE_TIMEOUT, - sqlx::query_scalar::<_, i64>( - "SELECT COUNT(*) FROM tasks WHERE status IN ('pending', 'processing')", + sqlx::query_as::<_, (i64, i64, i64, i64, i64)>( + r#" + SELECT + (SELECT COUNT(*) FROM tasks WHERE status IN ('pending', 'processing')), + (SELECT COUNT(*) FROM task_queue_outbox WHERE status IN ('pending', 'delivering')), + (SELECT COUNT(*) FROM task_queue_outbox WHERE status = 'dead'), + (SELECT COUNT(*) FROM storage_objects WHERE state = 'delete_pending'), + (SELECT COUNT(*) FROM storage_objects WHERE state = 'staging') + "#, ) .fetch_one(&state.db), ); let redis = tokio::time::timeout(SCRAPE_TIMEOUT, redis_queue_stats(state.redis.clone())); let (database, redis) = tokio::join!(database, redis); - let (database_up, active_tasks) = match database { - Ok(Ok(value)) => (1, value), - _ => (0, 0), - }; + let (database_up, active_tasks, outbox_pending, outbox_dead, delete_pending, staging) = + match database { + Ok(Ok((active, outbox_pending, outbox_dead, delete_pending, staging))) => ( + 1, + active, + outbox_pending, + outbox_dead, + delete_pending, + staging, + ), + _ => (0, 0, 0, 0, 0, 0), + }; let (redis_up, queue_length, pending, dead_length, cluster) = match redis { Ok(Ok((queue_length, pending, dead_length, cluster))) => { (1, queue_length, pending, dead_length, cluster) @@ -51,6 +66,28 @@ pub async fn metrics(State(state): State) -> impl IntoResponse { output.push_str("# HELP imageforge_active_tasks Current pending or processing tasks.\n"); output.push_str("# TYPE imageforge_active_tasks gauge\n"); let _ = writeln!(output, "imageforge_active_tasks {active_tasks}"); + output.push_str("# HELP imageforge_task_outbox Current durable task delivery states.\n"); + output.push_str("# TYPE imageforge_task_outbox gauge\n"); + let _ = writeln!( + output, + "imageforge_task_outbox{{state=\"pending\"}} {outbox_pending}" + ); + let _ = writeln!( + output, + "imageforge_task_outbox{{state=\"dead\"}} {outbox_dead}" + ); + output.push_str( + "# HELP imageforge_storage_object_lifecycle Current durable object cleanup states.\n", + ); + output.push_str("# TYPE imageforge_storage_object_lifecycle gauge\n"); + let _ = writeln!( + output, + "imageforge_storage_object_lifecycle{{state=\"delete_pending\"}} {delete_pending}" + ); + let _ = writeln!( + output, + "imageforge_storage_object_lifecycle{{state=\"staging\"}} {staging}" + ); output.push_str("# HELP imageforge_queue_messages Current Redis stream message counts.\n"); output.push_str("# TYPE imageforge_queue_messages gauge\n"); let _ = writeln!( diff --git a/src/api/mod.rs b/src/api/mod.rs index 6746c01..ba18165 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -53,12 +53,19 @@ pub async fn run(state: AppState) -> Result<(), AppError> { tracing::info!(addr = %addr, "API server listening"); let reconciliation_task = tokio::spawn(webhooks::reconciliation_loop(state.clone())); + let queue_dispatch_task = + tokio::spawn(crate::services::task_queue::dispatch_loop(state.clone())); + let object_lifecycle_task = tokio::spawn(crate::services::object_lifecycle::maintenance_loop( + state.clone(), + )); let serve_result = axum::serve( listener, app.into_make_service_with_connect_info::(), ) .await; reconciliation_task.abort(); + queue_dispatch_task.abort(); + object_lifecycle_task.abort(); serve_result .map_err(|err| AppError::new(ErrorCode::Internal, "HTTP 服务异常退出").with_source(err)) } diff --git a/src/api/tasks.rs b/src/api/tasks.rs index 3c535a9..99bbd38 100644 --- a/src/api/tasks.rs +++ b/src/api/tasks.rs @@ -8,8 +8,9 @@ use crate::services::compress; use crate::services::compress::{CompressionLevel, ImageFmt}; use crate::services::filename; use crate::services::idempotency; +use crate::services::object_lifecycle; use crate::services::quota; -use crate::services::storage; +use crate::services::task_queue; use crate::state::AppState; use axum::extract::{ConnectInfo, Multipart, Path, State}; @@ -127,6 +128,7 @@ async fn create_batch_task( } let mut idem_acquired = false; + let mut idem_owner = None; if let (Some(scope), Some(idem_key)) = (idempotency_scope, idempotency_key.as_deref()) { let begin_result = idempotency::begin( &state, @@ -181,8 +183,9 @@ async fn create_batch_task( "请求正在处理中，请稍后重试", )); } - idempotency::BeginResult::Acquired => { + idempotency::BeginResult::Acquired { owner } => { idem_acquired = true; + idem_owner = Some(owner); } } } @@ -190,8 +193,6 @@ 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 } => { @@ -299,75 +300,67 @@ async fn create_batch_task( })?; } + sqlx::query("INSERT INTO task_queue_outbox (task_id) VALUES ($1)") + .bind(task_id) + .execute(&mut *tx) + .await + .map_err(|err| { + AppError::new(ErrorCode::Internal, "创建任务队列 outbox 失败").with_source(err) + })?; + + let response = BatchCreateResponse { + task_id, + total_files: files.len() as i32, + status: "pending".to_string(), + status_url: format!("/api/v1/compress/tasks/{task_id}"), + }; + if let (Some(scope), Some(idem_key)) = (idempotency_scope, idempotency_key.as_deref()) { + if idem_acquired { + idempotency::complete_in_tx( + &mut tx, + scope, + idem_key, + &request_hash, + idem_owner.ok_or_else(|| { + AppError::new(ErrorCode::Internal, "幂等操作租约缺失") + })?, + 200, + serde_json::to_value(&response).map_err(|err| { + AppError::new(ErrorCode::Internal, "序列化幂等响应失败").with_source(err) + })?, + ) + .await?; + } + } + 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 { - 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"); - } + let dispatch_state = state.clone(); + tokio::spawn(async move { + if let Err(err) = task_queue::dispatch_task(&dispatch_state, task_id).await { + tracing::warn!(task_id = %task_id, error = %err, "immediate task queue dispatch deferred to outbox loop"); } - return Err(err); - } + }); - Ok(BatchCreateResponse { - task_id, - total_files: files.len() as i32, - status: "pending".to_string(), - status_url: format!("/api/v1/compress/tasks/{task_id}"), - }) + Ok(response) }) .await; match create_result { - Ok(resp) => { - if let (Some(scope), Some(idem_key)) = (idempotency_scope, idempotency_key.as_deref()) { - if idem_acquired { - let _ = idempotency::complete( - &state, - scope, - idem_key, - &request_hash, - 200, - serde_json::to_value(&resp).unwrap_or(serde_json::Value::Null), - ) - .await; - } - } - Ok(( - jar, - Json(Envelope { - success: true, - data: resp, - }), - )) - } + Ok(resp) => Ok(( + jar, + Json(Envelope { + success: true, + data: resp, + }), + )), Err(err) => { if anonymous_reserved_units > 0 { if let context::Principal::Anonymous { session_id } = &principal { - 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 - } - } - } else { - !task_persisted - }; - if should_refund_directly { + if !task_persisted { if let Some(date) = anonymous_quota_date { if let Err(refund_err) = quota::refund_anonymous_reservation_once( &state, @@ -387,10 +380,13 @@ async fn create_batch_task( } if let (Some(scope), Some(idem_key)) = (idempotency_scope, idempotency_key.as_deref()) { if idem_acquired { - let _ = idempotency::abort(&state, scope, idem_key, &request_hash).await; + if let Some(owner) = idem_owner { + let _ = + idempotency::abort(&state, scope, idem_key, &request_hash, owner).await; + } } } - if cleanup_inputs_on_error { + if !task_persisted { cleanup_task_input_dir(&state, task_id).await; } Err(err) @@ -484,107 +480,6 @@ fn plan_upload_limits(plan: &Plan) -> Result { }) } -async fn enqueue_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> { - let mut conn = state.redis.clone(); - let now = Utc::now().to_rfc3339(); - redis::cmd("XADD") - .arg("stream:compress_jobs") - .arg("MAXLEN") - .arg("~") - .arg(100_000) - .arg("*") - .arg("task_id") - .arg(task_id.to_string()) - .arg("created_at") - .arg(now) - .query_async::<_, redis::Value>(&mut conn) - .await - .map_err(|err| AppError::new(ErrorCode::Internal, "写入队列失败").with_source(err))?; - 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, @@ -948,7 +843,7 @@ async fn get_task( user_id, session_id FROM tasks - WHERE id = $1 + WHERE id = $1 AND deletion_started_at IS NULL "#, ) .bind(task_id) @@ -1051,7 +946,7 @@ async fn cancel_task( context::require_api_permission(&principal, &["compress", "batch_compress"])?; let task = sqlx::query_as::<_, TaskRow>( - "SELECT status::text AS status, total_files, completed_files, failed_files, created_at, completed_at, expires_at, user_id, session_id FROM tasks WHERE id = $1", + "SELECT status::text AS status, total_files, completed_files, failed_files, created_at, completed_at, expires_at, user_id, session_id FROM tasks WHERE id = $1 AND deletion_started_at IS NULL", ) .bind(task_id) .fetch_optional(&state.db) @@ -1077,7 +972,7 @@ async fn cancel_task( } let updated = sqlx::query( - "UPDATE tasks SET status = 'cancelled', completed_at = NOW() WHERE id = $1 AND status IN ('pending', 'processing')", + "UPDATE tasks SET status = 'cancelled', completed_at = NOW() WHERE id = $1 AND deletion_started_at IS NULL AND status IN ('pending', 'processing')", ) .bind(task_id) .execute(&state.db) @@ -1116,13 +1011,16 @@ async fn delete_task( let (jar, principal) = context::authenticate(&state, jar, &headers, ip).await?; context::require_api_permission(&principal, &["compress", "batch_compress"])?; + let mut tx = state.db.begin().await.map_err(|err| { + AppError::new(ErrorCode::Internal, "开启任务删除事务失败").with_source(err) + })?; let task = sqlx::query_as::<_, TaskRow>( - "SELECT status::text AS status, total_files, completed_files, failed_files, created_at, completed_at, expires_at, user_id, session_id FROM tasks WHERE id = $1", + "SELECT status::text AS status, total_files, completed_files, failed_files, created_at, completed_at, expires_at, user_id, session_id FROM tasks WHERE id = $1 FOR UPDATE", ) .bind(task_id) - .fetch_optional(&state.db) + .fetch_optional(&mut *tx) .await - .map_err(|err| AppError::new(ErrorCode::Internal, "查询任务失败").with_source(err))? + .map_err(|err| AppError::new(ErrorCode::Internal, "锁定待删除任务失败").with_source(err))? .ok_or_else(|| AppError::new(ErrorCode::NotFound, "任务不存在"))?; authorize_task( @@ -1138,87 +1036,18 @@ async fn delete_task( )); } - if task.status == "pending" { - let updated = sqlx::query( - "UPDATE tasks SET status = 'cancelled', completed_at = NOW() WHERE id = $1 AND status = 'pending'", - ) - .bind(task_id) - .execute(&state.db) - .await - .map_err(|err| AppError::new(ErrorCode::Internal, "锁定待删除任务失败").with_source(err))?; - if updated.rows_affected() == 0 { - return Err(AppError::new( - ErrorCode::InvalidRequest, - "任务状态已变化，请刷新后重试", - )); - } + if !object_lifecycle::mark_task_deleting(&mut tx, task_id, "user", false).await? { + return Err(AppError::new( + ErrorCode::InvalidRequest, + "任务状态已变化，请刷新后重试", + )); } + tx.commit().await.map_err(|err| { + AppError::new(ErrorCode::Internal, "提交任务删除状态失败").with_source(err) + })?; - quota::settle_anonymous_task_reservation(&state, task_id).await?; - - let files = sqlx::query_as::<_, TaskStorageRow>( - r#" - SELECT storage_backend, storage_endpoint_id, - COALESCE(storage_key, storage_path) AS storage_key, - input_path - FROM task_files - WHERE task_id = $1 - "#, - ) - .bind(task_id) - .fetch_all(&state.db) - .await - .map_err(|err| AppError::new(ErrorCode::Internal, "查询文件失败").with_source(err))?; - - for file in files { - if let Some(key) = file.storage_key { - storage::delete_object( - &state, - &storage::ObjectLocator { - backend: file.storage_backend, - endpoint_id: file.storage_endpoint_id, - key, - }, - ) - .await?; - } - if let Some(input_path) = file.input_path { - let _ = tokio::fs::remove_file(input_path).await; - } - } - - let zip = sqlx::query_as::<_, TaskZipStorageRow>( - "SELECT zip_storage_backend, zip_storage_endpoint_id, zip_storage_key FROM tasks WHERE id = $1", - ) - .bind(task_id) - .fetch_one(&state.db) - .await - .map_err(|err| AppError::new(ErrorCode::Internal, "查询 ZIP 文件失败").with_source(err))?; - if let (Some(backend), Some(key)) = (zip.zip_storage_backend, zip.zip_storage_key) { - storage::delete_object( - &state, - &storage::ObjectLocator { - backend, - endpoint_id: zip.zip_storage_endpoint_id, - key, - }, - ) - .await?; - } - - let legacy_zip_path = format!("{}/zips/{task_id}.zip", state.config.storage_path); - let _ = tokio::fs::remove_file(legacy_zip_path).await; - let orig_dir = format!("{}/orig/{task_id}", state.config.storage_path); - let _ = tokio::fs::remove_dir_all(orig_dir).await; - - let deleted = sqlx::query("DELETE FROM tasks WHERE id = $1") - .bind(task_id) - .execute(&state.db) - .await - .map_err(|err| AppError::new(ErrorCode::Internal, "删除任务失败").with_source(err))?; - - if deleted.rows_affected() == 0 { - return Err(AppError::new(ErrorCode::NotFound, "任务不存在")); + if let Err(err) = object_lifecycle::finalize_task_deletion(&state, task_id).await { + tracing::warn!(task_id = %task_id, error = %err, "task deletion persisted and will be retried by lifecycle maintenance"); } Ok(( @@ -1230,21 +1059,6 @@ async fn delete_task( )) } -#[derive(Debug, FromRow)] -struct TaskStorageRow { - storage_backend: String, - storage_endpoint_id: Option, - storage_key: Option, - input_path: Option, -} - -#[derive(Debug, FromRow)] -struct TaskZipStorageRow { - zip_storage_backend: Option, - zip_storage_endpoint_id: Option, - zip_storage_key: Option, -} - fn authorize_task( principal: &context::Principal, user_id: Option, @@ -1334,7 +1148,7 @@ mod tests { #[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() { + async fn outbox_retries_ambiguous_delivery_and_dead_letters_pending_task() { let database_url = std::env::var("IMAGEFORGE_TEST_DATABASE_URL") .expect("IMAGEFORGE_TEST_DATABASE_URL must be set"); assert!( @@ -1410,6 +1224,11 @@ mod tests { .await .expect("insert pending batch file"); } + sqlx::query("INSERT INTO task_queue_outbox (task_id) VALUES ($1)") + .bind(task_id) + .execute(&pool) + .await + .expect("insert task outbox"); let mut redis = state.redis.clone(); let _: i64 = redis::cmd("DEL") @@ -1424,17 +1243,48 @@ mod tests { .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, "队列提交失败") + task_queue::dispatch_task(&state, task_id) .await - .expect("finalize enqueue failure"), - "pending task was not finalized" + .expect("dispatch outbox through WRONGTYPE"), + "outbox was not claimed" + ); + let after_first_failure: (String, String, i32) = sqlx::query_as( + r#" + SELECT task.status::text, outbox.status, outbox.attempts + FROM tasks AS task + JOIN task_queue_outbox AS outbox ON outbox.task_id = task.id + WHERE task.id = $1 + "#, + ) + .bind(task_id) + .fetch_one(&pool) + .await + .expect("query retryable outbox state"); + assert_eq!( + after_first_failure, + ("pending".to_string(), "pending".to_string(), 1) + ); + assert!( + tokio::fs::try_exists(&input_dir) + .await + .expect("check retained input directory"), + "a retryable Redis error deleted task inputs" + ); + + sqlx::query( + "UPDATE task_queue_outbox SET attempts = 19, next_attempt_at = NOW() WHERE task_id = $1", + ) + .bind(task_id) + .execute(&pool) + .await + .expect("advance outbox to final attempt"); + assert!( + task_queue::dispatch_task(&state, task_id) + .await + .expect("dispatch final outbox attempt"), + "final outbox attempt was not claimed" ); - cleanup_task_input_dir(&state, task_id).await; let task: (String, bool, i32, i32) = sqlx::query_as( r#" @@ -1459,8 +1309,22 @@ mod tests { .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("队列提交失败") + row.0 == "failed" + && row.1 + && row.2 + && row + .3 + .as_deref() + .is_some_and(|message| message.starts_with("队列持续不可用：")) })); + let outbox_status: (String, i32, bool) = sqlx::query_as( + "SELECT status, attempts, last_error IS NOT NULL FROM task_queue_outbox WHERE task_id = $1", + ) + .bind(task_id) + .fetch_one(&pool) + .await + .expect("query dead outbox"); + assert_eq!(outbox_status, ("dead".to_string(), 20, true)); assert!( !tokio::fs::try_exists(&input_dir) .await @@ -1490,11 +1354,225 @@ mod tests { .query_async(&mut redis) .await .expect("remove WRONGTYPE fixture"); - sqlx::query("DELETE FROM tasks WHERE id = $1") - .bind(task_id) + + let ambiguous_task = Uuid::new_v4(); + sqlx::query( + r#" + INSERT INTO tasks ( + id, session_id, status, total_files, expires_at, retention_hours + ) VALUES ($1, $2, 'pending', 0, NOW() + INTERVAL '1 day', 24) + "#, + ) + .bind(ambiguous_task) + .bind(format!("ambiguous-{marker}")) + .execute(&pool) + .await + .expect("insert ambiguous delivery task"); + sqlx::query("INSERT INTO task_queue_outbox (task_id) VALUES ($1)") + .bind(ambiguous_task) .execute(&pool) .await - .expect("delete batch test task"); + .expect("insert ambiguous delivery outbox"); + task_queue::enqueue_task(&state, ambiguous_task) + .await + .expect("simulate XADD success with lost reply"); + assert!(task_queue::dispatch_task(&state, ambiguous_task) + .await + .expect("retry ambiguous delivery")); + let ambiguous_stream_len: i64 = redis::cmd("XLEN") + .arg("stream:compress_jobs") + .query_async(&mut redis) + .await + .expect("count duplicate ambiguous messages"); + assert_eq!(ambiguous_stream_len, 2); + sqlx::query("UPDATE tasks SET status = 'completed', completed_at = NOW() WHERE id = $1") + .bind(ambiguous_task) + .execute(&pool) + .await + .expect("simulate first duplicate message completing the task"); + assert_eq!( + crate::worker::process_task(&state, ambiguous_task, Uuid::new_v4()) + .await + .expect("process ambiguous delivery task"), + TaskProcessOutcome::Done + ); + assert_eq!( + crate::worker::process_task(&state, ambiguous_task, Uuid::new_v4()) + .await + .expect("reprocess duplicate ambiguous delivery task"), + TaskProcessOutcome::Done + ); + let ambiguous_status: String = + sqlx::query_scalar("SELECT status::text FROM tasks WHERE id = $1") + .bind(ambiguous_task) + .fetch_one(&pool) + .await + .expect("query ambiguous task status"); + assert_eq!(ambiguous_status, "completed"); + + let concurrent_task = Uuid::new_v4(); + sqlx::query( + "INSERT INTO tasks (id, session_id, status, total_files, expires_at, retention_hours) VALUES ($1, $2, 'pending', 0, NOW() + INTERVAL '1 day', 24)", + ) + .bind(concurrent_task) + .bind(format!("concurrent-{marker}")) + .execute(&pool) + .await + .expect("insert concurrent dispatcher task"); + sqlx::query("INSERT INTO task_queue_outbox (task_id) VALUES ($1)") + .bind(concurrent_task) + .execute(&pool) + .await + .expect("insert concurrent dispatcher outbox"); + let barrier = Arc::new(tokio::sync::Barrier::new(20)); + let mut dispatchers = Vec::new(); + for _ in 0..20 { + let state = state.clone(); + let barrier = barrier.clone(); + dispatchers.push(tokio::spawn(async move { + barrier.wait().await; + task_queue::dispatch_task(&state, concurrent_task).await + })); + } + let mut claimed = 0; + for dispatcher in dispatchers { + if dispatcher + .await + .expect("dispatcher task panicked") + .expect("concurrent dispatcher failed") + { + claimed += 1; + } + } + assert_eq!(claimed, 1, "multiple dispatchers owned one outbox row"); + + let takeover_task = Uuid::new_v4(); + sqlx::query( + "INSERT INTO tasks (id, session_id, status, total_files, expires_at, retention_hours) VALUES ($1, $2, 'pending', 0, NOW() + INTERVAL '1 day', 24)", + ) + .bind(takeover_task) + .bind(format!("takeover-{marker}")) + .execute(&pool) + .await + .expect("insert dispatcher takeover task"); + sqlx::query( + r#" + INSERT INTO task_queue_outbox ( + task_id, status, attempts, lease_owner, lease_until + ) VALUES ($1, 'delivering', 1, $2, NOW() - INTERVAL '1 second') + "#, + ) + .bind(takeover_task) + .bind(Uuid::new_v4()) + .execute(&pool) + .await + .expect("insert expired dispatcher lease"); + assert!(task_queue::dispatch_task(&state, takeover_task) + .await + .expect("take over expired dispatcher lease")); + let takeover_state: (String, i32) = + sqlx::query_as("SELECT status, attempts FROM task_queue_outbox WHERE task_id = $1") + .bind(takeover_task) + .fetch_one(&pool) + .await + .expect("query dispatcher takeover state"); + assert_eq!(takeover_state, ("delivered".to_string(), 2)); + + let idem_user: Uuid = sqlx::query_scalar( + r#" + INSERT INTO users (email, username, password_hash, email_verified_at) + VALUES ($1, $2, 'test', NOW()) + RETURNING id + "#, + ) + .bind(format!("outbox-idem-{marker}@example.test")) + .bind(format!("outbox-idem-{marker}")) + .fetch_one(&pool) + .await + .expect("insert outbox idempotency user"); + let idem_key = format!("outbox-{marker}"); + let idem_hash = "b".repeat(64); + let idem_owner = match idempotency::begin( + &state, + idempotency::Scope::User(idem_user), + &idem_key, + &idem_hash, + 24, + ) + .await + .expect("acquire outbox idempotency key") + { + idempotency::BeginResult::Acquired { owner } => owner, + other => panic!("unexpected outbox idempotency begin result: {other:?}"), + }; + let idem_task = Uuid::new_v4(); + let idem_response = serde_json::json!({"task_id": idem_task}); + let mut idem_tx = pool + .begin() + .await + .expect("begin atomic task/outbox response tx"); + sqlx::query( + "INSERT INTO tasks (id, user_id, status, total_files, expires_at, retention_hours) VALUES ($1, $2, 'pending', 0, NOW() + INTERVAL '1 day', 24)", + ) + .bind(idem_task) + .bind(idem_user) + .execute(&mut *idem_tx) + .await + .expect("insert idempotent task"); + sqlx::query("INSERT INTO task_queue_outbox (task_id) VALUES ($1)") + .bind(idem_task) + .execute(&mut *idem_tx) + .await + .expect("insert idempotent task outbox"); + idempotency::complete_in_tx( + &mut idem_tx, + idempotency::Scope::User(idem_user), + &idem_key, + &idem_hash, + idem_owner, + 200, + idem_response.clone(), + ) + .await + .expect("persist atomic idempotent task response"); + idem_tx.commit().await.expect("commit idempotent task"); + match idempotency::begin( + &state, + idempotency::Scope::User(idem_user), + &idem_key, + &idem_hash, + 24, + ) + .await + .expect("replay idempotent task") + { + idempotency::BeginResult::Replay { response_body } => { + assert_eq!(response_body, idem_response) + } + other => panic!("same idempotency key did not replay task: {other:?}"), + } + + sqlx::query("DELETE FROM tasks WHERE id = ANY($1)") + .bind(vec![ + task_id, + ambiguous_task, + concurrent_task, + takeover_task, + idem_task, + ]) + .execute(&pool) + .await + .expect("delete batch outbox test tasks"); + sqlx::query("DELETE FROM users WHERE id = $1") + .bind(idem_user) + .execute(&pool) + .await + .expect("delete outbox idempotency user"); + let _: i64 = redis::cmd("DEL") + .arg("stream:compress_jobs") + .query_async(&mut redis) + .await + .expect("clean compression stream"); let _ = tokio::fs::remove_dir_all(&storage_root).await; } } diff --git a/src/services/idempotency.rs b/src/services/idempotency.rs index fa0d9e6..405b1cb 100644 --- a/src/services/idempotency.rs +++ b/src/services/idempotency.rs @@ -6,6 +6,8 @@ use serde_json::Value as JsonValue; use sqlx::FromRow; use uuid::Uuid; +const OPERATION_LEASE_MINUTES: i64 = 30; + #[derive(Debug, Clone, Copy)] pub enum Scope { User(Uuid), @@ -14,7 +16,7 @@ pub enum Scope { #[derive(Debug)] pub enum BeginResult { - Acquired, + Acquired { owner: Uuid }, Replay { response_body: JsonValue }, InProgress, } @@ -54,6 +56,8 @@ pub async fn begin( let now = Utc::now(); let expires_at = now + Duration::hours(ttl_hours.max(1)); + let owner = Uuid::new_v4(); + let lease_until = now + Duration::minutes(OPERATION_LEASE_MINUTES); cleanup_expired_for_key(state, scope, idempotency_key, now).await?; @@ -64,11 +68,11 @@ pub async fn begin( INSERT INTO idempotency_keys ( user_id, idempotency_key, request_hash, response_status, response_body, - expires_at + expires_at, lease_owner, lease_until ) VALUES ( $1, $2, $3, 0, NULL, - $4 + $4, $5, $6 ) ON CONFLICT DO NOTHING "#, @@ -77,6 +81,8 @@ pub async fn begin( .bind(idempotency_key) .bind(request_hash) .bind(expires_at) + .bind(owner) + .bind(lease_until) .execute(&state.db) .await } @@ -86,11 +92,11 @@ pub async fn begin( INSERT INTO idempotency_keys ( api_key_id, idempotency_key, request_hash, response_status, response_body, - expires_at + expires_at, lease_owner, lease_until ) VALUES ( $1, $2, $3, 0, NULL, - $4 + $4, $5, $6 ) ON CONFLICT DO NOTHING "#, @@ -99,6 +105,8 @@ pub async fn begin( .bind(idempotency_key) .bind(request_hash) .bind(expires_at) + .bind(owner) + .bind(lease_until) .execute(&state.db) .await } @@ -106,12 +114,15 @@ pub async fn begin( .map_err(|err| AppError::new(ErrorCode::Internal, "写入幂等记录失败").with_source(err))?; if inserted.rows_affected() > 0 { - return Ok(BeginResult::Acquired); + return Ok(BeginResult::Acquired { owner }); } let row = get_row(state, scope, idempotency_key, now).await?; let Some(row) = row else { - return Ok(BeginResult::Acquired); + return Err(AppError::new( + ErrorCode::StorageUnavailable, + "幂等记录状态已变化，请重试", + )); }; if row.request_hash != request_hash { @@ -122,6 +133,19 @@ pub async fn begin( } if row.response_status == 0 || row.response_body.is_none() { + if take_over_stale_operation( + state, + scope, + idempotency_key, + request_hash, + owner, + lease_until, + now, + ) + .await? + { + return Ok(BeginResult::Acquired { owner }); + } return Ok(BeginResult::InProgress); } @@ -130,6 +154,152 @@ pub async fn begin( }) } +async fn take_over_stale_operation( + state: &AppState, + scope: Scope, + idempotency_key: &str, + request_hash: &str, + owner: Uuid, + lease_until: DateTime, + now: DateTime, +) -> Result { + let updated = match scope { + Scope::User(user_id) => { + sqlx::query( + r#" + UPDATE idempotency_keys + SET lease_owner = $4, lease_until = $5 + WHERE user_id = $1 AND idempotency_key = $2 + AND request_hash = $3 AND response_status = 0 + AND (lease_until IS NULL OR lease_until <= $6) + "#, + ) + .bind(user_id) + .bind(idempotency_key) + .bind(request_hash) + .bind(owner) + .bind(lease_until) + .bind(now) + .execute(&state.db) + .await + } + Scope::ApiKey(api_key_id) => { + sqlx::query( + r#" + UPDATE idempotency_keys + SET lease_owner = $4, lease_until = $5 + WHERE api_key_id = $1 AND idempotency_key = $2 + AND request_hash = $3 AND response_status = 0 + AND (lease_until IS NULL OR lease_until <= $6) + "#, + ) + .bind(api_key_id) + .bind(idempotency_key) + .bind(request_hash) + .bind(owner) + .bind(lease_until) + .bind(now) + .execute(&state.db) + .await + } + } + .map_err(|err| AppError::new(ErrorCode::Internal, "接管过期幂等操作失败").with_source(err))?; + Ok(updated.rows_affected() == 1) +} + +pub struct LeaseHeartbeat(Option>); + +impl Drop for LeaseHeartbeat { + fn drop(&mut self) { + if let Some(stop) = self.0.take() { + let _ = stop.send(()); + } + } +} + +pub fn start_lease_heartbeat( + state: AppState, + scope: Scope, + idempotency_key: String, + request_hash: String, + owner: Uuid, +) -> LeaseHeartbeat { + let (stop_tx, mut stop_rx) = tokio::sync::oneshot::channel(); + tokio::spawn(async move { + let mut interval = tokio::time::interval(std::time::Duration::from_secs(60)); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + interval.tick().await; + loop { + tokio::select! { + _ = &mut stop_rx => break, + _ = interval.tick() => { + match renew_lease( + &state, + scope, + &idempotency_key, + &request_hash, + owner, + ).await { + Ok(true) => {} + Ok(false) => break, + Err(err) => tracing::warn!(error = %err, "failed to renew idempotency operation lease"), + } + } + } + } + }); + LeaseHeartbeat(Some(stop_tx)) +} + +async fn renew_lease( + state: &AppState, + scope: Scope, + idempotency_key: &str, + request_hash: &str, + owner: Uuid, +) -> Result { + let updated = match scope { + Scope::User(user_id) => { + sqlx::query( + r#" + UPDATE idempotency_keys + SET lease_until = NOW() + ($5 * INTERVAL '1 minute') + WHERE user_id = $1 AND idempotency_key = $2 + AND request_hash = $3 AND lease_owner = $4 + AND response_status = 0 + "#, + ) + .bind(user_id) + .bind(idempotency_key) + .bind(request_hash) + .bind(owner) + .bind(OPERATION_LEASE_MINUTES) + .execute(&state.db) + .await + } + Scope::ApiKey(api_key_id) => { + sqlx::query( + r#" + UPDATE idempotency_keys + SET lease_until = NOW() + ($5 * INTERVAL '1 minute') + WHERE api_key_id = $1 AND idempotency_key = $2 + AND request_hash = $3 AND lease_owner = $4 + AND response_status = 0 + "#, + ) + .bind(api_key_id) + .bind(idempotency_key) + .bind(request_hash) + .bind(owner) + .bind(OPERATION_LEASE_MINUTES) + .execute(&state.db) + .await + } + } + .map_err(|err| AppError::new(ErrorCode::Internal, "续租幂等操作失败").with_source(err))?; + Ok(updated.rows_affected() == 1) +} + pub async fn wait_for_replay( state: &AppState, scope: Scope, @@ -165,11 +335,12 @@ pub async fn wait_for_replay( } } -pub async fn complete( - state: &AppState, +pub async fn complete_in_tx( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, scope: Scope, idempotency_key: &str, request_hash: &str, + owner: Uuid, response_status: i32, response_body: JsonValue, ) -> Result<(), AppError> { @@ -179,10 +350,13 @@ pub async fn complete( r#" UPDATE idempotency_keys SET response_status = $4, - response_body = $5 + response_body = $5, + lease_owner = NULL, + lease_until = NULL WHERE user_id = $1 AND idempotency_key = $2 AND request_hash = $3 + AND lease_owner = $6 AND response_status = 0 "#, ) @@ -191,7 +365,8 @@ pub async fn complete( .bind(request_hash) .bind(response_status) .bind(response_body) - .execute(&state.db) + .bind(owner) + .execute(&mut **tx) .await } Scope::ApiKey(api_key_id) => { @@ -199,10 +374,13 @@ pub async fn complete( r#" UPDATE idempotency_keys SET response_status = $4, - response_body = $5 + response_body = $5, + lease_owner = NULL, + lease_until = NULL WHERE api_key_id = $1 AND idempotency_key = $2 AND request_hash = $3 + AND lease_owner = $6 AND response_status = 0 "#, ) @@ -211,16 +389,19 @@ pub async fn complete( .bind(request_hash) .bind(response_status) .bind(response_body) - .execute(&state.db) + .bind(owner) + .execute(&mut **tx) .await } } - .map_err(|err| AppError::new(ErrorCode::Internal, "写入幂等结果失败").with_source(err))?; + .map_err(|err| AppError::new(ErrorCode::Internal, "事务内写入幂等结果失败").with_source(err))?; - if updated.rows_affected() == 0 { - tracing::warn!("idempotency record not updated (maybe already completed?)"); + if updated.rows_affected() != 1 { + return Err(AppError::new( + ErrorCode::IdempotencyConflict, + "幂等请求所有权已变化，请重试", + )); } - Ok(()) } @@ -229,25 +410,28 @@ pub async fn abort( scope: Scope, idempotency_key: &str, request_hash: &str, + owner: Uuid, ) -> Result<(), AppError> { match scope { Scope::User(user_id) => { let _ = sqlx::query( - "DELETE FROM idempotency_keys WHERE user_id = $1 AND idempotency_key = $2 AND request_hash = $3 AND response_status = 0", + "DELETE FROM idempotency_keys WHERE user_id = $1 AND idempotency_key = $2 AND request_hash = $3 AND lease_owner = $4 AND response_status = 0", ) .bind(user_id) .bind(idempotency_key) .bind(request_hash) + .bind(owner) .execute(&state.db) .await; } Scope::ApiKey(api_key_id) => { let _ = sqlx::query( - "DELETE FROM idempotency_keys WHERE api_key_id = $1 AND idempotency_key = $2 AND request_hash = $3 AND response_status = 0", + "DELETE FROM idempotency_keys WHERE api_key_id = $1 AND idempotency_key = $2 AND request_hash = $3 AND lease_owner = $4 AND response_status = 0", ) .bind(api_key_id) .bind(idempotency_key) .bind(request_hash) + .bind(owner) .execute(&state.db) .await; } @@ -334,3 +518,149 @@ async fn get_row( Ok(row) } + +#[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 sqlx::postgres::PgPoolOptions; + use std::sync::Arc; + use tokio::sync::Semaphore; + + async fn test_state(database_url: String, redis_url: String) -> AppState { + let mut config = Config::from_env().expect("load idempotency test config"); + config.database_url = database_url.clone(); + config.redis_url = redis_url; + config.mail_enabled = false; + config.mail_log_links_when_disabled = false; + let db = PgPoolOptions::new() + .max_connections(8) + .connect(&database_url) + .await + .expect("connect idempotency test database"); + sqlx::migrate!().run(&db).await.expect("run migrations"); + let redis = redis::Client::open(config.redis_url.clone()) + .expect("create idempotency test Redis client") + .get_connection_manager() + .await + .expect("connect idempotency test Redis"); + AppState { + mailer: Arc::new(Mailer::new(&config).expect("create disabled test mailer")), + image_processing_semaphore: Arc::new(Semaphore::new(1)), + zip_build_semaphore: Arc::new(Semaphore::new(1)), + runtime_policy_cache: RuntimePolicyCache::new(), + storage_cache: StorageCache::new(), + config, + db, + redis, + } + } + + #[tokio::test] + #[ignore = "requires IMAGEFORGE_TEST_DATABASE_URL and IMAGEFORGE_TEST_REDIS_URL"] + async fn stale_operation_is_fenced_and_replay_remains_atomic() { + let database_url = std::env::var("IMAGEFORGE_TEST_DATABASE_URL") + .expect("IMAGEFORGE_TEST_DATABASE_URL is required"); + let redis_url = std::env::var("IMAGEFORGE_TEST_REDIS_URL") + .expect("IMAGEFORGE_TEST_REDIS_URL is required"); + let state = test_state(database_url, redis_url).await; + let marker = Uuid::new_v4().simple().to_string(); + let user_id: Uuid = sqlx::query_scalar( + r#" + INSERT INTO users (email, username, password_hash, email_verified_at) + VALUES ($1, $2, 'test', NOW()) + RETURNING id + "#, + ) + .bind(format!("idem-{marker}@example.test")) + .bind(format!("idem-{marker}")) + .fetch_one(&state.db) + .await + .expect("insert idempotency test user"); + let scope = Scope::User(user_id); + let key = format!("idem-{marker}"); + let request_hash = "a".repeat(64); + let owner_one = match begin(&state, scope, &key, &request_hash, 24) + .await + .expect("acquire first operation") + { + BeginResult::Acquired { owner } => owner, + other => panic!("unexpected first begin result: {other:?}"), + }; + assert!(matches!( + begin(&state, scope, &key, &request_hash, 24) + .await + .expect("probe live operation"), + BeginResult::InProgress + )); + + sqlx::query( + "UPDATE idempotency_keys SET lease_until = NOW() - INTERVAL '1 second' WHERE user_id = $1 AND idempotency_key = $2", + ) + .bind(user_id) + .bind(&key) + .execute(&state.db) + .await + .expect("expire first operation lease"); + let owner_two = match begin(&state, scope, &key, &request_hash, 24) + .await + .expect("take over stale operation") + { + BeginResult::Acquired { owner } => owner, + other => panic!("unexpected takeover result: {other:?}"), + }; + assert_ne!(owner_one, owner_two); + + let mut stale_tx = state.db.begin().await.expect("begin stale completion tx"); + let stale_error = complete_in_tx( + &mut stale_tx, + scope, + &key, + &request_hash, + owner_one, + 200, + serde_json::json!({"owner": "stale"}), + ) + .await + .expect_err("stale operation completed after takeover"); + assert_eq!(stale_error.code, ErrorCode::IdempotencyConflict); + stale_tx + .rollback() + .await + .expect("rollback stale completion"); + + let expected = serde_json::json!({"owner": "current"}); + let mut current_tx = state.db.begin().await.expect("begin current completion tx"); + complete_in_tx( + &mut current_tx, + scope, + &key, + &request_hash, + owner_two, + 200, + expected.clone(), + ) + .await + .expect("complete current operation"); + current_tx + .commit() + .await + .expect("commit current completion"); + match begin(&state, scope, &key, &request_hash, 24) + .await + .expect("replay completed operation") + { + BeginResult::Replay { response_body } => assert_eq!(response_body, expected), + other => panic!("unexpected replay result: {other:?}"), + } + + sqlx::query("DELETE FROM users WHERE id = $1") + .bind(user_id) + .execute(&state.db) + .await + .expect("clean idempotency test user"); + } +} diff --git a/src/services/mod.rs b/src/services/mod.rs index eeb1c15..62ce47c 100644 --- a/src/services/mod.rs +++ b/src/services/mod.rs @@ -6,7 +6,9 @@ pub mod filename; pub mod idempotency; pub mod mail; pub mod metrics; +pub mod object_lifecycle; pub mod quota; pub mod rate_limit; pub mod settings; pub mod storage; +pub mod task_queue; diff --git a/src/services/object_lifecycle.rs b/src/services/object_lifecycle.rs new file mode 100644 index 0000000..fff8116 --- /dev/null +++ b/src/services/object_lifecycle.rs @@ -0,0 +1,1039 @@ +use crate::error::{AppError, ErrorCode}; +use crate::services::{quota, storage}; +use crate::state::AppState; + +use bytes::Bytes; +use sqlx::{FromRow, Postgres, Transaction}; +use std::path::{Path, PathBuf}; +use std::time::Duration; +use uuid::Uuid; + +const STAGING_LEASE_SECONDS: i64 = 15 * 60; +const DELETE_LEASE_SECONDS: i64 = 60; +const CLEANUP_INTERVAL: Duration = Duration::from_secs(5); +const CLEANUP_BATCH_SIZE: i64 = 100; + +#[derive(Debug, Clone)] +pub struct TrackedStoredObject { + pub lifecycle_id: Uuid, + pub lease_owner: Uuid, + pub task_id: Uuid, + pub stored: storage::StoredObject, +} + +#[derive(Debug)] +struct StagedObject { + id: Uuid, + lease_owner: Uuid, + task_id: Uuid, +} + +struct StagingHeartbeat(Option>); + +impl Drop for StagingHeartbeat { + fn drop(&mut self) { + if let Some(stop) = self.0.take() { + let _ = stop.send(()); + } + } +} + +#[derive(Debug, FromRow)] +struct DeleteClaim { + id: Uuid, + task_id: Uuid, + backend: String, + storage_endpoint_id: Option, + object_key: String, + delete_attempts: i32, +} + +pub async fn maintenance_loop(state: AppState) { + let mut interval = tokio::time::interval(CLEANUP_INTERVAL); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + interval.tick().await; + if let Err(err) = finalize_deleting_tasks(&state, CLEANUP_BATCH_SIZE).await { + tracing::error!(error = %err, "task deletion lifecycle iteration failed"); + } + if let Err(err) = cleanup_ready_objects(&state, CLEANUP_BATCH_SIZE, None).await { + tracing::error!(error = %err, "storage object lifecycle cleanup iteration failed"); + } + } +} + +pub async fn store_tracked_bytes( + state: &AppState, + task_id: Uuid, + task_file_id: Option, + object_kind: &str, + key: &str, + bytes: Bytes, + content_type: &str, +) -> Result { + if let Some(endpoint) = storage::active_endpoint(state).await? { + let stage = stage_object( + state, + task_id, + task_file_id, + object_kind, + "s3", + Some(endpoint.id), + key, + ) + .await?; + let heartbeat = start_staging_heartbeat(state.clone(), stage.id, stage.lease_owner); + let upload = + storage::store_bytes_s3(state, &endpoint, key, bytes.clone(), content_type).await; + drop(heartbeat); + match upload { + Ok(stored) => return Ok(tracked(stage, stored)), + Err(err) => { + schedule_delete(state, stage.id, stage.lease_owner, Some(&err)).await?; + storage::log_local_fallback(state, &endpoint, key, &err); + } + } + } + + let local_key = storage::local_path(state, key)?; + let local_key = local_key.to_string_lossy().to_string(); + let stage = stage_object( + state, + task_id, + task_file_id, + object_kind, + "local", + None, + &local_key, + ) + .await?; + let heartbeat = start_staging_heartbeat(state.clone(), stage.id, stage.lease_owner); + let upload = storage::store_bytes_local(state, key, bytes.as_ref()).await; + drop(heartbeat); + match upload { + Ok(stored) => Ok(tracked(stage, stored)), + Err(err) => { + schedule_delete(state, stage.id, stage.lease_owner, Some(&err)).await?; + Err(err) + } + } +} + +pub async fn store_tracked_file( + state: &AppState, + task_id: Uuid, + object_kind: &str, + key: &str, + path: &Path, + content_type: &str, +) -> Result { + let size = tokio::fs::metadata(path) + .await + .map_err(|err| { + AppError::new(ErrorCode::StorageUnavailable, "读取待上传文件失败").with_source(err) + })? + .len(); + if let Some(endpoint) = storage::active_endpoint(state).await? { + let stage = stage_object( + state, + task_id, + None, + object_kind, + "s3", + Some(endpoint.id), + key, + ) + .await?; + let heartbeat = start_staging_heartbeat(state.clone(), stage.id, stage.lease_owner); + let upload = storage::store_file_s3(state, &endpoint, key, path, content_type, size).await; + drop(heartbeat); + match upload { + Ok(stored) => return Ok(tracked(stage, stored)), + Err(err) => { + schedule_delete(state, stage.id, stage.lease_owner, Some(&err)).await?; + storage::log_local_fallback(state, &endpoint, key, &err); + } + } + } + + let local_key = storage::local_path(state, key)?; + let local_key = local_key.to_string_lossy().to_string(); + let stage = stage_object(state, task_id, None, object_kind, "local", None, &local_key).await?; + let heartbeat = start_staging_heartbeat(state.clone(), stage.id, stage.lease_owner); + let upload = storage::store_file_local(state, key, path, size).await; + drop(heartbeat); + match upload { + Ok(stored) => Ok(tracked(stage, stored)), + Err(err) => { + schedule_delete(state, stage.id, stage.lease_owner, Some(&err)).await?; + Err(err) + } + } +} + +pub async fn publish_in_tx( + tx: &mut Transaction<'_, Postgres>, + tracked: &TrackedStoredObject, +) -> Result<(), AppError> { + let updated = sqlx::query( + r#" + UPDATE storage_objects + SET state = 'published', storage_etag = $3, size_bytes = $4, + lease_owner = NULL, lease_until = NULL, + published_at = NOW(), last_error = NULL, updated_at = NOW() + WHERE id = $1 AND state = 'staging' + AND lease_owner = $2 AND lease_until > NOW() + "#, + ) + .bind(tracked.lifecycle_id) + .bind(tracked.lease_owner) + .bind(&tracked.stored.etag) + .bind(tracked.stored.size as i64) + .execute(&mut **tx) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "发布存储对象状态失败").with_source(err))?; + if updated.rows_affected() != 1 { + return Err(AppError::new( + ErrorCode::StorageUnavailable, + "存储对象发布租约已失效", + )); + } + Ok(()) +} + +pub async fn schedule_tracked_delete( + state: &AppState, + tracked: &TrackedStoredObject, + error: Option<&AppError>, +) -> Result<(), AppError> { + schedule_delete(state, tracked.lifecycle_id, tracked.lease_owner, error).await +} + +pub async fn mark_task_deleting( + tx: &mut Transaction<'_, Postgres>, + task_id: Uuid, + reason: &str, + allow_processing: bool, +) -> Result { + let updated = sqlx::query( + r#" + UPDATE tasks + SET deletion_started_at = COALESCE(deletion_started_at, NOW()), + deletion_reason = COALESCE(deletion_reason, $2), + status = CASE + WHEN status IN ('pending', 'processing') THEN 'cancelled'::task_status + ELSE status + END, + completed_at = CASE + WHEN status IN ('pending', 'processing') THEN COALESCE(completed_at, NOW()) + ELSE completed_at + END, + lease_owner = CASE WHEN status IN ('pending', 'processing') THEN NULL ELSE lease_owner END, + lease_until = CASE WHEN status IN ('pending', 'processing') THEN NULL ELSE lease_until END + WHERE id = $1 AND ($3 OR status <> 'processing') + "#, + ) + .bind(task_id) + .bind(reason) + .bind(allow_processing) + .execute(&mut **tx) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "标记任务删除状态失败").with_source(err))?; + Ok(updated.rows_affected() == 1) +} + +pub async fn mark_expired_task(state: &AppState, task_id: Uuid) -> Result { + let mut tx = state.db.begin().await.map_err(|err| { + AppError::new(ErrorCode::Internal, "开启过期任务删除事务失败").with_source(err) + })?; + let expired: Option = + sqlx::query_scalar("SELECT id FROM tasks WHERE id = $1 AND expires_at < NOW() FOR UPDATE") + .bind(task_id) + .fetch_optional(&mut *tx) + .await + .map_err(|err| { + AppError::new(ErrorCode::Internal, "锁定过期任务失败").with_source(err) + })?; + if expired.is_none() { + tx.rollback().await.ok(); + return Ok(false); + } + let marked = mark_task_deleting(&mut tx, task_id, "expired", true).await?; + tx.commit().await.map_err(|err| { + AppError::new(ErrorCode::Internal, "提交过期任务删除状态失败").with_source(err) + })?; + Ok(marked) +} + +pub async fn finalize_task_deletion(state: &AppState, task_id: Uuid) -> Result { + quota::settle_anonymous_task_reservation(state, task_id).await?; + let mut tx = state.db.begin().await.map_err(|err| { + AppError::new(ErrorCode::Internal, "开启任务删除收口事务失败").with_source(err) + })?; + let marked: Option = sqlx::query_scalar( + "SELECT deletion_started_at IS NOT NULL 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))?; + if marked != Some(true) { + tx.rollback().await.ok(); + return Ok(false); + } + + persist_task_object_pointers(&mut tx, state, task_id).await?; + sqlx::query( + r#" + UPDATE storage_objects + SET state = 'delete_pending', + lease_owner = CASE WHEN state = 'staging' THEN lease_owner ELSE NULL END, + lease_until = CASE WHEN state = 'staging' THEN lease_until ELSE NULL END, + next_attempt_at = NOW(), updated_at = NOW() + WHERE task_id = $1 AND state IN ('staging', 'published') + "#, + ) + .bind(task_id) + .execute(&mut *tx) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "安排任务对象删除失败").with_source(err))?; + sqlx::query("DELETE FROM tasks WHERE id = $1") + .bind(task_id) + .execute(&mut *tx) + .await + .map_err(|err| { + AppError::new(ErrorCode::Internal, "删除任务数据库记录失败").with_source(err) + })?; + tx.commit().await.map_err(|err| { + AppError::new(ErrorCode::Internal, "提交任务删除收口事务失败").with_source(err) + })?; + + cleanup_ready_objects(state, CLEANUP_BATCH_SIZE, Some(task_id)).await?; + Ok(true) +} + +pub async fn finalize_deleting_tasks(state: &AppState, limit: i64) -> Result { + let task_ids: Vec = sqlx::query_scalar( + r#" + SELECT id FROM tasks + WHERE deletion_started_at IS NOT NULL + ORDER BY deletion_started_at + LIMIT $1 + "#, + ) + .bind(limit.max(1)) + .fetch_all(&state.db) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "查询待收口删除任务失败").with_source(err))?; + let mut finalized = 0; + for task_id in task_ids { + if finalize_task_deletion(state, task_id).await? { + finalized += 1; + } + } + Ok(finalized) +} + +pub async fn cleanup_ready_objects( + state: &AppState, + limit: i64, + task_id: Option, +) -> Result { + let lease_owner = Uuid::new_v4(); + let claims = claim_deletions(state, lease_owner, limit.max(1), task_id).await?; + let count = claims.len(); + for claim in claims { + delete_claimed_object(state, lease_owner, claim).await?; + } + Ok(count) +} + +async fn stage_object( + state: &AppState, + task_id: Uuid, + task_file_id: Option, + object_kind: &str, + backend: &str, + storage_endpoint_id: Option, + object_key: &str, +) -> Result { + let id = Uuid::new_v4(); + let lease_owner = Uuid::new_v4(); + sqlx::query( + r#" + INSERT INTO storage_objects ( + id, task_id, task_file_id, object_kind, state, + backend, storage_endpoint_id, object_key, lease_owner, lease_until + ) VALUES ( + $1, $2, $3, $4, 'staging', $5, $6, $7, + $8, NOW() + ($9 * INTERVAL '1 second') + ) + "#, + ) + .bind(id) + .bind(task_id) + .bind(task_file_id) + .bind(object_kind) + .bind(backend) + .bind(storage_endpoint_id) + .bind(object_key) + .bind(lease_owner) + .bind(STAGING_LEASE_SECONDS) + .execute(&state.db) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "登记待上传对象失败").with_source(err))?; + Ok(StagedObject { + id, + lease_owner, + task_id, + }) +} + +fn tracked(stage: StagedObject, stored: storage::StoredObject) -> TrackedStoredObject { + TrackedStoredObject { + lifecycle_id: stage.id, + lease_owner: stage.lease_owner, + task_id: stage.task_id, + stored, + } +} + +fn start_staging_heartbeat( + state: AppState, + lifecycle_id: Uuid, + lease_owner: Uuid, +) -> StagingHeartbeat { + start_staging_heartbeat_with_interval( + state, + lifecycle_id, + lease_owner, + Duration::from_secs(60), + STAGING_LEASE_SECONDS, + ) +} + +fn start_staging_heartbeat_with_interval( + state: AppState, + lifecycle_id: Uuid, + lease_owner: Uuid, + heartbeat_interval: Duration, + lease_seconds: i64, +) -> StagingHeartbeat { + let (stop_tx, mut stop_rx) = tokio::sync::oneshot::channel(); + tokio::spawn(async move { + let mut interval = tokio::time::interval(heartbeat_interval); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + tokio::select! { + _ = &mut stop_rx => break, + _ = interval.tick() => { + match sqlx::query( + r#" + UPDATE storage_objects + SET lease_until = NOW() + ($3 * INTERVAL '1 second'), + updated_at = NOW() + WHERE id = $1 AND state = 'staging' AND lease_owner = $2 + "#, + ) + .bind(lifecycle_id) + .bind(lease_owner) + .bind(lease_seconds) + .execute(&state.db) + .await + { + Ok(result) if result.rows_affected() == 1 => {} + Ok(_) => break, + Err(error) => { + tracing::warn!( + storage_object_id = %lifecycle_id, + error = %error, + "storage object staging lease renewal failed" + ); + } + } + } + } + } + }); + StagingHeartbeat(Some(stop_tx)) +} + +async fn schedule_delete( + state: &AppState, + lifecycle_id: Uuid, + lease_owner: Uuid, + error: Option<&AppError>, +) -> Result<(), AppError> { + let last_error = error.map(|error| truncate_error(&error.to_string())); + let updated = sqlx::query( + r#" + UPDATE storage_objects + SET state = 'delete_pending', lease_owner = NULL, lease_until = NULL, + next_attempt_at = NOW(), last_error = COALESCE($3, last_error), + updated_at = NOW() + WHERE id = $1 + AND state IN ('staging', 'published', 'delete_pending') + AND (lease_owner = $2 OR lease_owner IS NULL) + "#, + ) + .bind(lifecycle_id) + .bind(lease_owner) + .bind(last_error) + .execute(&state.db) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "安排对象删除失败").with_source(err))?; + if updated.rows_affected() == 0 { + tracing::warn!(storage_object_id = %lifecycle_id, "storage object deletion was already owned or completed"); + } + Ok(()) +} + +async fn persist_task_object_pointers( + tx: &mut Transaction<'_, Postgres>, + state: &AppState, + task_id: Uuid, +) -> Result<(), AppError> { + sqlx::query( + r#" + INSERT INTO storage_objects ( + task_id, task_file_id, object_kind, state, + backend, storage_endpoint_id, object_key, storage_etag, + size_bytes, published_at + ) + SELECT + file.task_id, file.id, 'result', 'published', + file.storage_backend, file.storage_endpoint_id, + COALESCE(file.storage_key, file.storage_path), file.storage_etag, + file.compressed_size, COALESCE(file.completed_at, NOW()) + FROM task_files AS file + WHERE file.task_id = $1 + AND COALESCE(file.storage_key, file.storage_path) IS NOT NULL + ON CONFLICT DO NOTHING + "#, + ) + .bind(task_id) + .execute(&mut **tx) + .await + .map_err(|err| { + AppError::new(ErrorCode::Internal, "持久化任务结果对象指针失败").with_source(err) + })?; + sqlx::query( + r#" + INSERT INTO storage_objects ( + task_id, task_file_id, object_kind, state, + backend, object_key, published_at + ) + SELECT task_id, id, 'input', 'published', 'local', input_path, NOW() + FROM task_files + WHERE task_id = $1 AND input_path IS NOT NULL + ON CONFLICT DO NOTHING + "#, + ) + .bind(task_id) + .execute(&mut **tx) + .await + .map_err(|err| { + AppError::new(ErrorCode::Internal, "持久化任务输入对象指针失败").with_source(err) + })?; + sqlx::query( + r#" + INSERT INTO storage_objects ( + task_id, object_kind, state, + backend, storage_endpoint_id, object_key, storage_etag, + size_bytes, published_at + ) + SELECT + id, 'zip_attempt', 'published', + zip_storage_backend, zip_storage_endpoint_id, zip_storage_key, zip_storage_etag, + zip_size, NOW() + FROM tasks + WHERE id = $1 AND zip_storage_backend IS NOT NULL AND zip_storage_key IS NOT NULL + ON CONFLICT DO NOTHING + "#, + ) + .bind(task_id) + .execute(&mut **tx) + .await + .map_err(|err| { + AppError::new(ErrorCode::Internal, "持久化任务 ZIP 对象指针失败").with_source(err) + })?; + + let input_dir = PathBuf::from(&state.config.storage_path) + .join("orig") + .join(task_id.to_string()) + .to_string_lossy() + .to_string(); + let legacy_zip = PathBuf::from(&state.config.storage_path) + .join("zips") + .join(format!("{task_id}.zip")) + .to_string_lossy() + .to_string(); + for (kind, backend, key) in [ + ("legacy_zip", "local", legacy_zip), + ("input_dir", "local_dir", input_dir), + ] { + sqlx::query( + r#" + INSERT INTO storage_objects ( + task_id, object_kind, state, backend, object_key, published_at + ) VALUES ($1, $2, 'published', $3, $4, NOW()) + ON CONFLICT DO NOTHING + "#, + ) + .bind(task_id) + .bind(kind) + .bind(backend) + .bind(key) + .execute(&mut **tx) + .await + .map_err(|err| { + AppError::new(ErrorCode::Internal, "持久化本地清理目标失败").with_source(err) + })?; + } + Ok(()) +} + +async fn claim_deletions( + state: &AppState, + lease_owner: Uuid, + limit: i64, + task_id: Option, +) -> Result, AppError> { + sqlx::query_as::<_, DeleteClaim>( + r#" + WITH candidate AS ( + SELECT id + FROM storage_objects + WHERE state IN ('staging', 'delete_pending') + AND next_attempt_at <= NOW() + AND (lease_until IS NULL OR lease_until <= NOW()) + AND ($3::uuid IS NULL OR task_id = $3) + ORDER BY next_attempt_at, created_at + FOR UPDATE SKIP LOCKED + LIMIT $2 + ) + UPDATE storage_objects AS object + SET state = 'delete_pending', lease_owner = $1, + lease_until = NOW() + ($4 * INTERVAL '1 second'), + delete_attempts = object.delete_attempts + 1, updated_at = NOW() + FROM candidate + WHERE object.id = candidate.id + RETURNING object.id, object.task_id, object.backend, + object.storage_endpoint_id, object.object_key, object.delete_attempts + "#, + ) + .bind(lease_owner) + .bind(limit) + .bind(task_id) + .bind(DELETE_LEASE_SECONDS) + .fetch_all(&state.db) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "领取待删除对象失败").with_source(err)) +} + +async fn delete_claimed_object( + state: &AppState, + lease_owner: Uuid, + claim: DeleteClaim, +) -> Result<(), AppError> { + let result = if claim.backend == "local_dir" { + match tokio::fs::remove_dir_all(&claim.object_key).await { + Ok(()) => Ok(()), + Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(()), + Err(err) => Err( + AppError::new(ErrorCode::StorageUnavailable, "删除本地目录失败").with_source(err), + ), + } + } else { + storage::delete_object( + state, + &storage::ObjectLocator { + backend: claim.backend.clone(), + endpoint_id: claim.storage_endpoint_id, + key: claim.object_key.clone(), + }, + ) + .await + }; + + match result { + Ok(()) => { + sqlx::query( + r#" + UPDATE storage_objects + SET state = 'deleted', lease_owner = NULL, lease_until = NULL, + deleted_at = NOW(), last_error = NULL, updated_at = NOW() + WHERE id = $1 AND state = 'delete_pending' AND lease_owner = $2 + "#, + ) + .bind(claim.id) + .bind(lease_owner) + .execute(&state.db) + .await + .map_err(|err| { + AppError::new(ErrorCode::Internal, "完成对象删除状态失败").with_source(err) + })?; + } + Err(error) => { + let delay = delete_retry_delay(claim.delete_attempts); + sqlx::query( + r#" + UPDATE storage_objects + SET lease_owner = NULL, lease_until = NULL, + next_attempt_at = NOW() + ($3 * INTERVAL '1 second'), + last_error = $4, updated_at = NOW() + WHERE id = $1 AND state = 'delete_pending' AND lease_owner = $2 + "#, + ) + .bind(claim.id) + .bind(lease_owner) + .bind(delay.as_secs() as i64) + .bind(truncate_error(&error.to_string())) + .execute(&state.db) + .await + .map_err(|err| { + AppError::new(ErrorCode::Internal, "记录对象删除重试失败").with_source(err) + })?; + tracing::warn!(storage_object_id = %claim.id, task_id = %claim.task_id, attempts = claim.delete_attempts, retry_seconds = delay.as_secs(), error = %error, "storage object deletion deferred"); + } + } + Ok(()) +} + +fn delete_retry_delay(attempts: i32) -> Duration { + let exponent = attempts.saturating_sub(1).min(10) as u32; + Duration::from_secs(2_u64.saturating_pow(exponent).min(60 * 60)) +} + +fn truncate_error(value: &str) -> String { + value.chars().take(2_000).collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::Config; + use crate::services::mail::Mailer; + use crate::services::settings::{self, RuntimePolicyCache}; + use crate::services::storage::StorageCache; + use sqlx::postgres::PgPoolOptions; + use std::sync::Arc; + use tokio::sync::Semaphore; + + async fn test_state(database_url: String, redis_url: String, storage_path: String) -> AppState { + let mut config = Config::from_env().expect("load object lifecycle test config"); + config.database_url = database_url.clone(); + config.redis_url = redis_url; + config.storage_path = storage_path; + config.mail_enabled = false; + config.mail_log_links_when_disabled = false; + let db = PgPoolOptions::new() + .max_connections(16) + .connect(&database_url) + .await + .expect("connect object lifecycle test database"); + sqlx::migrate!().run(&db).await.expect("run migrations"); + let redis = redis::Client::open(config.redis_url.clone()) + .expect("create object lifecycle Redis client") + .get_connection_manager() + .await + .expect("connect object lifecycle Redis"); + 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(1)), + runtime_policy_cache: RuntimePolicyCache::new(), + storage_cache: StorageCache::new(), + config, + db, + redis, + } + } + + 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()?; + sqlx::query("UPDATE storage_endpoints SET is_active = false WHERE is_active = true") + .execute(&state.db) + .await + .expect("disable previous test storage endpoint"); + let endpoint_id = Uuid::new_v4(); + 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!("lifecycle-test-{marker}")) + .bind(endpoint) + .bind(bucket) + .bind(settings::encrypt_secret(state, &access_key).expect("encrypt test access key")) + .bind(settings::encrypt_secret(state, &secret_key).expect("encrypt test secret key")) + .execute(&state.db) + .await + .expect("insert lifecycle test S3 endpoint"); + state.storage_cache.invalidate(); + Some(endpoint_id) + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + #[ignore = "requires IMAGEFORGE_TEST_DATABASE_URL, IMAGEFORGE_TEST_REDIS_URL and optional IMAGEFORGE_TEST_S3_*"] + async fn staged_objects_and_task_deletion_are_durable_for_local_and_s3() { + let database_url = std::env::var("IMAGEFORGE_TEST_DATABASE_URL") + .expect("IMAGEFORGE_TEST_DATABASE_URL is required"); + let redis_url = std::env::var("IMAGEFORGE_TEST_REDIS_URL") + .expect("IMAGEFORGE_TEST_REDIS_URL is required"); + let marker = Uuid::new_v4().simple().to_string(); + let storage_root = std::env::temp_dir().join(format!("imageforge-lifecycle-{marker}")); + tokio::fs::create_dir_all(&storage_root) + .await + .expect("create lifecycle storage root"); + let state = test_state( + database_url, + redis_url, + storage_root.to_string_lossy().to_string(), + ) + .await; + sqlx::query("UPDATE storage_endpoints SET is_active = false WHERE is_active = true") + .execute(&state.db) + .await + .expect("disable S3 for local lifecycle case"); + state.storage_cache.invalidate(); + + let local_task = Uuid::new_v4(); + let local = store_tracked_bytes( + &state, + local_task, + None, + "result", + &format!("results/1d/test/{marker}-aborted-local.bin"), + Bytes::from_static(b"aborted-local"), + "application/octet-stream", + ) + .await + .expect("store staged local object"); + assert_eq!(local.stored.backend, "local"); + assert!(tokio::fs::try_exists(&local.stored.key) + .await + .expect("check staged local object")); + sqlx::query( + "UPDATE storage_objects SET lease_until = NOW() + INTERVAL '1 second' WHERE id = $1", + ) + .bind(local.lifecycle_id) + .execute(&state.db) + .await + .expect("shorten staged local lease"); + let heartbeat = start_staging_heartbeat_with_interval( + state.clone(), + local.lifecycle_id, + local.lease_owner, + Duration::from_millis(100), + 1, + ); + tokio::time::sleep(Duration::from_millis(1_500)).await; + let lease_is_current: bool = + sqlx::query_scalar("SELECT lease_until > NOW() FROM storage_objects WHERE id = $1") + .bind(local.lifecycle_id) + .fetch_one(&state.db) + .await + .expect("query renewed staged local lease"); + assert!(lease_is_current); + assert_eq!( + cleanup_ready_objects(&state, 10, Some(local_task)) + .await + .expect("skip actively uploaded local object"), + 0 + ); + drop(heartbeat); + sqlx::query( + "UPDATE storage_objects SET lease_until = NOW() - INTERVAL '1 second' WHERE id = $1", + ) + .bind(local.lifecycle_id) + .execute(&state.db) + .await + .expect("expire staged local object"); + cleanup_ready_objects(&state, 10, Some(local_task)) + .await + .expect("clean staged local object"); + assert!(!tokio::fs::try_exists(&local.stored.key) + .await + .expect("check cleaned local object")); + + let published_task = Uuid::new_v4(); + let published_file = Uuid::new_v4(); + let published = store_tracked_bytes( + &state, + published_task, + Some(published_file), + "result", + &format!("results/1d/test/{marker}-published-local.bin"), + Bytes::from_static(b"published-local"), + "application/octet-stream", + ) + .await + .expect("store publishable local object"); + let mut publish_tx = state.db.begin().await.expect("begin local publish tx"); + 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, 15, 15, NOW(), NOW(), NOW() + INTERVAL '1 day', 24) + "#, + ) + .bind(published_task) + .bind(format!("lifecycle-{marker}")) + .execute(&mut *publish_tx) + .await + .expect("insert published lifecycle task"); + sqlx::query( + r#" + INSERT INTO task_files ( + id, task_id, original_name, original_format, output_format, + original_size, compressed_size, storage_path, + storage_backend, storage_key, status, completed_at + ) VALUES ($1, $2, 'test.bin', 'png', 'png', 15, 15, $3, 'local', $3, 'completed', NOW()) + "#, + ) + .bind(published_file) + .bind(published_task) + .bind(&published.stored.key) + .execute(&mut *publish_tx) + .await + .expect("insert published lifecycle file"); + publish_in_tx(&mut publish_tx, &published) + .await + .expect("publish local lifecycle row"); + publish_tx.commit().await.expect("commit local publication"); + + let mut delete_tx = state.db.begin().await.expect("begin task deletion tx"); + let _: Uuid = sqlx::query_scalar("SELECT id FROM tasks WHERE id = $1 FOR UPDATE") + .bind(published_task) + .fetch_one(&mut *delete_tx) + .await + .expect("lock published task"); + assert!( + mark_task_deleting(&mut delete_tx, published_task, "test", false) + .await + .expect("mark published task deleting") + ); + delete_tx + .commit() + .await + .expect("commit task deletion marker"); + assert!(finalize_task_deletion(&state, published_task) + .await + .expect("finalize published task deletion")); + assert!(!tokio::fs::try_exists(&published.stored.key) + .await + .expect("check deleted published object")); + + let retry_dir = storage_root.join(format!("retry-{marker}")); + tokio::fs::create_dir_all(&retry_dir) + .await + .expect("create delete retry directory"); + let retry_id: Uuid = sqlx::query_scalar( + r#" + INSERT INTO storage_objects ( + task_id, object_kind, state, backend, object_key + ) VALUES ($1, 'result', 'delete_pending', 'local', $2) + RETURNING id + "#, + ) + .bind(Uuid::new_v4()) + .bind(retry_dir.to_string_lossy().to_string()) + .fetch_one(&state.db) + .await + .expect("insert deletion retry fixture"); + cleanup_ready_objects(&state, 1, None) + .await + .expect("run failing deletion attempt"); + let retry_state: (String, i32, bool) = sqlx::query_as( + "SELECT state, delete_attempts, last_error IS NOT NULL FROM storage_objects WHERE id = $1", + ) + .bind(retry_id) + .fetch_one(&state.db) + .await + .expect("query persisted deletion error"); + assert_eq!(retry_state, ("delete_pending".to_string(), 1, true)); + tokio::fs::remove_dir_all(&retry_dir) + .await + .expect("remove retry fixture directory"); + sqlx::query("UPDATE storage_objects SET next_attempt_at = NOW() WHERE id = $1") + .bind(retry_id) + .execute(&state.db) + .await + .expect("make deletion retry ready"); + cleanup_ready_objects(&state, 1, None) + .await + .expect("retry object deletion"); + let retried_state: String = + sqlx::query_scalar("SELECT state FROM storage_objects WHERE id = $1") + .bind(retry_id) + .fetch_one(&state.db) + .await + .expect("query successful deletion retry"); + assert_eq!(retried_state, "deleted"); + + let endpoint_id = configure_test_s3(&state, &marker).await; + if endpoint_id.is_some() { + let s3_task = Uuid::new_v4(); + let staged_s3 = store_tracked_bytes( + &state, + s3_task, + None, + "result", + &format!("results/1d/test/{marker}-aborted-s3.bin"), + Bytes::from_static(b"aborted-s3"), + "application/octet-stream", + ) + .await + .expect("store staged S3 object"); + assert_eq!(staged_s3.stored.backend, "s3"); + sqlx::query("UPDATE storage_objects SET lease_until = NOW() - INTERVAL '1 second' WHERE id = $1") + .bind(staged_s3.lifecycle_id) + .execute(&state.db) + .await + .expect("expire staged S3 object"); + cleanup_ready_objects(&state, 10, Some(s3_task)) + .await + .expect("clean staged S3 object"); + assert!(storage::read_bytes( + &state, + &storage::ObjectLocator { + backend: staged_s3.stored.backend, + endpoint_id: staged_s3.stored.endpoint_id, + key: staged_s3.stored.key, + }, + ) + .await + .is_err()); + } + + if let Some(endpoint_id) = endpoint_id { + sqlx::query("DELETE FROM storage_objects WHERE storage_endpoint_id = $1") + .bind(endpoint_id) + .execute(&state.db) + .await + .expect("clean S3 lifecycle rows"); + sqlx::query("DELETE FROM storage_endpoints WHERE id = $1") + .bind(endpoint_id) + .execute(&state.db) + .await + .expect("clean lifecycle S3 endpoint"); + } + sqlx::query("DELETE FROM storage_objects WHERE object_key LIKE $1") + .bind(format!("%{marker}%")) + .execute(&state.db) + .await + .expect("clean lifecycle object rows"); + let _ = tokio::fs::remove_dir_all(&storage_root).await; + } +} diff --git a/src/services/storage.rs b/src/services/storage.rs index dd90687..bc53486 100644 --- a/src/services/storage.rs +++ b/src/services/storage.rs @@ -330,27 +330,7 @@ fn retention_prefix(hours: i64) -> String { } } -pub async fn store_bytes( - state: &AppState, - key: &str, - bytes: B, - content_type: &str, -) -> Result -where - B: Into, -{ - let bytes = bytes.into(); - if let Some(endpoint) = active_endpoint(state).await? { - match store_bytes_s3(state, &endpoint, key, bytes.clone(), content_type).await { - Ok(stored) => return Ok(stored), - Err(err) => log_local_fallback(state, &endpoint, key, &err), - } - } - - store_bytes_local(state, key, bytes.as_ref()).await -} - -async fn store_bytes_s3( +pub(crate) async fn store_bytes_s3( state: &AppState, endpoint: &StorageEndpoint, key: &str, @@ -378,7 +358,7 @@ async fn store_bytes_s3( }) } -async fn store_bytes_local( +pub(crate) async fn store_bytes_local( state: &AppState, key: &str, bytes: &[u8], @@ -402,27 +382,7 @@ async fn store_bytes_local( }) } -pub async fn store_file( - state: &AppState, - key: &str, - path: &Path, - content_type: &str, -) -> Result { - let metadata = tokio::fs::metadata(path).await.map_err(|err| { - AppError::new(ErrorCode::StorageUnavailable, "读取待上传文件失败").with_source(err) - })?; - - if let Some(endpoint) = active_endpoint(state).await? { - match store_file_s3(state, &endpoint, key, path, content_type, metadata.len()).await { - Ok(stored) => return Ok(stored), - Err(err) => log_local_fallback(state, &endpoint, key, &err), - } - } - - store_file_local(state, key, path, metadata.len()).await -} - -async fn store_file_s3( +pub(crate) async fn store_file_s3( state: &AppState, endpoint: &StorageEndpoint, key: &str, @@ -459,7 +419,7 @@ async fn store_file_s3( }) } -async fn store_file_local( +pub(crate) async fn store_file_local( state: &AppState, key: &str, path: &Path, @@ -484,7 +444,12 @@ async fn store_file_local( }) } -fn log_local_fallback(state: &AppState, endpoint: &StorageEndpoint, key: &str, err: &AppError) { +pub(crate) fn log_local_fallback( + state: &AppState, + endpoint: &StorageEndpoint, + key: &str, + err: &AppError, +) { crate::services::metrics::record_storage_fallback(state); tracing::warn!( storage_endpoint_id = %endpoint.id, @@ -816,7 +781,7 @@ async fn endpoint_for_object( get_endpoint(state, endpoint_id).await } -fn local_path(state: &AppState, key: &str) -> Result { +pub(crate) fn local_path(state: &AppState, key: &str) -> Result { if key.is_empty() || key.starts_with('/') || key.starts_with('\\') diff --git a/src/services/task_queue.rs b/src/services/task_queue.rs new file mode 100644 index 0000000..668176d --- /dev/null +++ b/src/services/task_queue.rs @@ -0,0 +1,367 @@ +use crate::error::{AppError, ErrorCode}; +use crate::services::{metrics, object_lifecycle, quota}; +use crate::state::AppState; + +use chrono::Utc; +use sqlx::FromRow; +use std::time::Duration; +use uuid::Uuid; + +const DISPATCH_INTERVAL: Duration = Duration::from_secs(1); +const DISPATCH_BATCH_SIZE: i64 = 50; +const DELIVERY_LEASE_SECONDS: i64 = 30; +const MAX_DELIVERY_ATTEMPTS: i32 = 20; +const MAX_RETRY_SECONDS: u64 = 60; + +#[derive(Debug, FromRow)] +struct OutboxClaim { + task_id: Uuid, + attempts: i32, +} + +pub async fn dispatch_loop(state: AppState) { + let mut interval = tokio::time::interval(DISPATCH_INTERVAL); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + interval.tick().await; + if let Err(err) = dispatch_ready(&state, DISPATCH_BATCH_SIZE).await { + tracing::error!(error = %err, "task queue outbox dispatch iteration failed"); + } + } +} + +pub async fn dispatch_ready(state: &AppState, limit: i64) -> Result { + reconcile_non_pending_tasks(state).await?; + let lease_owner = Uuid::new_v4(); + let claims = claim_ready(state, lease_owner, limit.max(1), None).await?; + let count = claims.len(); + for claim in claims { + dispatch_claim(state, lease_owner, claim).await?; + } + Ok(count) +} + +pub async fn dispatch_task(state: &AppState, task_id: Uuid) -> Result { + let lease_owner = Uuid::new_v4(); + let mut claims = claim_ready(state, lease_owner, 1, Some(task_id)).await?; + let Some(claim) = claims.pop() else { + return Ok(false); + }; + dispatch_claim(state, lease_owner, claim).await?; + Ok(true) +} + +async fn claim_ready( + state: &AppState, + lease_owner: Uuid, + limit: i64, + task_id: Option, +) -> Result, AppError> { + sqlx::query_as::<_, OutboxClaim>( + r#" + WITH candidate AS ( + SELECT outbox.task_id + FROM task_queue_outbox AS outbox + JOIN tasks AS task ON task.id = outbox.task_id + WHERE outbox.status IN ('pending', 'delivering') + AND outbox.next_attempt_at <= NOW() + AND (outbox.lease_until IS NULL OR outbox.lease_until <= NOW()) + AND task.status = 'pending' + AND task.deletion_started_at IS NULL + AND ($3::uuid IS NULL OR outbox.task_id = $3) + ORDER BY outbox.next_attempt_at, outbox.created_at + FOR UPDATE OF outbox SKIP LOCKED + LIMIT $2 + ) + UPDATE task_queue_outbox AS outbox + SET status = 'delivering', + attempts = outbox.attempts + 1, + lease_owner = $1, + lease_until = NOW() + ($4 * INTERVAL '1 second'), + updated_at = NOW() + FROM candidate + WHERE outbox.task_id = candidate.task_id + RETURNING outbox.task_id, outbox.attempts + "#, + ) + .bind(lease_owner) + .bind(limit) + .bind(task_id) + .bind(DELIVERY_LEASE_SECONDS) + .fetch_all(&state.db) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "领取任务队列 outbox 失败").with_source(err)) +} + +async fn dispatch_claim( + state: &AppState, + lease_owner: Uuid, + claim: OutboxClaim, +) -> Result<(), AppError> { + match enqueue_task(state, claim.task_id).await { + Ok(()) => mark_delivered(state, claim.task_id, lease_owner).await, + Err(err) if claim.attempts >= MAX_DELIVERY_ATTEMPTS => { + dead_letter_pending_task(state, claim.task_id, lease_owner, &err).await + } + Err(err) => release_for_retry(state, claim, lease_owner, &err).await, + } +} + +pub(crate) async fn enqueue_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> { + let mut connection = state.redis.clone(); + redis::cmd("XADD") + .arg(metrics::QUEUE_STREAM_KEY) + .arg("MAXLEN") + .arg("~") + .arg(100_000) + .arg("*") + .arg("task_id") + .arg(task_id.to_string()) + .arg("created_at") + .arg(Utc::now().to_rfc3339()) + .query_async::<_, redis::Value>(&mut connection) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "写入队列失败").with_source(err))?; + Ok(()) +} + +async fn mark_delivered( + state: &AppState, + task_id: Uuid, + lease_owner: Uuid, +) -> Result<(), AppError> { + sqlx::query( + r#" + UPDATE task_queue_outbox + SET status = 'delivered', + delivered_at = COALESCE(delivered_at, NOW()), + lease_owner = NULL, + lease_until = NULL, + last_error = NULL, + updated_at = NOW() + WHERE task_id = $1 + AND status = 'delivering' + AND lease_owner = $2 + "#, + ) + .bind(task_id) + .bind(lease_owner) + .execute(&state.db) + .await + .map_err(|err| { + AppError::new(ErrorCode::Internal, "完成任务队列 outbox 失败").with_source(err) + })?; + Ok(()) +} + +async fn release_for_retry( + state: &AppState, + claim: OutboxClaim, + lease_owner: Uuid, + error: &AppError, +) -> Result<(), AppError> { + let delay = retry_delay(claim.attempts); + let message = truncate_error(error); + sqlx::query( + r#" + UPDATE task_queue_outbox + SET status = 'pending', + next_attempt_at = NOW() + ($3 * INTERVAL '1 second'), + lease_owner = NULL, + lease_until = NULL, + last_error = $4, + updated_at = NOW() + WHERE task_id = $1 + AND status = 'delivering' + AND lease_owner = $2 + "#, + ) + .bind(claim.task_id) + .bind(lease_owner) + .bind(delay.as_secs() as i64) + .bind(&message) + .execute(&state.db) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "记录任务队列重试失败").with_source(err))?; + tracing::warn!(task_id = %claim.task_id, attempts = claim.attempts, retry_seconds = delay.as_secs(), error = %error, "task queue delivery deferred"); + Ok(()) +} + +async fn dead_letter_pending_task( + state: &AppState, + task_id: Uuid, + lease_owner: Uuid, + error: &AppError, +) -> Result<(), AppError> { + let message = format!("队列持续不可用：{}", truncate_error(error)); + let input_dir = std::path::PathBuf::from(&state.config.storage_path) + .join("orig") + .join(task_id.to_string()) + .to_string_lossy() + .to_string(); + let mut tx = state.db.begin().await.map_err(|err| { + AppError::new(ErrorCode::Internal, "开启 outbox 死信事务失败").with_source(err) + })?; + let task: Option = + sqlx::query_scalar("SELECT status::text FROM tasks WHERE id = $1 FOR UPDATE") + .bind(task_id) + .fetch_optional(&mut *tx) + .await + .map_err(|err| { + AppError::new(ErrorCode::Internal, "锁定 outbox 死信任务失败").with_source(err) + })?; + let Some(status) = task else { + tx.rollback().await.ok(); + return Ok(()); + }; + let owned: Option = sqlx::query_scalar( + r#" + SELECT task_id + FROM task_queue_outbox + WHERE task_id = $1 + AND status = 'delivering' + AND lease_owner = $2 + FOR UPDATE + "#, + ) + .bind(task_id) + .bind(lease_owner) + .fetch_optional(&mut *tx) + .await + .map_err(|err| { + AppError::new(ErrorCode::Internal, "校验 outbox 死信租约失败").with_source(err) + })?; + if owned.is_none() { + tx.rollback().await.ok(); + return Ok(()); + } + + let failed = status == "pending"; + if failed { + sqlx::query( + r#" + UPDATE task_files + SET status = 'failed', + error_message = $2, + completed_at = NOW(), + input_path = NULL, + lease_owner = NULL, + lease_until = NULL + WHERE task_id = $1 + AND status IN ('pending', 'processing') + "#, + ) + .bind(task_id) + .bind(&message) + .execute(&mut *tx) + .await + .map_err(|err| { + AppError::new(ErrorCode::Internal, "收口 outbox 死信文件失败").with_source(err) + })?; + sqlx::query( + r#" + INSERT INTO storage_objects ( + task_id, object_kind, state, backend, object_key + ) VALUES ($1, 'input_dir', 'delete_pending', 'local_dir', $2) + ON CONFLICT DO NOTHING + "#, + ) + .bind(task_id) + .bind(&input_dir) + .execute(&mut *tx) + .await + .map_err(|err| { + AppError::new(ErrorCode::Internal, "安排 outbox 死信输入清理失败").with_source(err) + })?; + sqlx::query( + r#" + UPDATE tasks + SET status = 'failed', + completed_files = 0, + failed_files = total_files, + error_message = $2, + completed_at = NOW(), + lease_owner = NULL, + lease_until = NULL + WHERE id = $1 AND status = 'pending' + "#, + ) + .bind(task_id) + .bind(&message) + .execute(&mut *tx) + .await + .map_err(|err| { + AppError::new(ErrorCode::Internal, "收口 outbox 死信任务失败").with_source(err) + })?; + } + + sqlx::query( + r#" + UPDATE task_queue_outbox + SET status = CASE WHEN $3 THEN 'dead' ELSE 'delivered' END, + delivered_at = CASE WHEN $3 THEN delivered_at ELSE COALESCE(delivered_at, NOW()) END, + lease_owner = NULL, + lease_until = NULL, + last_error = $4, + updated_at = NOW() + WHERE task_id = $1 + AND lease_owner = $2 + "#, + ) + .bind(task_id) + .bind(lease_owner) + .bind(failed) + .bind(&message) + .execute(&mut *tx) + .await + .map_err(|err| { + AppError::new(ErrorCode::Internal, "提交 outbox 死信状态失败").with_source(err) + })?; + tx.commit().await.map_err(|err| { + AppError::new(ErrorCode::Internal, "提交 outbox 死信事务失败").with_source(err) + })?; + + if failed { + metrics::record_dead_letter(state); + if let Err(err) = object_lifecycle::cleanup_ready_objects(state, 10, Some(task_id)).await { + tracing::warn!(task_id = %task_id, error = %err, "dead outbox input cleanup deferred"); + } + quota::settle_anonymous_task_reservation(state, task_id).await?; + } + Ok(()) +} + +async fn reconcile_non_pending_tasks(state: &AppState) -> Result<(), AppError> { + sqlx::query( + r#" + UPDATE task_queue_outbox AS outbox + SET status = 'delivered', + delivered_at = COALESCE(outbox.delivered_at, NOW()), + lease_owner = NULL, + lease_until = NULL, + updated_at = NOW() + FROM tasks AS task + WHERE task.id = outbox.task_id + AND task.status <> 'pending' + AND outbox.status IN ('pending', 'delivering') + "#, + ) + .execute(&state.db) + .await + .map_err(|err| { + AppError::new(ErrorCode::Internal, "对账任务队列 outbox 失败").with_source(err) + })?; + Ok(()) +} + +fn retry_delay(attempts: i32) -> Duration { + let exponent = attempts.saturating_sub(1).min(6) as u32; + Duration::from_secs(2_u64.saturating_pow(exponent).min(MAX_RETRY_SECONDS)) +} + +fn truncate_error(error: &AppError) -> String { + format!("{}: {}", error.code.as_str(), error.message) + .chars() + .take(2_000) + .collect() +} diff --git a/src/worker/mod.rs b/src/worker/mod.rs index 634904a..be2eb9b 100644 --- a/src/worker/mod.rs +++ b/src/worker/mod.rs @@ -2,6 +2,7 @@ use crate::error::{AppError, ErrorCode}; use crate::services::billing; use crate::services::compress; use crate::services::metrics; +use crate::services::object_lifecycle; use crate::services::quota; use crate::services::storage; use crate::state::AppState; @@ -40,6 +41,10 @@ pub async fn run(state: AppState) -> Result<(), AppError> { let consumer = format!("worker_{worker_id}"); ensure_group(&state).await?; tokio::spawn(maintenance_loop(state.clone())); + tokio::spawn(crate::services::task_queue::dispatch_loop(state.clone())); + tokio::spawn(crate::services::object_lifecycle::maintenance_loop( + state.clone(), + )); let task_concurrency = state.config.worker_task_concurrency.max(1) as usize; let mut inflight = JoinSet::new(); @@ -675,21 +680,6 @@ struct TaskFileProcRow { output_format: String, } -#[derive(Debug, FromRow)] -struct CleanupFileRow { - storage_backend: String, - storage_endpoint_id: Option, - storage_key: Option, - input_path: Option, -} - -#[derive(Debug, FromRow)] -struct CleanupZipRow { - zip_storage_backend: Option, - zip_storage_endpoint_id: Option, - zip_storage_key: Option, -} - #[derive(Clone)] struct TaskContext { api_key_id: Option, @@ -729,6 +719,7 @@ pub(crate) async fn process_task( lease_owner = $2, lease_until = NOW() + $3 * INTERVAL '1 second' WHERE id = $1 + AND deletion_started_at IS NULL AND ( status = 'pending' OR ( @@ -940,6 +931,7 @@ async fn file_attempt_is_current(state: &AppState, fence: &FileFence) -> Result< JOIN task_files f ON f.task_id = t.id WHERE t.id = $1 AND t.status = 'processing' + AND t.deletion_started_at IS NULL AND t.processing_attempt = $2 AND t.lease_owner = $5 AND t.lease_until > NOW() @@ -1103,10 +1095,13 @@ async fn process_task_file( file_attempt, format_out.extension(), ); - let stored = match storage::store_bytes( + let tracked = match object_lifecycle::store_tracked_bytes( &state, + task_id, + Some(file.id), + "result", &object_key, - compressed, + compressed.into(), format_out.content_type(), ) .await @@ -1122,19 +1117,19 @@ async fn process_task_file( if ctx.is_anonymous && charge_units && !ctx.anonymous_quota_reserved { let Some(session_id) = ctx.session_id.as_deref() else { - let _ = storage::delete_object(&state, &stored_locator(&stored)).await; + discard_tracked_result(&state, &tracked, None).await; mark_file_failed_and_cleanup(&state, &fence, "匿名任务缺少 session_id", &input_path) .await?; return Ok(()); }; let Some(ip) = ctx.anon_ip else { - let _ = storage::delete_object(&state, &stored_locator(&stored)).await; + discard_tracked_result(&state, &tracked, None).await; mark_file_failed_and_cleanup(&state, &fence, "匿名任务缺少 client_ip", &input_path) .await?; return Ok(()); }; if let Err(err) = quota::consume_anonymous_units(&state, session_id, ip, 1).await { - let _ = storage::delete_object(&state, &stored_locator(&stored)).await; + discard_tracked_result(&state, &tracked, Some(&err)).await; mark_file_failed_and_cleanup(&state, &fence, &err.message, &input_path).await?; return Ok(()); } @@ -1146,7 +1141,7 @@ async fn process_task_file( ctx.api_key_id, &ctx.source, &fence, - &stored, + &tracked, original_size as i64, compressed_size as i64, saved_percent, @@ -1160,16 +1155,85 @@ async fn process_task_file( let _ = tokio::fs::remove_file(&input_path).await; } Ok(FinalizeFileOutcome::LeaseLost) => { - let _ = storage::delete_object(&state, &stored_locator(&stored)).await; - } - Err(err) => { - let _ = storage::delete_object(&state, &stored_locator(&stored)).await; - mark_file_failed_and_cleanup(&state, &fence, &err.message, &input_path).await?; + discard_tracked_result(&state, &tracked, None).await; } + Err(err) => match worker_result_was_committed(&state, &fence, &tracked).await { + Ok(true) => { + tracing::warn!(task_id = %task_id, file_id = %fence.file_id, error = %err, "worker result commit response was lost; recovered committed publication"); + let _ = tokio::fs::remove_file(&input_path).await; + } + Ok(false) => { + discard_tracked_result(&state, &tracked, Some(&err)).await; + mark_file_failed_and_cleanup(&state, &fence, &err.message, &input_path).await?; + } + Err(probe_err) => { + tracing::error!(task_id = %task_id, file_id = %fence.file_id, error = %probe_err, original_error = %err, "worker result commit state is unknown; staging lease will reconcile object"); + return Err(err); + } + }, } Ok(()) } +async fn discard_tracked_result( + state: &AppState, + tracked: &object_lifecycle::TrackedStoredObject, + error: Option<&AppError>, +) { + if let Err(schedule_err) = + object_lifecycle::schedule_tracked_delete(state, tracked, error).await + { + tracing::error!(storage_object_id = %tracked.lifecycle_id, error = %schedule_err, "failed to persist discarded worker object cleanup"); + return; + } + if let Err(cleanup_err) = + object_lifecycle::cleanup_ready_objects(state, 1, Some(tracked.task_id)).await + { + tracing::warn!(storage_object_id = %tracked.lifecycle_id, error = %cleanup_err, "discarded worker object cleanup deferred"); + } +} + +async fn worker_result_was_committed( + state: &AppState, + fence: &FileFence, + tracked: &object_lifecycle::TrackedStoredObject, +) -> Result { + sqlx::query_scalar( + r#" + SELECT EXISTS( + SELECT 1 + FROM tasks AS task + JOIN task_files AS file ON file.task_id = task.id + JOIN storage_objects AS object ON object.id = $3 + WHERE task.id = $1 + AND file.id = $2 + AND file.status = 'completed' + AND file.storage_backend = $4 + AND file.storage_endpoint_id IS NOT DISTINCT FROM $5 + AND COALESCE(file.storage_key, file.storage_path) = $6 + AND object.state = 'published' + AND object.task_id = task.id + AND object.task_file_id = file.id + ) + "#, + ) + .bind(fence.task_id) + .bind(fence.file_id) + .bind(tracked.lifecycle_id) + .bind(&tracked.stored.backend) + .bind(tracked.stored.endpoint_id) + .bind(&tracked.stored.key) + .fetch_one(&state.db) + .await + .map_err(|err| { + AppError::new( + ErrorCode::StorageUnavailable, + "核验 Worker 结果提交状态失败", + ) + .with_source(err) + }) +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum FinalizeFileOutcome { Committed, @@ -1183,7 +1247,7 @@ async fn finalize_file( api_key_id: Option, source: &str, fence: &FileFence, - stored: &storage::StoredObject, + tracked: &object_lifecycle::TrackedStoredObject, bytes_in: i64, bytes_out: i64, saved_percent: f64, @@ -1191,6 +1255,7 @@ async fn finalize_file( format_out: compress::ImageFmt, charge_units: bool, ) -> Result { + let stored = &tracked.stored; let mut tx = state .db .begin() @@ -1332,6 +1397,8 @@ async fn finalize_file( return Ok(FinalizeFileOutcome::LeaseLost); } + object_lifecycle::publish_in_tx(&mut tx, tracked).await?; + tx.commit() .await .map_err(|err| AppError::new(ErrorCode::Internal, "提交事务失败").with_source(err))?; @@ -1840,6 +1907,12 @@ async fn cleanup_expired_records(state: &AppState) -> Result<(), AppError> { .execute(&state.db) .await; + let _ = sqlx::query( + "DELETE FROM storage_objects WHERE state = 'deleted' AND deleted_at < NOW() - INTERVAL '7 days'", + ) + .execute(&state.db) + .await; + let _ = sqlx::query("DELETE FROM webhook_events WHERE received_at < NOW() - INTERVAL '90 days'") .execute(&state.db) @@ -1851,6 +1924,7 @@ async fn cleanup_expired_records(state: &AppState) -> Result<(), AppError> { WHERE e.deleted_at < NOW() - INTERVAL '30 days' AND NOT EXISTS (SELECT 1 FROM task_files f WHERE f.storage_endpoint_id = e.id) AND NOT EXISTS (SELECT 1 FROM tasks t WHERE t.zip_storage_endpoint_id = e.id) + AND NOT EXISTS (SELECT 1 FROM storage_objects o WHERE o.storage_endpoint_id = e.id) "#, ) .execute(&state.db) @@ -1894,82 +1968,13 @@ async fn cleanup_expired_tasks(state: &AppState) -> Result<(), AppError> { } async fn cleanup_expired_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> { - sqlx::query( - "UPDATE tasks SET status = 'cancelled', completed_at = COALESCE(completed_at, NOW()) WHERE id = $1 AND expires_at < NOW() AND status IN ('pending', 'processing')", - ) - .bind(task_id) - .execute(&state.db) - .await - .map_err(|err| AppError::new(ErrorCode::Internal, "终止过期任务失败").with_source(err))?; - quota::settle_anonymous_task_reservation(state, task_id).await?; - - let files: Vec = sqlx::query_as( - r#" - SELECT storage_backend, storage_endpoint_id, - COALESCE(storage_key, storage_path) AS storage_key, - input_path - FROM task_files - WHERE task_id = $1 - "#, - ) - .bind(task_id) - .fetch_all(&state.db) - .await - .map_err(|err| AppError::new(ErrorCode::Internal, "查询过期任务文件失败").with_source(err))?; - - for file in files { - if let Some(key) = file.storage_key { - storage::delete_object( - state, - &storage::ObjectLocator { - backend: file.storage_backend, - endpoint_id: file.storage_endpoint_id, - key, - }, - ) - .await?; - } - if let Some(input_path) = file.input_path { - let _ = tokio::fs::remove_file(input_path).await; - } + if object_lifecycle::mark_expired_task(state, task_id).await? { + object_lifecycle::finalize_task_deletion(state, task_id).await?; } - - let zip: 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))?; - if let Some(zip) = zip { - if let (Some(backend), Some(key)) = (zip.zip_storage_backend, zip.zip_storage_key) { - storage::delete_object( - state, - &storage::ObjectLocator { - backend, - endpoint_id: zip.zip_storage_endpoint_id, - key, - }, - ) - .await?; - } - } - - let legacy_zip_path = format!("{}/zips/{task_id}.zip", state.config.storage_path); - let _ = tokio::fs::remove_file(legacy_zip_path).await; - let orig_dir = format!("{}/orig/{task_id}", state.config.storage_path); - let _ = tokio::fs::remove_dir_all(orig_dir).await; - - sqlx::query("DELETE FROM tasks WHERE id = $1 AND expires_at < NOW()") - .bind(task_id) - .execute(&state.db) - .await - .map_err(|err| { - AppError::new(ErrorCode::Internal, "删除过期任务记录失败").with_source(err) - })?; Ok(()) } +#[cfg(test)] fn stored_locator(stored: &storage::StoredObject) -> storage::ObjectLocator { storage::ObjectLocator { backend: stored.backend.clone(), @@ -2148,16 +2153,22 @@ mod tests { let stale_key = storage::result_attempt_key(24, task_id, file_id, 1, 1, "png"); let winning_key = storage::result_attempt_key(24, task_id, file_id, 2, 2, "png"); - let stale_object = storage::store_bytes( + let stale_object = object_lifecycle::store_tracked_bytes( &state, + task_id, + Some(file_id), + "result", &stale_key, Bytes::from_static(b"stale-attempt"), "image/png", ) .await .expect("store stale attempt object"); - let winning_object = storage::store_bytes( + let winning_object = object_lifecycle::store_tracked_bytes( &state, + task_id, + Some(file_id), + "result", &winning_key, Bytes::from_static(b"winning-attempt"), "image/png", @@ -2165,8 +2176,8 @@ mod tests { .await .expect("store winning attempt object"); if let Ok(expected_backend) = std::env::var("IMAGEFORGE_TEST_EXPECT_STORAGE_BACKEND") { - assert_eq!(stale_object.backend, expected_backend); - assert_eq!(winning_object.backend, expected_backend); + assert_eq!(stale_object.stored.backend, expected_backend); + assert_eq!(winning_object.stored.backend, expected_backend); } let period_start = Utc::now() - chrono::Duration::hours(1); @@ -2259,14 +2270,14 @@ mod tests { assert_eq!(stale_result, FinalizeFileOutcome::LeaseLost); assert_eq!(winning_result, FinalizeFileOutcome::Committed); - storage::delete_object(&state, &stored_locator(&stale_object)) - .await - .expect("delete stale attempt object"); - assert!(storage::read_bytes(&state, &stored_locator(&stale_object)) - .await - .is_err()); + discard_tracked_result(&state, &stale_object, None).await; + assert!( + storage::read_bytes(&state, &stored_locator(&stale_object.stored)) + .await + .is_err() + ); assert_eq!( - storage::read_bytes(&state, &stored_locator(&winning_object)) + storage::read_bytes(&state, &stored_locator(&winning_object.stored)) .await .expect("read winning object"), b"winning-attempt" @@ -2292,7 +2303,11 @@ mod tests { .expect("query test file"); assert_eq!( file, - ("completed".to_string(), winning_object.key.clone(), 40) + ( + "completed".to_string(), + winning_object.stored.key.clone(), + 40 + ) ); let usage_event_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM usage_events WHERE task_file_id = $1") @@ -2312,9 +2327,7 @@ mod tests { .expect("query used units"); assert_eq!(used_units, 1); - storage::delete_object(&state, &stored_locator(&winning_object)) - .await - .expect("delete winning object"); + discard_tracked_result(&state, &winning_object, None).await; sqlx::query("DELETE FROM usage_events WHERE task_id = $1") .bind(task_id) .execute(&pool) @@ -2325,6 +2338,11 @@ mod tests { .execute(&pool) .await .expect("delete test task"); + sqlx::query("DELETE FROM storage_objects WHERE task_id = $1") + .bind(task_id) + .execute(&pool) + .await + .expect("delete test storage lifecycle rows"); sqlx::query("DELETE FROM usage_periods WHERE user_id = $1") .bind(user_id) .execute(&pool)