From 57892f6509018f9c3aab905e8cacfe6f8cb04409 Mon Sep 17 00:00:00 2001 From: 237899745 <237899745@users.noreply.git.workyai.cn> Date: Sun, 26 Jul 2026 02:03:57 +0800 Subject: [PATCH] fix: fence worker retries and object publication --- migrations/015_worker_attempt_leases.sql | 17 + src/api/admin.rs | 56 +- src/services/storage.rs | 22 + src/worker/mod.rs | 1288 ++++++++++++++++++---- 4 files changed, 1146 insertions(+), 237 deletions(-) create mode 100644 migrations/015_worker_attempt_leases.sql diff --git a/migrations/015_worker_attempt_leases.sql b/migrations/015_worker_attempt_leases.sql new file mode 100644 index 0000000..db0061a --- /dev/null +++ b/migrations/015_worker_attempt_leases.sql @@ -0,0 +1,17 @@ +ALTER TABLE tasks + ADD COLUMN IF NOT EXISTS processing_attempt BIGINT NOT NULL DEFAULT 0, + ADD COLUMN IF NOT EXISTS lease_owner UUID, + ADD COLUMN IF NOT EXISTS lease_until TIMESTAMPTZ; + +ALTER TABLE task_files + ADD COLUMN IF NOT EXISTS processing_attempt BIGINT NOT NULL DEFAULT 0, + ADD COLUMN IF NOT EXISTS lease_owner UUID, + ADD COLUMN IF NOT EXISTS lease_until TIMESTAMPTZ; + +CREATE INDEX IF NOT EXISTS idx_tasks_processing_lease + ON tasks(lease_until) + WHERE status = 'processing'; + +CREATE INDEX IF NOT EXISTS idx_task_files_processing_lease + ON task_files(task_id, lease_until) + WHERE status = 'processing'; diff --git a/src/api/admin.rs b/src/api/admin.rs index eda5b1c..fc30b31 100644 --- a/src/api/admin.rs +++ b/src/api/admin.rs @@ -556,28 +556,14 @@ async fn cancel_task( .await .map_err(|err| AppError::new(ErrorCode::Internal, "开启事务失败").with_source(err))?; - sqlx::query( - r#" - UPDATE task_files - SET status = 'failed', - error_message = '任务已取消', - completed_at = NOW() - WHERE task_id = $1 AND status IN ('pending', 'processing') - "#, - ) - .bind(task_id) - .execute(&mut *tx) - .await - .map_err(|err| AppError::new(ErrorCode::Internal, "更新任务文件失败").with_source(err))?; - let updated = sqlx::query( r#" UPDATE tasks SET status = 'cancelled', error_message = '管理员取消任务', completed_at = NOW(), - completed_files = (SELECT COUNT(*) FROM task_files WHERE task_id = $1 AND status = 'completed'), - failed_files = (SELECT COUNT(*) FROM task_files WHERE task_id = $1 AND status = 'failed') + lease_owner = NULL, + lease_until = NULL WHERE id = $1 AND status IN ('pending', 'processing') "#, ) @@ -590,6 +576,44 @@ async fn cancel_task( return Err(AppError::new(ErrorCode::InvalidRequest, "任务状态无法取消")); } + sqlx::query( + r#" + UPDATE task_files + SET status = 'failed', + error_message = '任务已取消', + completed_at = NOW(), + lease_owner = NULL, + lease_until = NULL + WHERE task_id = $1 AND status IN ('pending', 'processing') + "#, + ) + .bind(task_id) + .execute(&mut *tx) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "更新任务文件失败").with_source(err))?; + + sqlx::query( + r#" + UPDATE tasks + SET completed_files = ( + SELECT COUNT(*) FROM task_files WHERE task_id = $1 AND status = 'completed' + ), + failed_files = ( + SELECT COUNT(*) FROM task_files WHERE task_id = $1 AND status = 'failed' + ), + total_compressed_size = COALESCE(( + SELECT SUM(compressed_size) + FROM task_files + WHERE task_id = $1 AND status = 'completed' + ), 0)::bigint + WHERE id = $1 AND status = 'cancelled' + "#, + ) + .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))?; diff --git a/src/services/storage.rs b/src/services/storage.rs index a8fca68..fa5a667 100644 --- a/src/services/storage.rs +++ b/src/services/storage.rs @@ -291,6 +291,26 @@ pub fn result_key(retention_hours: i64, task_id: Uuid, file_id: Uuid, extension: ) } +pub fn result_attempt_key( + retention_hours: i64, + task_id: Uuid, + file_id: Uuid, + task_attempt: i64, + file_attempt: i64, + extension: &str, +) -> String { + let now = Utc::now(); + format!( + "results/{}/{:04}/{:02}/{task_id}/{file_id}-t{}-f{}.{}", + retention_prefix(retention_hours), + now.year(), + now.month(), + task_attempt.max(1), + file_attempt.max(1), + extension.trim_start_matches('.') + ) +} + pub fn archive_key(retention_hours: i64, task_id: Uuid) -> String { let now = Utc::now(); format!( @@ -964,6 +984,8 @@ mod tests { assert!(key.starts_with("results/7d/")); assert!(key.ends_with("/00000000-0000-0000-0000-000000000001.webp")); assert!(archive_key(360, task_id).starts_with("archives/15d/")); + let attempt_key = result_attempt_key(24, task_id, file_id, 2, 3, "avif"); + assert!(attempt_key.contains("-t2-f3.avif")); } #[test] diff --git a/src/worker/mod.rs b/src/worker/mod.rs index 475c8bc..d04033d 100644 --- a/src/worker/mod.rs +++ b/src/worker/mod.rs @@ -21,7 +21,9 @@ const GROUP_NAME: &str = metrics::QUEUE_GROUP_NAME; const DEAD_STREAM_KEY: &str = metrics::DEAD_STREAM_KEY; const MAX_DELIVERIES: usize = 3; const STALE_MESSAGE_IDLE_MS: usize = 5 * 60 * 1000; -const MESSAGE_HEARTBEAT_SECONDS: u64 = 60; +const MESSAGE_HEARTBEAT_SECONDS: u64 = 30; +const PROCESSING_LEASE_SECONDS: i64 = 120; +const LEASE_BUSY_RETRY_SECONDS: u64 = 5; const QUEUE_BLOCK_MS: usize = 1_000; const MAINTENANCE_INTERVAL_SECONDS: u64 = 300; const MAINTENANCE_BATCH_SIZE: i64 = 1_000; @@ -34,7 +36,8 @@ pub async fn run(state: AppState) -> Result<(), AppError> { crate::services::bootstrap::ensure_schema(&state).await?; - let consumer = format!("worker_{}", Uuid::new_v4()); + let worker_id = Uuid::new_v4(); + let consumer = format!("worker_{worker_id}"); ensure_group(&state).await?; tokio::spawn(maintenance_loop(state.clone())); @@ -56,8 +59,12 @@ pub async fn run(state: AppState) -> Result<(), AppError> { for message in messages { let state = state.clone(); let consumer = consumer.clone(); + // A process can receive duplicate messages for the same task. Keep their + // database leases distinct even though they share one Redis consumer. + let lease_owner = Uuid::new_v4(); inflight.spawn(async move { - process_message_with_retries(state, consumer, message).await + process_message_with_retries(state, lease_owner, consumer, message) + .await }); } } @@ -166,12 +173,17 @@ async fn read_messages( async fn process_message_with_retries( state: AppState, + worker_id: Uuid, consumer: String, mut message: StreamId, ) -> Result<(), AppError> { loop { - match handle_message_with_heartbeat(&state, &consumer, &message).await { - Ok(()) => return Ok(()), + match handle_message_with_heartbeat(&state, worker_id, &consumer, &message).await { + Ok(MessageOutcome::Done | MessageOutcome::OwnershipLost) => return Ok(()), + Ok(MessageOutcome::LeaseBusy) => { + tokio::time::sleep(Duration::from_secs(LEASE_BUSY_RETRY_SECONDS)).await; + continue; + } Err(err) => { let mut conn = state.redis.clone(); let deliveries = match pending_delivery_count(&mut conn, &message.id).await { @@ -222,12 +234,16 @@ async fn process_message_with_retries( async fn handle_message_with_heartbeat( state: &AppState, + worker_id: Uuid, consumer: &str, message: &StreamId, -) -> Result<(), AppError> { +) -> Result { let mut conn = state.redis.clone(); - let handling = handle_message(state, &mut conn, message); + let handling = handle_message(state, worker_id, &mut conn, message); tokio::pin!(handling); + let task_id = message + .get::("task_id") + .and_then(|value| Uuid::parse_str(&value).ok()); let mut heartbeat = tokio::time::interval_at( tokio::time::Instant::now() + Duration::from_secs(MESSAGE_HEARTBEAT_SECONDS), Duration::from_secs(MESSAGE_HEARTBEAT_SECONDS), @@ -238,12 +254,27 @@ async fn handle_message_with_heartbeat( tokio::select! { result = &mut handling => return result, _ = heartbeat.tick() => { - if let Err(err) = touch_pending_message(state, consumer, &message.id).await { - tracing::warn!( - message_id = %message.id, - error = ?err, - "failed to refresh worker message heartbeat" - ); + match touch_pending_message(state, consumer, &message.id).await { + Ok(true) => {} + Ok(false) => { + if let Some(task_id) = task_id { + release_processing_lease(state, task_id, worker_id).await; + } + return Ok(MessageOutcome::OwnershipLost); + } + Err(err) => { + tracing::warn!( + message_id = %message.id, + error = ?err, + "failed to refresh worker message heartbeat" + ); + continue; + } + } + if let Some(task_id) = task_id { + if !renew_processing_lease(state, task_id, worker_id).await? { + return Ok(MessageOutcome::OwnershipLost); + } } } } @@ -254,47 +285,68 @@ async fn touch_pending_message( state: &AppState, consumer: &str, message_id: &str, -) -> Result<(), AppError> { +) -> Result { let mut conn = state.redis.clone(); - redis::cmd("XCLAIM") - .arg(STREAM_KEY) + let script = redis::Script::new( + r#" + local pending = redis.call('XPENDING', KEYS[1], ARGV[1], ARGV[2], ARGV[2], 1) + if #pending == 0 then return 0 end + if pending[1][2] ~= ARGV[3] then return -1 end + redis.call('XCLAIM', KEYS[1], ARGV[1], ARGV[3], 0, ARGV[2], 'JUSTID') + return 1 + "#, + ); + let owned: i64 = script + .key(STREAM_KEY) .arg(GROUP_NAME) - .arg(consumer) - .arg(0) .arg(message_id) - .arg("JUSTID") - .query_async::<_, redis::Value>(&mut conn) + .arg(consumer) + .invoke_async(&mut conn) .await .map_err(|err| { AppError::new(ErrorCode::Internal, "刷新队列任务心跳失败").with_source(err) })?; - Ok(()) + Ok(owned == 1) +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum MessageOutcome { + Done, + LeaseBusy, + OwnershipLost, } async fn handle_message( state: &AppState, + worker_id: Uuid, conn: &mut redis::aio::ConnectionManager, msg: &StreamId, -) -> Result<(), AppError> { +) -> Result { let Some(task_id_str) = msg.get::("task_id") else { ack_message(conn, &msg.id).await?; - return Ok(()); + return Ok(MessageOutcome::Done); }; let task_id = match Uuid::parse_str(&task_id_str) { Ok(value) => value, Err(_) => { ack_message(conn, &msg.id).await?; - return Ok(()); + return Ok(MessageOutcome::Done); } }; - match process_task(state, task_id).await { - Ok(()) => ack_message(conn, &msg.id).await, + match process_task(state, task_id, worker_id).await { + Ok(TaskProcessOutcome::Done) => { + ack_message(conn, &msg.id).await?; + Ok(MessageOutcome::Done) + } + Ok(TaskProcessOutcome::LeaseBusy) => Ok(MessageOutcome::LeaseBusy), Err(err) => { let deliveries = pending_delivery_count(conn, &msg.id).await?; if should_dead_letter(deliveries) { + if !mark_task_dead_letter(state, task_id, worker_id, &err.message).await? { + return Ok(MessageOutcome::OwnershipLost); + } write_dead_letter(conn, &msg.id, task_id, deliveries, &err).await?; - mark_task_dead_letter(state, task_id, &err.message).await?; ack_message(conn, &msg.id).await?; metrics::record_dead_letter(state); tracing::error!( @@ -303,7 +355,7 @@ async fn handle_message( error = %err, "task moved to dead-letter stream" ); - return Ok(()); + return Ok(MessageOutcome::Done); } Err(err) @@ -428,19 +480,44 @@ async fn write_dead_letter( async fn mark_task_dead_letter( state: &AppState, task_id: Uuid, + worker_id: Uuid, message: &str, -) -> Result<(), AppError> { +) -> Result { let mut tx = state.db.begin().await.map_err(|err| { AppError::new(ErrorCode::Internal, "开启死信事务失败").with_source(err) })?; + let owned_task: Option = sqlx::query_scalar( + r#" + SELECT id + FROM tasks + WHERE id = $1 + AND status = 'processing' + AND lease_owner = $2 + AND lease_until > NOW() + FOR UPDATE + "#, + ) + .bind(task_id) + .bind(worker_id) + .fetch_optional(&mut *tx) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "锁定死信任务失败").with_source(err))?; + if owned_task.is_none() { + tx.rollback().await.ok(); + return Ok(false); + } + sqlx::query( r#" UPDATE task_files SET status = 'failed', error_message = $2, - completed_at = NOW() - WHERE task_id = $1 AND status IN ('pending', 'processing') + completed_at = NOW(), + lease_owner = NULL, + lease_until = NULL + WHERE task_id = $1 + AND status IN ('pending', 'processing') "#, ) .bind(task_id) @@ -452,14 +529,27 @@ async fn mark_task_dead_letter( r#" UPDATE tasks SET status = 'failed', - failed_files = GREATEST(total_files - completed_files, failed_files), + completed_files = ( + SELECT COUNT(*) FROM task_files WHERE task_id = $1 AND status = 'completed' + ), + failed_files = ( + SELECT COUNT(*) FROM task_files WHERE task_id = $1 AND status = 'failed' + ), + total_compressed_size = COALESCE(( + SELECT SUM(compressed_size) + FROM task_files + WHERE task_id = $1 AND status = 'completed' + ), 0)::bigint, error_message = $2, - completed_at = NOW() - WHERE id = $1 AND status NOT IN ('completed', 'failed', 'cancelled') + completed_at = NOW(), + lease_owner = NULL, + lease_until = NULL + WHERE id = $1 AND status = 'processing' AND lease_owner = $3 "#, ) .bind(task_id) .bind(message) + .bind(worker_id) .execute(&mut *tx) .await .map_err(|err| AppError::new(ErrorCode::Internal, "标记死信任务失败").with_source(err))?; @@ -467,7 +557,83 @@ async fn mark_task_dead_letter( .await .map_err(|err| AppError::new(ErrorCode::Internal, "提交死信事务失败").with_source(err))?; quota::settle_anonymous_task_reservation(state, task_id).await?; - Ok(()) + Ok(true) +} + +async fn renew_processing_lease( + state: &AppState, + task_id: Uuid, + worker_id: Uuid, +) -> Result { + let mut tx = + state.db.begin().await.map_err(|err| { + AppError::new(ErrorCode::Internal, "开启续租事务失败").with_source(err) + })?; + let renewed = sqlx::query( + r#" + UPDATE tasks + SET lease_until = NOW() + $3 * INTERVAL '1 second' + WHERE id = $1 + AND status = 'processing' + AND lease_owner = $2 + AND lease_until > NOW() + "#, + ) + .bind(task_id) + .bind(worker_id) + .bind(PROCESSING_LEASE_SECONDS) + .execute(&mut *tx) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "续租任务失败").with_source(err))?; + if renewed.rows_affected() == 0 { + tx.rollback().await.ok(); + return Ok(false); + } + sqlx::query( + r#" + UPDATE task_files + SET lease_until = NOW() + $3 * INTERVAL '1 second' + WHERE task_id = $1 + AND status = 'processing' + AND lease_owner = $2 + AND lease_until > NOW() + "#, + ) + .bind(task_id) + .bind(worker_id) + .bind(PROCESSING_LEASE_SECONDS) + .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))?; + Ok(true) +} + +async fn release_processing_lease(state: &AppState, task_id: Uuid, worker_id: Uuid) { + let result = async { + let mut tx = state.db.begin().await?; + sqlx::query( + "UPDATE tasks SET lease_until = NOW() WHERE id = $1 AND status = 'processing' AND lease_owner = $2", + ) + .bind(task_id) + .bind(worker_id) + .execute(&mut *tx) + .await?; + sqlx::query( + "UPDATE task_files SET lease_until = NOW() WHERE task_id = $1 AND status = 'processing' AND lease_owner = $2", + ) + .bind(task_id) + .bind(worker_id) + .execute(&mut *tx) + .await?; + tx.commit().await + } + .await; + if let Err(err) = result { + tracing::warn!(task_id = %task_id, worker_id = %worker_id, error = ?err, "failed to release processing lease"); + } } async fn ack_message( @@ -486,7 +652,6 @@ async fn ack_message( #[derive(Debug, FromRow)] struct TaskProcRow { - status: String, compression_level: String, compression_rate: Option, max_width: Option, @@ -499,6 +664,7 @@ struct TaskProcRow { client_ip: Option, retention_hours: i32, anonymous_units_reserved: i32, + processing_attempt: i64, } #[derive(Debug, FromRow)] @@ -507,7 +673,6 @@ struct TaskFileProcRow { input_path: Option, original_format: String, output_format: String, - status: String, } #[derive(Debug, FromRow)] @@ -537,11 +702,46 @@ struct TaskContext { anonymous_quota_reserved: bool, } -async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> { - let Some(mut task): Option = sqlx::query_as( +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum TaskProcessOutcome { + Done, + LeaseBusy, +} + +async fn process_task( + state: &AppState, + task_id: Uuid, + worker_id: Uuid, +) -> Result { + let task: Option = sqlx::query_as( r#" - SELECT - status::text AS status, + UPDATE tasks + SET status = 'processing', + started_at = COALESCE(started_at, NOW()), + processing_attempt = CASE + WHEN status = 'pending' + OR lease_owner IS DISTINCT FROM $2 + OR lease_until IS NULL + OR lease_until <= NOW() + THEN processing_attempt + 1 + ELSE processing_attempt + END, + lease_owner = $2, + lease_until = NOW() + $3 * INTERVAL '1 second' + WHERE id = $1 + AND ( + status = 'pending' + OR ( + status = 'processing' + AND ( + lease_owner = $2 + OR lease_owner IS NULL + OR lease_until IS NULL + OR lease_until <= NOW() + ) + ) + ) + RETURNING compression_level::text AS compression_level, compression_rate, max_width, @@ -553,61 +753,43 @@ async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> { source::text AS source, host(client_ip) AS client_ip, retention_hours, - anonymous_units_reserved - FROM tasks - WHERE id = $1 + anonymous_units_reserved, + processing_attempt "#, ) .bind(task_id) + .bind(worker_id) + .bind(PROCESSING_LEASE_SECONDS) .fetch_optional(&state.db) .await - .map_err(|err| AppError::new(ErrorCode::Internal, "查询任务失败").with_source(err))? - else { - tracing::info!(task_id = %task_id, "task was deleted before processing; acknowledging message"); - return Ok(()); + .map_err(|err| AppError::new(ErrorCode::Internal, "领取任务失败").with_source(err))?; + let Some(task) = task else { + let status: Option = + sqlx::query_scalar("SELECT status::text FROM tasks WHERE id = $1") + .bind(task_id) + .fetch_optional(&state.db) + .await + .map_err(|err| { + AppError::new(ErrorCode::Internal, "查询任务状态失败").with_source(err) + })?; + match status.as_deref() { + None => { + tracing::info!(task_id = %task_id, "task was deleted before processing; acknowledging message"); + return Ok(TaskProcessOutcome::Done); + } + Some("cancelled") => { + finalize_cancelled_task(state, task_id).await?; + return Ok(TaskProcessOutcome::Done); + } + Some("completed" | "failed") => { + quota::settle_anonymous_task_reservation(state, task_id).await?; + return Ok(TaskProcessOutcome::Done); + } + Some("processing") => return Ok(TaskProcessOutcome::LeaseBusy), + Some(_) => return Ok(TaskProcessOutcome::LeaseBusy), + } }; - if task.status == "cancelled" { - finalize_task_status(state, task_id).await?; - return Ok(()); - } - if matches!(task.status.as_str(), "completed" | "failed") { - quota::settle_anonymous_task_reservation(state, task_id).await?; - return Ok(()); - } - - let is_retry = task.status == "processing"; - - let updated = sqlx::query( - r#" - UPDATE tasks - SET status = 'processing', started_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 && task.status == "pending" { - // Another worker might have taken it. - return Ok(()); - } - - // Refresh task row after status change - task.status = "processing".to_string(); - - if is_retry { - sqlx::query( - "UPDATE task_files SET status = 'pending' WHERE task_id = $1 AND status = 'processing'", - ) - .bind(task_id) - .execute(&state.db) - .await - .map_err(|err| AppError::new(ErrorCode::Internal, "恢复待重试文件失败").with_source(err))?; - } - let compression_rate = task.compression_rate.and_then(|v| u8::try_from(v).ok()); let level = compression_rate .map(compress::rate_to_level) @@ -621,10 +803,10 @@ async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> { id, input_path, original_format, - output_format, - status::text AS status + output_format FROM task_files WHERE task_id = $1 + AND status IN ('pending', 'processing') ORDER BY created_at ASC "#, ) @@ -660,10 +842,6 @@ async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> { let mut join_set = JoinSet::new(); for file in files.drain(..) { - if file.status != "pending" { - continue; - } - let permit = semaphore.clone().acquire_owned().await.unwrap(); let state = state.clone(); let ctx = ctx.clone(); @@ -675,6 +853,8 @@ async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> { let result = process_task_file( state, task_id, + task.processing_attempt, + worker_id, file, level, compression_rate, @@ -712,8 +892,11 @@ async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> { return Err(err); } - finalize_task_status(state, task_id).await?; - Ok(()) + if finalize_task_status(state, task_id, task.processing_attempt, worker_id).await? { + Ok(TaskProcessOutcome::Done) + } else { + Ok(TaskProcessOutcome::LeaseBusy) + } } fn parse_image_fmt(value: &str) -> Result { @@ -739,10 +922,51 @@ async fn is_task_cancelled(state: &AppState, task_id: Uuid) -> Result Result { + sqlx::query_scalar( + r#" + SELECT EXISTS( + SELECT 1 + FROM tasks t + JOIN task_files f ON f.task_id = t.id + WHERE t.id = $1 + AND t.status = 'processing' + AND t.processing_attempt = $2 + AND t.lease_owner = $5 + AND t.lease_until > NOW() + AND f.id = $3 + AND f.status = 'processing' + AND f.processing_attempt = $4 + AND f.lease_owner = $5 + AND f.lease_until > NOW() + ) + "#, + ) + .bind(fence.task_id) + .bind(fence.task_attempt) + .bind(fence.file_id) + .bind(fence.file_attempt) + .bind(fence.worker_id) + .fetch_one(&state.db) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "检查文件处理租约失败").with_source(err)) +} + #[allow(clippy::too_many_arguments)] async fn process_task_file( state: AppState, task_id: Uuid, + task_attempt: i64, + worker_id: Uuid, file: TaskFileProcRow, level: compress::CompressionLevel, compression_rate: Option, @@ -751,36 +975,65 @@ async fn process_task_file( ctx: TaskContext, billing_ctx: Option, ) -> Result<(), AppError> { - let updated = sqlx::query( + let file_attempt: Option = sqlx::query_scalar( r#" UPDATE task_files AS f - SET status = 'processing' + SET status = 'processing', + processing_attempt = processing_attempt + 1, + lease_owner = $4, + lease_until = NOW() + $5 * INTERVAL '1 second', + error_message = NULL FROM tasks AS t WHERE f.id = $1 AND f.task_id = t.id AND t.id = $2 - AND f.status = 'pending' - AND t.status IN ('pending', 'processing') + AND t.status = 'processing' + AND t.processing_attempt = $3 + AND t.lease_owner = $4 + AND t.lease_until > NOW() + AND ( + f.status = 'pending' + OR ( + f.status = 'processing' + AND ( + f.lease_owner = $4 + OR f.lease_owner IS NULL + OR f.lease_until IS NULL + OR f.lease_until <= NOW() + ) + ) + ) + RETURNING f.processing_attempt "#, ) .bind(file.id) .bind(task_id) - .execute(&state.db) + .bind(task_attempt) + .bind(worker_id) + .bind(PROCESSING_LEASE_SECONDS) + .fetch_optional(&state.db) .await .map_err(|err| AppError::new(ErrorCode::Internal, "更新文件处理状态失败").with_source(err))?; - if updated.rows_affected() == 0 { + let Some(file_attempt) = file_attempt else { return Ok(()); - } + }; + let fence = FileFence { + task_id, + task_attempt, + file_id: file.id, + file_attempt, + worker_id, + }; let Some(input_path) = file.input_path.clone() else { - mark_file_failed(&state, task_id, file.id, "原文件不存在").await?; + mark_file_failed(&state, &fence, "原文件不存在").await?; return Ok(()); }; let input_bytes = match tokio::fs::read(&input_path).await { Ok(v) => v, Err(_) => { - mark_file_failed(&state, task_id, file.id, "读取原文件失败").await?; + mark_file_failed_and_cleanup(&state, &fence, "读取原文件失败", &input_path).await?; return Ok(()); } }; @@ -788,16 +1041,14 @@ async fn process_task_file( let format_in = match parse_image_fmt(&file.original_format) { Ok(format) => format, Err(err) => { - mark_file_failed(&state, task_id, file.id, &err.message).await?; - let _ = tokio::fs::remove_file(&input_path).await; + mark_file_failed_and_cleanup(&state, &fence, &err.message, &input_path).await?; return Ok(()); } }; let format_out = match parse_image_fmt(&file.output_format) { Ok(format) => format, Err(err) => { - mark_file_failed(&state, task_id, file.id, &err.message).await?; - let _ = tokio::fs::remove_file(&input_path).await; + mark_file_failed_and_cleanup(&state, &fence, &err.message, &input_path).await?; return Ok(()); } }; @@ -819,15 +1070,13 @@ async fn process_task_file( { Ok(v) => v, Err(err) => { - mark_file_failed(&state, task_id, file.id, &err.message).await?; - let _ = tokio::fs::remove_file(&input_path).await; + mark_file_failed_and_cleanup(&state, &fence, &err.message, &input_path).await?; return Ok(()); } }; - if is_task_cancelled(&state, task_id).await? { - mark_file_failed(&state, task_id, file.id, "已取消").await?; - let _ = tokio::fs::remove_file(&input_path).await; + if !file_attempt_is_current(&state, &fence).await? || is_task_cancelled(&state, task_id).await? + { return Ok(()); } @@ -843,10 +1092,12 @@ async fn process_task_file( && max_height.is_none(); let charge_units = !skip_charge && compressed_size < original_size; - let object_key = storage::result_key( + let object_key = storage::result_attempt_key( ctx.retention_hours as i64, task_id, file.id, + task_attempt, + file_attempt, format_out.extension(), ); let stored = match storage::store_bytes( @@ -859,39 +1110,39 @@ async fn process_task_file( { Ok(value) => value, Err(err) => { - reset_file_for_retry(&state, file.id, "对象存储暂时不可用").await?; - return Err(err); + if reset_file_for_retry(&state, &fence, "对象存储暂时不可用").await? { + return Err(err); + } + return Ok(()); } }; 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; - mark_file_failed(&state, task_id, file.id, "匿名任务缺少 session_id").await?; - let _ = tokio::fs::remove_file(&input_path).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; - mark_file_failed(&state, task_id, file.id, "匿名任务缺少 client_ip").await?; - let _ = tokio::fs::remove_file(&input_path).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; - mark_file_failed(&state, task_id, file.id, &err.message).await?; - let _ = tokio::fs::remove_file(&input_path).await; + mark_file_failed_and_cleanup(&state, &fence, &err.message, &input_path).await?; return Ok(()); } } - if let Err(err) = finalize_file( + match finalize_file( &state, &billing_ctx, ctx.api_key_id, &ctx.source, - task_id, - file.id, + &fence, &stored, original_size as i64, compressed_size as i64, @@ -902,24 +1153,33 @@ async fn process_task_file( ) .await { - let _ = storage::delete_object(&state, &stored_locator(&stored)).await; - mark_file_failed(&state, task_id, file.id, &err.message).await?; - let _ = tokio::fs::remove_file(&input_path).await; - return Ok(()); + Ok(FinalizeFileOutcome::Committed) => { + 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?; + } } - - let _ = tokio::fs::remove_file(&input_path).await; Ok(()) } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum FinalizeFileOutcome { + Committed, + LeaseLost, +} + #[allow(clippy::too_many_arguments)] async fn finalize_file( state: &AppState, billing_ctx: &Option, api_key_id: Option, source: &str, - task_id: Uuid, - task_file_id: Uuid, + fence: &FileFence, stored: &storage::StoredObject, bytes_in: i64, bytes_out: i64, @@ -927,26 +1187,60 @@ async fn finalize_file( format_in: compress::ImageFmt, format_out: compress::ImageFmt, charge_units: bool, -) -> Result<(), AppError> { +) -> Result { let mut tx = state .db .begin() .await .map_err(|err| AppError::new(ErrorCode::Internal, "开启事务失败").with_source(err))?; - let task_status: 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, "锁定任务状态失败").with_source(err) - })?; - if matches!(task_status.as_deref(), Some("cancelled")) { - return Err(AppError::new(ErrorCode::InvalidRequest, "已取消")); + let task_fence: Option<(String, i64, Option, bool)> = sqlx::query_as( + r#" + SELECT status::text, processing_attempt, lease_owner, + COALESCE(lease_until > NOW(), false) + FROM tasks + WHERE id = $1 + FOR UPDATE + "#, + ) + .bind(fence.task_id) + .fetch_optional(&mut *tx) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "锁定任务状态失败").with_source(err))?; + let Some((task_status, task_attempt, task_owner, task_lease_valid)) = task_fence else { + return Ok(FinalizeFileOutcome::LeaseLost); + }; + if task_status != "processing" + || task_attempt != fence.task_attempt + || task_owner != Some(fence.worker_id) + || !task_lease_valid + { + return Ok(FinalizeFileOutcome::LeaseLost); } - if !matches!(task_status.as_deref(), Some("pending" | "processing")) { - return Err(AppError::new(ErrorCode::InvalidRequest, "任务已结束")); + + let file_fence: Option<(String, i64, Option, bool)> = sqlx::query_as( + r#" + SELECT status::text, processing_attempt, lease_owner, + COALESCE(lease_until > NOW(), false) + FROM task_files + WHERE id = $1 AND task_id = $2 + FOR UPDATE + "#, + ) + .bind(fence.file_id) + .bind(fence.task_id) + .fetch_optional(&mut *tx) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "锁定文件状态失败").with_source(err))?; + let Some((file_status, file_attempt, file_owner, file_lease_valid)) = file_fence else { + return Ok(FinalizeFileOutcome::LeaseLost); + }; + if file_status != "processing" + || file_attempt != fence.file_attempt + || file_owner != Some(fence.worker_id) + || !file_lease_valid + { + return Ok(FinalizeFileOutcome::LeaseLost); } // Paid users: charge before marking file completed (atomic w/ status update). @@ -957,8 +1251,8 @@ async fn finalize_file( billing, api_key_id, source, - task_id, - task_file_id, + fence.task_id, + fence.file_id, format_in, format_out, bytes_in as u64, @@ -968,7 +1262,7 @@ async fn finalize_file( } } - sqlx::query( + let file_updated = sqlx::query( r#" UPDATE task_files SET storage_path = $2, @@ -980,11 +1274,17 @@ async fn finalize_file( compressed_size = $7, saved_percent = $8, status = 'completed', - completed_at = NOW() + completed_at = NOW(), + lease_owner = NULL, + lease_until = NULL WHERE id = $1 + AND task_id = $9 + AND status = 'processing' + AND processing_attempt = $10 + AND lease_owner = $11 "#, ) - .bind(task_file_id) + .bind(fence.file_id) .bind(if stored.backend == "local" { Some(stored.key.as_str()) } else { @@ -996,43 +1296,80 @@ async fn finalize_file( .bind(&stored.etag) .bind(bytes_out) .bind(saved_percent) + .bind(fence.task_id) + .bind(fence.file_attempt) + .bind(fence.worker_id) .execute(&mut *tx) .await .map_err(|err| AppError::new(ErrorCode::Internal, "更新文件失败").with_source(err))?; - sqlx::query( + if file_updated.rows_affected() == 0 { + return Ok(FinalizeFileOutcome::LeaseLost); + } + + let task_updated = sqlx::query( r#" UPDATE tasks - SET completed_files = completed_files + 1, + SET completed_files = LEAST(total_files, completed_files + 1), total_compressed_size = total_compressed_size + $2 WHERE id = $1 + AND status = 'processing' + AND processing_attempt = $3 + AND lease_owner = $4 "#, ) - .bind(task_id) + .bind(fence.task_id) .bind(bytes_out) + .bind(fence.task_attempt) + .bind(fence.worker_id) .execute(&mut *tx) .await .map_err(|err| AppError::new(ErrorCode::Internal, "更新任务统计失败").with_source(err))?; + if task_updated.rows_affected() == 0 { + return Ok(FinalizeFileOutcome::LeaseLost); + } tx.commit() .await .map_err(|err| AppError::new(ErrorCode::Internal, "提交事务失败").with_source(err))?; - Ok(()) + Ok(FinalizeFileOutcome::Committed) } async fn mark_file_failed( state: &AppState, - task_id: Uuid, - task_file_id: Uuid, + fence: &FileFence, message: &str, -) -> Result<(), AppError> { +) -> Result { let mut tx = state .db .begin() .await .map_err(|err| AppError::new(ErrorCode::Internal, "开启事务失败").with_source(err))?; + let task_owned: Option = sqlx::query_scalar( + r#" + SELECT id + FROM tasks + WHERE id = $1 + AND status = 'processing' + AND processing_attempt = $2 + AND lease_owner = $3 + AND lease_until > NOW() + FOR UPDATE + "#, + ) + .bind(fence.task_id) + .bind(fence.task_attempt) + .bind(fence.worker_id) + .fetch_optional(&mut *tx) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "锁定任务状态失败").with_source(err))?; + if task_owned.is_none() { + tx.rollback().await.ok(); + return Ok(false); + } + let updated = sqlx::query( r#" UPDATE task_files @@ -1040,13 +1377,21 @@ async fn mark_file_failed( error_message = $2, storage_path = NULL, input_path = NULL, - completed_at = NOW() + completed_at = NOW(), + lease_owner = NULL, + lease_until = NULL WHERE id = $1 - AND status NOT IN ('completed', 'failed') + AND task_id = $3 + AND status = 'processing' + AND processing_attempt = $4 + AND lease_owner = $5 "#, ) - .bind(task_file_id) + .bind(fence.file_id) .bind(message) + .bind(fence.task_id) + .bind(fence.file_attempt) + .bind(fence.worker_id) .execute(&mut *tx) .await .map_err(|err| AppError::new(ErrorCode::Internal, "更新文件失败").with_source(err))?; @@ -1055,11 +1400,16 @@ async fn mark_file_failed( sqlx::query( r#" UPDATE tasks - SET failed_files = failed_files + 1 + SET failed_files = LEAST(total_files, failed_files + 1) WHERE id = $1 + AND status = 'processing' + AND processing_attempt = $2 + AND lease_owner = $3 "#, ) - .bind(task_id) + .bind(fence.task_id) + .bind(fence.task_attempt) + .bind(fence.worker_id) .execute(&mut *tx) .await .map_err(|err| AppError::new(ErrorCode::Internal, "更新任务统计失败").with_source(err))?; @@ -1069,89 +1419,251 @@ async fn mark_file_failed( .await .map_err(|err| AppError::new(ErrorCode::Internal, "提交事务失败").with_source(err))?; + Ok(updated.rows_affected() > 0) +} + +async fn mark_file_failed_and_cleanup( + state: &AppState, + fence: &FileFence, + message: &str, + input_path: &str, +) -> Result<(), AppError> { + if mark_file_failed(state, fence, message).await? { + let _ = tokio::fs::remove_file(input_path).await; + } Ok(()) } async fn reset_file_for_retry( state: &AppState, - task_file_id: Uuid, + fence: &FileFence, message: &str, -) -> Result<(), AppError> { - sqlx::query( - "UPDATE task_files SET status = 'pending', error_message = $2 WHERE id = $1 AND status = 'processing'", +) -> Result { + let mut tx = + state.db.begin().await.map_err(|err| { + AppError::new(ErrorCode::Internal, "开启重试事务失败").with_source(err) + })?; + let task_owned: Option = sqlx::query_scalar( + r#" + SELECT id FROM tasks + WHERE id = $1 + AND status = 'processing' + AND processing_attempt = $2 + AND lease_owner = $3 + AND lease_until > NOW() + FOR UPDATE + "#, ) - .bind(task_file_id) + .bind(fence.task_id) + .bind(fence.task_attempt) + .bind(fence.worker_id) + .fetch_optional(&mut *tx) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "锁定重试任务失败").with_source(err))?; + if task_owned.is_none() { + tx.rollback().await.ok(); + return Ok(false); + } + let updated = sqlx::query( + r#" + UPDATE task_files + SET status = 'pending', + error_message = $2, + lease_owner = NULL, + lease_until = NULL + WHERE id = $1 + AND task_id = $3 + AND status = 'processing' + AND processing_attempt = $4 + AND lease_owner = $5 + "#, + ) + .bind(fence.file_id) .bind(message) - .execute(&state.db) + .bind(fence.task_id) + .bind(fence.file_attempt) + .bind(fence.worker_id) + .execute(&mut *tx) .await .map_err(|err| AppError::new(ErrorCode::Internal, "恢复待重试文件失败").with_source(err))?; - Ok(()) + tx.commit() + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "提交重试状态失败").with_source(err))?; + Ok(updated.rows_affected() > 0) } -async fn finalize_task_status(state: &AppState, task_id: Uuid) -> Result<(), AppError> { - let row: Option<(i32, i32, i32, String)> = sqlx::query_as( - "SELECT total_files, completed_files, failed_files, status::text AS status FROM tasks WHERE id = $1", +async fn finalize_task_status( + state: &AppState, + task_id: Uuid, + task_attempt: i64, + worker_id: Uuid, +) -> Result { + let mut tx = state.db.begin().await.map_err(|err| { + AppError::new(ErrorCode::Internal, "开启任务终结事务失败").with_source(err) + })?; + let task: Option<(String, i64, Option, bool, i32)> = sqlx::query_as( + r#" + SELECT status::text, processing_attempt, lease_owner, + COALESCE(lease_until > NOW(), false), total_files + 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))?; - - let Some((total, completed, failed, status)) = row else { - return Ok(()); + .map_err(|err| AppError::new(ErrorCode::Internal, "锁定任务失败").with_source(err))?; + let Some((status, current_attempt, current_owner, lease_valid, total)) = task else { + return Ok(true); }; if status == "cancelled" { - let paths: Vec> = sqlx::query_scalar( - "SELECT input_path FROM task_files WHERE task_id = $1 AND status IN ('pending','processing')", - ) - .bind(task_id) - .fetch_all(&state.db) - .await - .unwrap_or_default(); - for p in paths.into_iter().flatten() { - let _ = tokio::fs::remove_file(p).await; - } + tx.rollback().await.ok(); + finalize_cancelled_task(state, task_id).await?; + return Ok(true); + } + if status != "processing" + || current_attempt != task_attempt + || current_owner != Some(worker_id) + || !lease_valid + { + return Ok(false); + } - let _ = sqlx::query( - "UPDATE task_files SET status = 'failed', error_message = '已取消', storage_path = NULL, input_path = NULL, completed_at = NOW() WHERE task_id = $1 AND status IN ('pending','processing')", - ) - .bind(task_id) - .execute(&state.db) - .await; - - let _ = sqlx::query( + let (completed, failed, total_compressed_size): (i64, i64, i64) = sqlx::query_as( + r#" + SELECT + COUNT(*) FILTER (WHERE status = 'completed'), + COUNT(*) FILTER (WHERE status = 'failed'), + COALESCE(SUM(compressed_size) FILTER (WHERE status = 'completed'), 0)::bigint + FROM task_files + WHERE task_id = $1 + "#, + ) + .bind(task_id) + .fetch_one(&mut *tx) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "统计任务文件失败").with_source(err))?; + let finished = completed + failed >= i64::from(total) && total > 0; + if !finished { + sqlx::query( r#" UPDATE tasks - SET completed_files = (SELECT COUNT(*) FROM task_files WHERE task_id = $1 AND status = 'completed'), - failed_files = (SELECT COUNT(*) FROM task_files WHERE task_id = $1 AND status = 'failed'), - completed_at = COALESCE(completed_at, NOW()) - WHERE id = $1 + SET completed_files = $4, failed_files = $5, total_compressed_size = $6 + WHERE id = $1 AND status = 'processing' + AND processing_attempt = $2 AND lease_owner = $3 "#, ) .bind(task_id) - .execute(&state.db) - .await; - quota::settle_anonymous_task_reservation(state, task_id).await?; + .bind(task_attempt) + .bind(worker_id) + .bind(completed as i32) + .bind(failed as i32) + .bind(total_compressed_size) + .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) + })?; + return Ok(false); + } + + let final_status = if completed == 0 && failed == i64::from(total) { + "failed" + } else { + "completed" + }; + let updated = sqlx::query( + r#" + UPDATE tasks + SET status = $4::task_status, + completed_files = $5, + failed_files = $6, + total_compressed_size = $7, + completed_at = NOW(), + lease_owner = NULL, + lease_until = NULL + WHERE id = $1 AND status = 'processing' + AND processing_attempt = $2 AND lease_owner = $3 + "#, + ) + .bind(task_id) + .bind(task_attempt) + .bind(worker_id) + .bind(final_status) + .bind(completed as i32) + .bind(failed as i32) + .bind(total_compressed_size) + .execute(&mut *tx) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "更新任务状态失败").with_source(err))?; + if updated.rows_affected() == 0 { + return Ok(false); + } + tx.commit() + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "提交任务状态失败").with_source(err))?; + quota::settle_anonymous_task_reservation(state, task_id).await?; + Ok(true) +} + +async fn finalize_cancelled_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> { + let mut tx = state.db.begin().await.map_err(|err| { + AppError::new(ErrorCode::Internal, "开启取消清理事务失败").with_source(err) + })?; + let status: 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, "锁定取消任务失败").with_source(err) + })?; + if status.as_deref() != Some("cancelled") { + tx.rollback().await.ok(); return Ok(()); } - - if completed + failed >= total && total > 0 { - let final_status = if completed == 0 && failed == total { - "failed" - } else { - "completed" - }; - sqlx::query( - "UPDATE tasks SET status = $2::task_status, completed_at = NOW() WHERE id = $1", - ) - .bind(task_id) - .bind(final_status) - .execute(&state.db) + let paths: Vec> = sqlx::query_scalar( + "SELECT input_path FROM task_files WHERE task_id = $1 AND status IN ('pending','processing') FOR UPDATE", + ) + .bind(task_id) + .fetch_all(&mut *tx) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "锁定取消文件失败").with_source(err))?; + sqlx::query( + r#" + UPDATE task_files + SET status = 'failed', error_message = '已取消', storage_path = NULL, + input_path = NULL, completed_at = NOW(), lease_owner = NULL, lease_until = NULL + WHERE task_id = $1 AND status IN ('pending','processing') + "#, + ) + .bind(task_id) + .execute(&mut *tx) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "取消任务文件失败").with_source(err))?; + sqlx::query( + r#" + UPDATE tasks + SET completed_files = (SELECT COUNT(*) FROM task_files WHERE task_id = $1 AND status = 'completed'), + failed_files = (SELECT COUNT(*) FROM task_files WHERE task_id = $1 AND status = 'failed'), + total_compressed_size = COALESCE((SELECT SUM(compressed_size) FROM task_files WHERE task_id = $1 AND status = 'completed'), 0)::bigint, + completed_at = COALESCE(completed_at, NOW()), lease_owner = NULL, lease_until = NULL + WHERE id = $1 AND status = 'cancelled' + "#, + ) + .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))?; - quota::settle_anonymous_task_reservation(state, task_id).await?; + .map_err(|err| AppError::new(ErrorCode::Internal, "提交取消清理失败").with_source(err))?; + for path in paths.into_iter().flatten() { + let _ = tokio::fs::remove_file(path).await; } - + quota::settle_anonymous_task_reservation(state, task_id).await?; Ok(()) } @@ -1295,6 +1807,12 @@ async fn cleanup_expired_records(state: &AppState) -> Result<(), AppError> { .execute(&state.db) .await; + let _ = sqlx::query( + "DELETE FROM email_change_requests WHERE expires_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) @@ -1436,6 +1954,12 @@ fn stored_locator(stored: &storage::StoredObject) -> storage::ObjectLocator { #[cfg(test)] mod tests { use super::*; + use crate::config::Config; + use crate::services::mail::Mailer; + use bytes::Bytes; + use chrono::Utc; + use sqlx::postgres::PgPoolOptions; + use tokio::sync::Barrier; #[test] fn third_delivery_moves_message_to_dead_letter() { @@ -1461,4 +1985,326 @@ mod tests { fn message_heartbeat_precedes_stale_claim_threshold() { assert!(MESSAGE_HEARTBEAT_SECONDS * 1_000 < STALE_MESSAGE_IDLE_MS as u64); } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + #[ignore = "requires isolated IMAGEFORGE_TEST_DATABASE_URL and IMAGEFORGE_TEST_REDIS_URL"] + async fn concurrent_attempts_finalize_once_and_keep_winning_object() { + let database_url = std::env::var("IMAGEFORGE_TEST_DATABASE_URL") + .expect("IMAGEFORGE_TEST_DATABASE_URL must be set"); + assert!( + database_url.to_ascii_lowercase().contains("test"), + "refusing to run destructive integration test outside a test database" + ); + let redis_url = std::env::var("IMAGEFORGE_TEST_REDIS_URL") + .expect("IMAGEFORGE_TEST_REDIS_URL must be set"); + let pool = PgPoolOptions::new() + .max_connections(16) + .connect(&database_url) + .await + .expect("connect test database"); + sqlx::migrate!().run(&pool).await.expect("run migrations"); + + let mut config = Config::from_env().expect("load test config"); + config.database_url = database_url; + config.redis_url = redis_url; + config.storage_path = format!( + "{}/imageforge-worker-fence-{}", + std::env::temp_dir().display(), + Uuid::new_v4() + ); + tokio::fs::create_dir_all(&config.storage_path) + .await + .expect("create test storage directory"); + let redis = redis::Client::open(config.redis_url.clone()) + .expect("create test redis client") + .get_connection_manager() + .await + .expect("connect test redis"); + let state = AppState { + mailer: Arc::new(Mailer::new(&config).expect("create disabled test mailer")), + image_processing_semaphore: Arc::new(Semaphore::new(2)), + runtime_policy_cache: crate::services::settings::RuntimePolicyCache::new(), + storage_cache: storage::StorageCache::new(), + config, + db: pool.clone(), + redis, + }; + + ensure_group(&state) + .await + .expect("create test stream group"); + let stream_task_id = Uuid::new_v4(); + let mut redis_conn = state.redis.clone(); + let stream_message_id: String = redis::cmd("XADD") + .arg(STREAM_KEY) + .arg("*") + .arg("task_id") + .arg(stream_task_id.to_string()) + .query_async(&mut redis_conn) + .await + .expect("append test stream message"); + let stream_owner = format!("test-owner-{}", Uuid::new_v4()); + let options = StreamReadOptions::default() + .group(GROUP_NAME, &stream_owner) + .count(100); + let _: redis::streams::StreamReadReply = redis_conn + .xread_options(&[STREAM_KEY], &[">"], &options) + .await + .expect("claim test stream message"); + assert!( + !touch_pending_message(&state, "test-intruder", &stream_message_id) + .await + .expect("reject intruder heartbeat") + ); + assert!( + touch_pending_message(&state, &stream_owner, &stream_message_id) + .await + .expect("accept owner heartbeat") + ); + ack_message(&mut redis_conn, &stream_message_id) + .await + .expect("ack test stream message"); + let _: i64 = redis::cmd("XDEL") + .arg(STREAM_KEY) + .arg(&stream_message_id) + .query_async(&mut redis_conn) + .await + .expect("delete test stream message"); + + let marker = Uuid::new_v4().simple().to_string(); + let user_id = Uuid::new_v4(); + let task_id = Uuid::new_v4(); + let file_id = Uuid::new_v4(); + let stale_owner = Uuid::new_v4(); + let winning_owner = Uuid::new_v4(); + sqlx::query( + "INSERT INTO users (id, email, username, password_hash) VALUES ($1, $2, $3, 'test-only')", + ) + .bind(user_id) + .bind(format!("worker-{marker}@example.test")) + .bind(format!("worker_{marker}")) + .execute(&pool) + .await + .expect("insert test user"); + sqlx::query( + r#" + INSERT INTO tasks ( + id, user_id, status, total_files, total_original_size, + processing_attempt, lease_owner, lease_until + ) VALUES ($1, $2, 'processing', 1, 100, 2, $3, NOW() + INTERVAL '5 minutes') + "#, + ) + .bind(task_id) + .bind(user_id) + .bind(winning_owner) + .execute(&pool) + .await + .expect("insert test task"); + sqlx::query( + r#" + INSERT INTO task_files ( + id, task_id, original_name, original_format, output_format, + original_size, status, processing_attempt, lease_owner, lease_until + ) VALUES ( + $1, $2, 'fence.png', 'png', 'png', + 100, 'processing', 2, $3, NOW() + INTERVAL '5 minutes' + ) + "#, + ) + .bind(file_id) + .bind(task_id) + .bind(winning_owner) + .execute(&pool) + .await + .expect("insert test task file"); + + 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( + &state, + &stale_key, + Bytes::from_static(b"stale-attempt"), + "image/png", + ) + .await + .expect("store stale attempt object"); + let winning_object = storage::store_bytes( + &state, + &winning_key, + Bytes::from_static(b"winning-attempt"), + "image/png", + ) + .await + .expect("store winning attempt object"); + + let period_start = Utc::now() - chrono::Duration::hours(1); + let period_end = Utc::now() + chrono::Duration::days(30); + let billing = billing::BillingContext { + user_id, + subscription_id: None, + plan: billing::Plan { + included_units_per_period: 100, + max_file_size_mb: 10, + max_files_per_batch: 10, + retention_days: 1, + feature_api_enabled: false, + }, + period_start, + period_end, + }; + let stale_fence = FileFence { + task_id, + task_attempt: 1, + file_id, + file_attempt: 1, + worker_id: stale_owner, + }; + let winning_fence = FileFence { + task_id, + task_attempt: 2, + file_id, + file_attempt: 2, + worker_id: winning_owner, + }; + let barrier = Arc::new(Barrier::new(2)); + + let stale_join = { + let state = state.clone(); + let billing = billing.clone(); + let stored = stale_object.clone(); + let barrier = barrier.clone(); + tokio::spawn(async move { + barrier.wait().await; + finalize_file( + &state, + &Some(billing), + None, + "web", + &stale_fence, + &stored, + 100, + 50, + 50.0, + compress::ImageFmt::Png, + compress::ImageFmt::Png, + true, + ) + .await + }) + }; + let winning_join = { + let state = state.clone(); + let billing = billing.clone(); + let stored = winning_object.clone(); + let barrier = barrier.clone(); + tokio::spawn(async move { + barrier.wait().await; + finalize_file( + &state, + &Some(billing), + None, + "web", + &winning_fence, + &stored, + 100, + 40, + 60.0, + compress::ImageFmt::Png, + compress::ImageFmt::Png, + true, + ) + .await + }) + }; + let stale_result = stale_join + .await + .expect("stale finalize task") + .expect("stale finalize"); + let winning_result = winning_join + .await + .expect("winning finalize task") + .expect("winning finalize"); + 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()); + assert_eq!( + storage::read_bytes(&state, &stored_locator(&winning_object)) + .await + .expect("read winning object"), + b"winning-attempt" + ); + + assert!(finalize_task_status(&state, task_id, 2, winning_owner) + .await + .expect("finalize test task")); + let task: (String, i32, i32, i64) = sqlx::query_as( + "SELECT status::text, completed_files, failed_files, total_compressed_size FROM tasks WHERE id = $1", + ) + .bind(task_id) + .fetch_one(&pool) + .await + .expect("query test task"); + assert_eq!(task, ("completed".to_string(), 1, 0, 40)); + let file: (String, String, i64) = sqlx::query_as( + "SELECT status::text, storage_key, compressed_size FROM task_files WHERE id = $1", + ) + .bind(file_id) + .fetch_one(&pool) + .await + .expect("query test file"); + assert_eq!( + file, + ("completed".to_string(), winning_object.key.clone(), 40) + ); + let usage_event_count: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM usage_events WHERE task_file_id = $1") + .bind(file_id) + .fetch_one(&pool) + .await + .expect("count usage events"); + assert_eq!(usage_event_count, 1); + let used_units: i32 = sqlx::query_scalar( + "SELECT used_units FROM usage_periods WHERE user_id = $1 AND period_start = $2 AND period_end = $3", + ) + .bind(user_id) + .bind(period_start) + .bind(period_end) + .fetch_one(&pool) + .await + .expect("query used units"); + assert_eq!(used_units, 1); + + storage::delete_object(&state, &stored_locator(&winning_object)) + .await + .expect("delete winning object"); + sqlx::query("DELETE FROM usage_events WHERE task_id = $1") + .bind(task_id) + .execute(&pool) + .await + .expect("delete test usage events"); + sqlx::query("DELETE FROM tasks WHERE id = $1") + .bind(task_id) + .execute(&pool) + .await + .expect("delete test task"); + sqlx::query("DELETE FROM usage_periods WHERE user_id = $1") + .bind(user_id) + .execute(&pool) + .await + .expect("delete test usage period"); + sqlx::query("DELETE FROM users WHERE id = $1") + .bind(user_id) + .execute(&pool) + .await + .expect("delete test user"); + tokio::fs::remove_dir_all(&state.config.storage_path) + .await + .expect("delete test storage directory"); + } }