use crate::error::{AppError, ErrorCode}; use crate::services::billing; use crate::services::compress; use crate::services::quota; use crate::services::storage; use crate::state::AppState; use redis::streams::{StreamClaimReply, StreamId, StreamPendingCountReply, StreamReadOptions}; use redis::AsyncCommands; use sqlx::FromRow; use std::net::IpAddr; use std::sync::Arc; use std::time::Instant; use tokio::sync::Semaphore; use tokio::task::JoinSet; use uuid::Uuid; const STREAM_KEY: &str = "stream:compress_jobs"; const GROUP_NAME: &str = "compress_workers"; const DEAD_STREAM_KEY: &str = "stream:compress_jobs:dead"; const MAX_DELIVERIES: usize = 3; const STALE_MESSAGE_IDLE_MS: usize = 5 * 60 * 1000; pub async fn run(state: AppState) -> Result<(), AppError> { tracing::info!("Worker started"); crate::services::bootstrap::ensure_schema(&state).await?; let consumer = format!("worker_{}", Uuid::new_v4()); ensure_group(&state, &consumer).await?; let mut last_maintenance = Instant::now(); loop { if let Err(err) = poll_once(&state, &consumer).await { tracing::error!(error = ?err, "worker poll error"); tokio::time::sleep(std::time::Duration::from_secs(2)).await; } if last_maintenance.elapsed().as_secs() >= 300 { if let Err(err) = maintenance(&state).await { tracing::error!(error = ?err, "maintenance failed"); } last_maintenance = Instant::now(); } } } async fn ensure_group(state: &AppState, _consumer: &str) -> Result<(), AppError> { let mut conn = state.redis.clone(); let res: Result = redis::cmd("XGROUP") .arg("CREATE") .arg(STREAM_KEY) .arg(GROUP_NAME) .arg("0") .arg("MKSTREAM") .query_async(&mut conn) .await; match res { Ok(_) => Ok(()), Err(err) => { let msg = err.to_string(); if msg.contains("BUSYGROUP") { return Ok(()); } Err(AppError::new(ErrorCode::Internal, "初始化队列失败").with_source(err)) } } } async fn poll_once(state: &AppState, consumer: &str) -> Result<(), AppError> { let mut conn = state.redis.clone(); if let Some(msg) = claim_stale_message(&mut conn, consumer).await? { return handle_message(state, &mut conn, msg).await; } // Retry messages already delivered to this consumer before taking new work. let pending_opts = StreamReadOptions::default() .group(GROUP_NAME, consumer) .count(1); let mut reply: redis::streams::StreamReadReply = conn .xread_options(&[STREAM_KEY], &["0"], &pending_opts) .await .map_err(|err| AppError::new(ErrorCode::Internal, "读取待重试任务失败").with_source(err))?; if !reply.keys.iter().any(|key| !key.ids.is_empty()) { let opts = StreamReadOptions::default() .group(GROUP_NAME, consumer) .count(1) .block(5000); reply = conn .xread_options(&[STREAM_KEY], &[">"], &opts) .await .map_err(|err| AppError::new(ErrorCode::Internal, "读取队列失败").with_source(err))?; } if reply.keys.is_empty() { return Ok(()); } for key in reply.keys { for msg in key.ids { handle_message(state, &mut conn, msg).await?; } } Ok(()) } async fn handle_message( state: &AppState, conn: &mut redis::aio::ConnectionManager, msg: StreamId, ) -> Result<(), AppError> { let Some(task_id_str) = msg.get::("task_id") else { ack_message(conn, &msg.id).await?; return Ok(()); }; let task_id = match Uuid::parse_str(&task_id_str) { Ok(value) => value, Err(_) => { ack_message(conn, &msg.id).await?; return Ok(()); } }; match process_task(state, task_id).await { Ok(()) => ack_message(conn, &msg.id).await, Err(err) => { let deliveries = pending_delivery_count(conn, &msg.id).await?; if should_dead_letter(deliveries) { 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?; tracing::error!( task_id = %task_id, deliveries, error = %err, "task moved to dead-letter stream" ); return Ok(()); } tracing::warn!( task_id = %task_id, deliveries, error = %err, "task processing failed; message left pending for retry" ); Err(err) } } } async fn claim_stale_message( conn: &mut redis::aio::ConnectionManager, consumer: &str, ) -> Result, AppError> { let pending: StreamPendingCountReply = conn .xpending_count(STREAM_KEY, GROUP_NAME, "-", "+", 100) .await .map_err(|err| { AppError::new(ErrorCode::Internal, "查询停滞队列消息失败").with_source(err) })?; let Some(stale) = pending .ids .into_iter() .find(|item| item.consumer != consumer && item.last_delivered_ms >= STALE_MESSAGE_IDLE_MS) else { return Ok(None); }; let claimed: StreamClaimReply = conn .xclaim( STREAM_KEY, GROUP_NAME, consumer, STALE_MESSAGE_IDLE_MS, &[stale.id], ) .await .map_err(|err| { AppError::new(ErrorCode::Internal, "认领停滞队列消息失败").with_source(err) })?; Ok(claimed.ids.into_iter().next()) } async fn pending_delivery_count( conn: &mut redis::aio::ConnectionManager, msg_id: &str, ) -> Result { let pending: StreamPendingCountReply = conn .xpending_count(STREAM_KEY, GROUP_NAME, msg_id, msg_id, 1) .await .map_err(|err| { AppError::new(ErrorCode::Internal, "查询队列重试次数失败").with_source(err) })?; Ok(pending .ids .first() .map(|item| item.times_delivered) .unwrap_or(1)) } fn should_dead_letter(deliveries: usize) -> bool { deliveries >= MAX_DELIVERIES } async fn write_dead_letter( conn: &mut redis::aio::ConnectionManager, message_id: &str, task_id: Uuid, deliveries: usize, error: &AppError, ) -> Result<(), AppError> { redis::cmd("XADD") .arg(DEAD_STREAM_KEY) .arg("MAXLEN") .arg("~") .arg(10_000) .arg("*") .arg("task_id") .arg(task_id.to_string()) .arg("source_message_id") .arg(message_id) .arg("deliveries") .arg(deliveries) .arg("error_code") .arg(error.code.as_str()) .arg("error_message") .arg(&error.message) .arg("failed_at") .arg(chrono::Utc::now().to_rfc3339()) .query_async::<_, redis::Value>(conn) .await .map_err(|err| AppError::new(ErrorCode::Internal, "写入死信队列失败").with_source(err))?; Ok(()) } async fn mark_task_dead_letter( state: &AppState, task_id: Uuid, message: &str, ) -> Result<(), AppError> { let mut tx = state.db.begin().await.map_err(|err| { AppError::new(ErrorCode::Internal, "开启死信事务失败").with_source(err) })?; sqlx::query( r#" UPDATE task_files SET status = 'failed', error_message = $2, completed_at = NOW() 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, "标记死信文件失败").with_source(err))?; sqlx::query( r#" UPDATE tasks SET status = 'failed', failed_files = GREATEST(total_files - completed_files, failed_files), error_message = $2, completed_at = NOW() WHERE id = $1 AND status NOT IN ('completed', 'failed', 'cancelled') "#, ) .bind(task_id) .bind(message) .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(()) } async fn ack_message( conn: &mut redis::aio::ConnectionManager, msg_id: &str, ) -> Result<(), AppError> { let _: i64 = redis::cmd("XACK") .arg(STREAM_KEY) .arg(GROUP_NAME) .arg(msg_id) .query_async(conn) .await .map_err(|err| AppError::new(ErrorCode::Internal, "确认队列消息失败").with_source(err))?; Ok(()) } #[derive(Debug, FromRow)] struct TaskProcRow { status: String, compression_level: String, compression_rate: Option, max_width: Option, max_height: Option, preserve_metadata: bool, user_id: Option, session_id: Option, api_key_id: Option, source: String, client_ip: Option, retention_hours: i32, anonymous_units_reserved: i32, } #[derive(Debug, FromRow)] struct TaskFileProcRow { id: Uuid, input_path: Option, original_format: String, output_format: String, status: 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, source: String, preserve_metadata: bool, session_id: Option, anon_ip: Option, is_anonymous: bool, retention_hours: i32, anonymous_quota_reserved: bool, } async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> { let Some(mut task): Option = sqlx::query_as( r#" SELECT status::text AS status, compression_level::text AS compression_level, compression_rate, max_width, max_height, preserve_metadata, user_id, session_id, api_key_id, source::text AS source, host(client_ip) AS client_ip, retention_hours, anonymous_units_reserved FROM tasks WHERE id = $1 "#, ) .bind(task_id) .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(()); }; if matches!(task.status.as_str(), "completed" | "failed" | "cancelled") { 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) .unwrap_or(compress::parse_level(&task.compression_level)?); let max_width = task.max_width.and_then(|v| u32::try_from(v).ok()); let max_height = task.max_height.and_then(|v| u32::try_from(v).ok()); let mut files: Vec = sqlx::query_as( r#" SELECT id, input_path, original_format, output_format, status::text AS status FROM task_files WHERE task_id = $1 ORDER BY created_at ASC "#, ) .bind(task_id) .fetch_all(&state.db) .await .map_err(|err| AppError::new(ErrorCode::Internal, "查询任务文件失败").with_source(err))?; let billing_ctx = if let Some(user_id) = task.user_id { Some(billing::get_user_billing(state, user_id).await?) } else { None }; let anon_ip: Option = task .client_ip .as_deref() .and_then(|s| s.parse::().ok()); let ctx = TaskContext { api_key_id: task.api_key_id, source: task.source.clone(), preserve_metadata: task.preserve_metadata, session_id: task.session_id.clone(), anon_ip, is_anonymous: task.user_id.is_none(), retention_hours: task.retention_hours, anonymous_quota_reserved: task.anonymous_units_reserved > 0, }; let concurrency = state.config.worker_concurrency.max(1) as usize; let semaphore = Arc::new(Semaphore::new(concurrency)); 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(); let billing_ctx = billing_ctx.clone(); let file_id = file.id; join_set.spawn(async move { let _permit = permit; let result = process_task_file( state, task_id, file, level, compression_rate, max_width, max_height, ctx, billing_ctx, ) .await; if let Err(err) = &result { tracing::error!(task_id = %task_id, file_id = %file_id, error = %err, "file processing failed"); } result }); } let mut first_error = None; while let Some(result) = join_set.join_next().await { match result { Ok(Ok(())) => {} Ok(Err(err)) => { if first_error.is_none() { first_error = Some(err); } } Err(err) => { return Err( AppError::new(ErrorCode::Internal, "文件处理线程异常退出").with_source(err) ); } } } if let Some(err) = first_error { return Err(err); } finalize_task_status(state, task_id).await?; Ok(()) } fn parse_image_fmt(value: &str) -> Result { match value.trim().to_ascii_lowercase().as_str() { "png" => Ok(compress::ImageFmt::Png), "jpeg" | "jpg" => Ok(compress::ImageFmt::Jpeg), "webp" => Ok(compress::ImageFmt::Webp), "avif" => Ok(compress::ImageFmt::Avif), "gif" => Ok(compress::ImageFmt::Gif), "bmp" => Ok(compress::ImageFmt::Bmp), "tif" | "tiff" => Ok(compress::ImageFmt::Tiff), "ico" => Ok(compress::ImageFmt::Ico), _ => Err(AppError::new(ErrorCode::InvalidRequest, "未知图片格式")), } } async fn is_task_cancelled(state: &AppState, task_id: Uuid) -> Result { 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))?; Ok(matches!(status.as_deref(), Some("cancelled"))) } #[allow(clippy::too_many_arguments)] async fn process_task_file( state: AppState, task_id: Uuid, file: TaskFileProcRow, level: compress::CompressionLevel, compression_rate: Option, max_width: Option, max_height: Option, ctx: TaskContext, billing_ctx: Option, ) -> Result<(), AppError> { if is_task_cancelled(&state, task_id).await? { return Ok(()); } let updated = sqlx::query( "UPDATE task_files SET status = 'processing' WHERE id = $1 AND status = 'pending'", ) .bind(file.id) .execute(&state.db) .await .map_err(|err| AppError::new(ErrorCode::Internal, "更新文件处理状态失败").with_source(err))?; if updated.rows_affected() == 0 { return Ok(()); } if is_task_cancelled(&state, task_id).await? { mark_file_failed(&state, task_id, file.id, "已取消").await?; return Ok(()); } let Some(input_path) = file.input_path.clone() else { mark_file_failed(&state, task_id, file.id, "原文件不存在").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?; return Ok(()); } }; 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; 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; return Ok(()); } }; let original_size = input_bytes.len() as u64; let compressed = match compress::compress_image_bytes( &state, input_bytes, format_in, format_out, level, compression_rate, None, // target_size_bytes: worker 批量任务不支持精确大小 max_width, max_height, ctx.preserve_metadata, ) .await { Ok(v) => v, Err(err) => { mark_file_failed(&state, task_id, file.id, &err.message).await?; let _ = tokio::fs::remove_file(&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; return Ok(()); } let compressed_size = compressed.len() as u64; let saved_percent = if original_size == 0 { 0.0 } else { (original_size.saturating_sub(compressed_size) as f64) * 100.0 / (original_size as f64) }; let skip_charge = compression_rate == Some(100) && format_in == format_out && max_width.is_none() && max_height.is_none(); let charge_units = !skip_charge && compressed_size < original_size; let object_key = storage::result_key( ctx.retention_hours as i64, task_id, file.id, format_out.extension(), ); let stored = match storage::store_bytes( &state, &object_key, compressed, format_out.content_type(), ) .await { Ok(value) => value, Err(err) => { reset_file_for_retry(&state, file.id, "对象存储暂时不可用").await?; return Err(err); } }; 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; 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; 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; return Ok(()); } } if is_task_cancelled(&state, task_id).await? { let _ = storage::delete_object(&state, &stored_locator(&stored)).await; mark_file_failed(&state, task_id, file.id, "已取消").await?; let _ = tokio::fs::remove_file(&input_path).await; return Ok(()); } if let Err(err) = finalize_file( &state, &billing_ctx, ctx.api_key_id, &ctx.source, task_id, file.id, &stored, original_size as i64, compressed_size as i64, saved_percent, format_in, format_out, charge_units, ) .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(()); } let _ = tokio::fs::remove_file(&input_path).await; Ok(()) } #[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, stored: &storage::StoredObject, bytes_in: i64, bytes_out: i64, saved_percent: f64, format_in: compress::ImageFmt, format_out: compress::ImageFmt, charge_units: bool, ) -> Result<(), AppError> { let mut tx = state .db .begin() .await .map_err(|err| AppError::new(ErrorCode::Internal, "开启事务失败").with_source(err))?; // Paid users: charge before marking file completed (atomic w/ status update). if charge_units { if let Some(billing) = billing_ctx { charge_one_unit( &mut tx, billing, api_key_id, source, task_id, task_file_id, format_in, format_out, bytes_in as u64, bytes_out as u64, ) .await?; } } sqlx::query( r#" UPDATE task_files SET storage_path = $2, storage_backend = $3, storage_endpoint_id = $4, storage_key = $5, storage_etag = $6, input_path = NULL, compressed_size = $7, saved_percent = $8, status = 'completed', completed_at = NOW() WHERE id = $1 "#, ) .bind(task_file_id) .bind(if stored.backend == "local" { Some(stored.key.as_str()) } else { None }) .bind(&stored.backend) .bind(stored.endpoint_id) .bind(&stored.key) .bind(&stored.etag) .bind(bytes_out) .bind(saved_percent) .execute(&mut *tx) .await .map_err(|err| AppError::new(ErrorCode::Internal, "更新文件失败").with_source(err))?; sqlx::query( r#" UPDATE tasks SET completed_files = completed_files + 1, total_compressed_size = total_compressed_size + $2 WHERE id = $1 "#, ) .bind(task_id) .bind(bytes_out) .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(()) } async fn mark_file_failed( state: &AppState, task_id: Uuid, task_file_id: Uuid, message: &str, ) -> Result<(), AppError> { let mut tx = state .db .begin() .await .map_err(|err| AppError::new(ErrorCode::Internal, "开启事务失败").with_source(err))?; let updated = sqlx::query( r#" UPDATE task_files SET status = 'failed', error_message = $2, storage_path = NULL, input_path = NULL, completed_at = NOW() WHERE id = $1 AND status NOT IN ('completed', 'failed') "#, ) .bind(task_file_id) .bind(message) .execute(&mut *tx) .await .map_err(|err| AppError::new(ErrorCode::Internal, "更新文件失败").with_source(err))?; if updated.rows_affected() > 0 { sqlx::query( r#" UPDATE tasks SET failed_files = failed_files + 1 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))?; Ok(()) } async fn reset_file_for_retry( state: &AppState, task_file_id: Uuid, message: &str, ) -> Result<(), AppError> { sqlx::query( "UPDATE task_files SET status = 'pending', error_message = $2 WHERE id = $1 AND status = 'processing'", ) .bind(task_file_id) .bind(message) .execute(&state.db) .await .map_err(|err| AppError::new(ErrorCode::Internal, "恢复待重试文件失败").with_source(err))?; Ok(()) } 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", ) .bind(task_id) .fetch_optional(&state.db) .await .map_err(|err| AppError::new(ErrorCode::Internal, "查询任务失败").with_source(err))?; let Some((total, completed, failed, status)) = row else { return Ok(()); }; 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; } 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( "UPDATE tasks SET failed_files = GREATEST(total_files - completed_files, 0), completed_at = NOW() WHERE id = $1 AND completed_at IS NULL", ) .bind(task_id) .execute(&state.db) .await; 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) .await .map_err(|err| AppError::new(ErrorCode::Internal, "更新任务状态失败").with_source(err))?; } Ok(()) } #[allow(clippy::too_many_arguments)] async fn charge_one_unit( tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, billing: &billing::BillingContext, api_key_id: Option, source: &str, task_id: Uuid, task_file_id: Uuid, format_in: compress::ImageFmt, format_out: compress::ImageFmt, bytes_in: u64, bytes_out: u64, ) -> Result<(), AppError> { quota::consume_user_unit(tx, billing, bytes_in, bytes_out).await?; sqlx::query( r#" INSERT INTO usage_events ( user_id, api_key_id, source, task_id, task_file_id, units, bytes_in, bytes_out, format_in, format_out ) VALUES ( $1, $2, $3::task_source, $4, $5, 1, $6, $7, $8, $9 ) "#, ) .bind(billing.user_id) .bind(api_key_id) .bind(source) .bind(task_id) .bind(task_file_id) .bind(bytes_in as i64) .bind(bytes_out as i64) .bind(format_in.as_str()) .bind(format_out.as_str()) .execute(&mut **tx) .await .map_err(|err| AppError::new(ErrorCode::Internal, "写入用量明细失败").with_source(err))?; Ok(()) } async fn maintenance(state: &AppState) -> Result<(), AppError> { cleanup_expired_tasks(state).await?; cleanup_stale_zip_temp(state).await?; cleanup_expired_records(state).await?; Ok(()) } async fn cleanup_stale_zip_temp(state: &AppState) -> Result<(), AppError> { let root = std::path::Path::new(&state.config.storage_path).join("tmp/zips"); let mut entries = match tokio::fs::read_dir(&root).await { Ok(entries) => entries, Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(()), Err(err) => { return Err( AppError::new(ErrorCode::StorageUnavailable, "读取 ZIP 临时目录失败") .with_source(err), ) } }; let cutoff = std::time::SystemTime::now() .checked_sub(std::time::Duration::from_secs(6 * 60 * 60)) .unwrap_or(std::time::UNIX_EPOCH); while let Some(entry) = entries.next_entry().await.map_err(|err| { AppError::new(ErrorCode::StorageUnavailable, "遍历 ZIP 临时目录失败").with_source(err) })? { let metadata = match entry.metadata().await { Ok(metadata) => metadata, Err(_) => continue, }; if !metadata.is_dir() || metadata .modified() .map(|time| time >= cutoff) .unwrap_or(true) { continue; } let _ = tokio::fs::remove_dir_all(entry.path()).await; } Ok(()) } async fn cleanup_expired_records(state: &AppState) -> Result<(), AppError> { let _ = sqlx::query("DELETE FROM idempotency_keys WHERE expires_at < NOW()") .execute(&state.db) .await; let _ = sqlx::query( "DELETE FROM email_verifications WHERE expires_at < NOW() AND verified_at IS NULL", ) .execute(&state.db) .await; let _ = sqlx::query("DELETE FROM password_resets 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) .await; let _ = sqlx::query( r#" DELETE FROM storage_endpoints e 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) "#, ) .execute(&state.db) .await; Ok(()) } async fn cleanup_expired_tasks(state: &AppState) -> Result<(), AppError> { let task_ids: Vec = sqlx::query_scalar("SELECT id FROM tasks WHERE expires_at < NOW() LIMIT 200") .fetch_all(&state.db) .await .unwrap_or_default(); if task_ids.is_empty() { return Ok(()); } for task_id in task_ids { if let Err(err) = cleanup_expired_task(state, task_id).await { tracing::warn!(task_id = %task_id, error = %err, "expired task cleanup deferred"); } } Ok(()) } async fn cleanup_expired_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> { 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; } } 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(()) } fn stored_locator(stored: &storage::StoredObject) -> storage::ObjectLocator { storage::ObjectLocator { backend: stored.backend.clone(), endpoint_id: stored.endpoint_id, key: stored.key.clone(), } } #[cfg(test)] mod tests { use super::*; #[test] fn third_delivery_moves_message_to_dead_letter() { assert!(!should_dead_letter(1)); assert!(!should_dead_letter(2)); assert!(should_dead_letter(3)); assert!(should_dead_letter(10)); } }