fix: harden queue retries and quota reservation

src/worker/mod.rs

@@ -5,7 +5,7 @@ use crate::services::quota;
 use crate::services::storage;
 use crate::state::AppState;
 
-use redis::streams::StreamReadOptions;
+use redis::streams::{StreamClaimReply, StreamId, StreamPendingCountReply, StreamReadOptions};
 use redis::AsyncCommands;
 use sqlx::FromRow;
 use std::net::IpAddr;
@@ -17,6 +17,9 @@ 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");
@@ -70,6 +73,10 @@ async fn ensure_group(state: &AppState, _consumer: &str) -> Result<(), AppError>
 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)
@@ -97,31 +104,187 @@ async fn poll_once(state: &AppState, consumer: &str) -> Result<(), AppError> {
 
     for key in reply.keys {
         for msg in key.ids {
-            let Some(task_id_str) = msg.get::<String>("task_id") else {
-                ack_message(&mut conn, &msg.id).await?;
-                continue;
-            };
-
-            let task_id = match Uuid::parse_str(&task_id_str) {
-                Ok(v) => v,
-                Err(_) => {
-                    ack_message(&mut conn, &msg.id).await?;
-                    continue;
-                }
-            };
-
-            process_task(state, task_id).await.map_err(|err| {
-                tracing::error!(task_id = %task_id, error = %err, "task processing failed; message left pending for retry");
-                err
-            })?;
-
-            ack_message(&mut conn, &msg.id).await?;
+            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::<String>("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<Option<StreamId>, 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<usize, AppError> {
+    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,
@@ -150,6 +313,7 @@ struct TaskProcRow {
     source: String,
     client_ip: Option<String>,
     retention_hours: i32,
+    anonymous_units_reserved: i32,
 }
 
 #[derive(Debug, FromRow)]
@@ -185,10 +349,11 @@ struct TaskContext {
     anon_ip: Option<IpAddr>,
     is_anonymous: bool,
     retention_hours: i32,
+    anonymous_quota_reserved: bool,
 }
 
 async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
-    let mut task: TaskProcRow = sqlx::query_as(
+    let Some(mut task): Option<TaskProcRow> = sqlx::query_as(
         r#"
         SELECT
           status::text AS status,
@@ -202,7 +367,8 @@ async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
           api_key_id,
           source::text AS source,
           host(client_ip) AS client_ip,
-          retention_hours
+          retention_hours,
+          anonymous_units_reserved
         FROM tasks
         WHERE id = $1
         "#,
@@ -211,7 +377,10 @@ async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
     .fetch_optional(&state.db)
     .await
     .map_err(|err| AppError::new(ErrorCode::Internal, "查询任务失败").with_source(err))?
-    .ok_or_else(|| AppError::new(ErrorCode::NotFound, "任务不存在"))?;
+    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(());
@@ -293,6 +462,7 @@ async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
         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;
@@ -503,7 +673,7 @@ async fn process_task_file(
         }
     };
 
-    if ctx.is_anonymous && charge_units {
+    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?;
@@ -994,3 +1164,16 @@ fn stored_locator(stored: &storage::StoredObject) -> storage::ObjectLocator {
         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));
+    }
+}