From 72f36c631e3492f8e8893a733c8f976bb7f1eef9 Mon Sep 17 00:00:00 2001
From: 237899745 <237899745@users.noreply.git.workyai.cn>
Date: Sun, 26 Jul 2026 11:06:53 +0800
Subject: [PATCH] fix(worker): qualify file attempt claim

---
 src/worker/mod.rs | 64 ++++++++++++++++++++++++++++++++++-------------
 1 file changed, 47 insertions(+), 17 deletions(-)

diff --git a/src/worker/mod.rs b/src/worker/mod.rs
index be2eb9b..1169d79 100644
--- a/src/worker/mod.rs
+++ b/src/worker/mod.rs
@@ -953,25 +953,18 @@ async fn file_attempt_is_current(state: &AppState, fence: &FileFence) -> Result<
     .map_err(|err| AppError::new(ErrorCode::Internal, "检查文件处理租约失败").with_source(err))
 }
 
-#[allow(clippy::too_many_arguments)]
-async fn process_task_file(
-    state: AppState,
+async fn claim_task_file_attempt(
+    state: &AppState,
     task_id: Uuid,
     task_attempt: i64,
+    file_id: Uuid,
     worker_id: Uuid,
-    file: TaskFileProcRow,
-    level: compress::CompressionLevel,
-    compression_rate: Option<u8>,
-    max_width: Option<u32>,
-    max_height: Option<u32>,
-    ctx: TaskContext,
-    billing_ctx: Option<billing::BillingContext>,
-) -> Result<(), AppError> {
-    let file_attempt: Option<i64> = sqlx::query_scalar(
+) -> Result<Option<i64>, AppError> {
+    sqlx::query_scalar(
         r#"
         UPDATE task_files AS f
         SET status = 'processing',
-            processing_attempt = processing_attempt + 1,
+            processing_attempt = f.processing_attempt + 1,
             lease_owner = $4,
             lease_until = NOW() + $5 * INTERVAL '1 second',
             error_message = NULL
@@ -998,14 +991,32 @@ async fn process_task_file(
         RETURNING f.processing_attempt
         "#,
     )
-    .bind(file.id)
+    .bind(file_id)
     .bind(task_id)
     .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))?;
+    .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<u8>,
+    max_width: Option<u32>,
+    max_height: Option<u32>,
+    ctx: TaskContext,
+    billing_ctx: Option<billing::BillingContext>,
+) -> Result<(), AppError> {
+    let file_attempt =
+        claim_task_file_attempt(&state, task_id, task_attempt, file.id, worker_id).await?;
     let Some(file_attempt) = file_attempt else {
         return Ok(());
     };
@@ -2140,16 +2151,35 @@ mod tests {
               original_size, status, processing_attempt, lease_owner, lease_until
             ) VALUES (
               $1, $2, 'fence.png', 'png', 'png',
-              100, 'processing', 2, $3, NOW() + INTERVAL '5 minutes'
+              100, 'pending', 0, NULL, NULL
             )
             "#,
         )
         .bind(file_id)
         .bind(task_id)
-        .bind(winning_owner)
         .execute(&pool)
         .await
         .expect("insert test task file");
+        assert_eq!(
+            claim_task_file_attempt(&state, task_id, 2, file_id, winning_owner)
+                .await
+                .expect("claim pending task file"),
+            Some(1)
+        );
+        sqlx::query(
+            r#"
+            UPDATE task_files
+            SET processing_attempt = 2,
+                lease_owner = $2,
+                lease_until = NOW() + INTERVAL '5 minutes'
+            WHERE id = $1
+            "#,
+        )
+        .bind(file_id)
+        .bind(winning_owner)
+        .execute(&pool)
+        .await
+        .expect("prepare winning file fence");
 
         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");