Parallelize worker batch processing

src/worker/mod.rs

@@ -8,7 +8,10 @@ use redis::streams::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";
@@ -149,6 +152,16 @@ struct TaskFileProcRow {
     status: String,
 }
 
+#[derive(Clone)]
+struct TaskContext {
+    api_key_id: Option<Uuid>,
+    source: String,
+    preserve_metadata: bool,
+    session_id: Option<String>,
+    anon_ip: Option<IpAddr>,
+    is_anonymous: bool,
+}
+
 async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
     let mut task: TaskProcRow = sqlx::query_as(
         r#"
@@ -242,138 +255,54 @@ async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
         .as_deref()
         .and_then(|s| s.parse::<IpAddr>().ok());
 
-    for file in &mut files {
-        // Stop early if cancelled.
-        let status: Option<String> = sqlx::query_scalar("SELECT status::text FROM tasks WHERE id = $1")
-            .bind(task_id)
-            .fetch_optional(&state.db)
-            .await
-            .unwrap_or(None);
-        if matches!(status.as_deref(), Some("cancelled")) {
-            break;
-        }
+    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(),
+    };
 
+    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 updated = sqlx::query("UPDATE task_files SET status = 'processing' WHERE id = $1 AND status = 'pending'")
-            .bind(file.id)
-            .execute(&state.db)
+        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;
+            if let Err(err) = process_task_file(
+                state,
+                task_id,
+                file,
+                level,
+                compression_rate,
+                max_width,
+                max_height,
+                ctx,
+                billing_ctx,
+            )
             .await
-            .unwrap_or_else(|_| sqlx::postgres::PgQueryResult::default());
-        if updated.rows_affected() == 0 {
-            continue;
-        }
-
-        let Some(input_path) = file.storage_path.clone() else {
-            mark_file_failed(state, task_id, file.id, "原文件不存在").await?;
-            continue;
-        };
-
-        let input_bytes = match tokio::fs::read(&input_path).await {
-            Ok(v) => v,
-            Err(_) => {
-                mark_file_failed(state, task_id, file.id, "读取原文件失败").await?;
-                continue;
+            {
+                tracing::error!(task_id = %task_id, file_id = %file_id, error = %err, "file processing failed");
             }
-        };
+        });
+    }
 
-        let format_in = parse_image_fmt(&file.original_format)?;
-        let format_out = parse_image_fmt(&file.output_format)?;
-
-        let compressed = match compress::compress_image_bytes(
-            state,
-            &input_bytes,
-            format_in,
-            format_out,
-            level,
-            compression_rate,
-            max_width,
-            max_height,
-            task.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;
-                continue;
-            }
-        };
-
-        let original_size = input_bytes.len() as u64;
-        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);
-        let charge_units = !skip_charge && compressed_size < original_size;
-
-        // Anonymous quota enforcement requires session_id + client_ip.
-        if task.user_id.is_none() && charge_units {
-            let Some(session_id) = task.session_id.as_deref() else {
-                mark_file_failed(state, task_id, file.id, "匿名任务缺少 session_id").await?;
-                let _ = tokio::fs::remove_file(&input_path).await;
-                continue;
-            };
-            let Some(ip) = anon_ip else {
-                mark_file_failed(state, task_id, file.id, "匿名任务缺少 client_ip").await?;
-                let _ = tokio::fs::remove_file(&input_path).await;
-                continue;
-            };
-            if let Err(err) = quota::consume_anonymous_units(state, session_id, ip, 1).await {
-                mark_file_failed(state, task_id, file.id, &err.message).await?;
-                let _ = tokio::fs::remove_file(&input_path).await;
-                continue;
-            }
+    while let Some(result) = join_set.join_next().await {
+        if let Err(err) = result {
+            tracing::error!(task_id = %task_id, error = %err, "file worker panicked");
         }
-
-        let output_path = format!(
-            "{}/{}.{}",
-            state.config.storage_path,
-            file.id,
-            format_out.extension()
-        );
-        if let Err(err) = tokio::fs::write(&output_path, &compressed).await {
-            mark_file_failed(state, task_id, file.id, "写入压缩文件失败").await?;
-            let _ = tokio::fs::remove_file(&input_path).await;
-            return Err(AppError::new(ErrorCode::StorageUnavailable, "写入压缩文件失败").with_source(err));
-        }
-
-        if let Err(err) = finalize_file(
-            state,
-            &billing_ctx,
-            task.api_key_id,
-            &task.source,
-            task_id,
-            file.id,
-            &output_path,
-            original_size as i64,
-            compressed_size as i64,
-            saved_percent,
-            format_in,
-            format_out,
-            charge_units,
-        )
-        .await
-        {
-            // If quota exceeded for paid users, don't leave output behind.
-            if err.code == ErrorCode::QuotaExceeded {
-                let _ = tokio::fs::remove_file(&output_path).await;
-                mark_file_failed(state, task_id, file.id, &err.message).await?;
-            } else {
-                mark_file_failed(state, task_id, file.id, &err.message).await?;
-            }
-            let _ = tokio::fs::remove_file(&input_path).await;
-            continue;
-        }
-
-        // Success: remove original.
-        let _ = tokio::fs::remove_file(&input_path).await;
     }
 
     finalize_task_status(state, task_id).await?;
@@ -394,6 +323,165 @@ fn parse_image_fmt(value: &str) -> Result<compress::ImageFmt, AppError> {
     }
 }
 
+async fn is_task_cancelled(state: &AppState, task_id: Uuid) -> Result<bool, AppError> {
+    let status: Option<String> = 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")))
+}
+
+async fn process_task_file(
+    state: AppState,
+    task_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> {
+    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
+        .unwrap_or_else(|_| sqlx::postgres::PgQueryResult::default());
+    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.storage_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 = parse_image_fmt(&file.original_format)?;
+    let format_out = parse_image_fmt(&file.output_format)?;
+
+    let compressed = match compress::compress_image_bytes(
+        &state,
+        &input_bytes,
+        format_in,
+        format_out,
+        level,
+        compression_rate,
+        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 original_size = input_bytes.len() as u64;
+    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);
+    let charge_units = !skip_charge && compressed_size < original_size;
+
+    if ctx.is_anonymous && charge_units {
+        let Some(session_id) = ctx.session_id.as_deref() else {
+            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 {
+            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 {
+            mark_file_failed(&state, task_id, file.id, &err.message).await?;
+            let _ = tokio::fs::remove_file(&input_path).await;
+            return Ok(());
+        }
+    }
+
+    let output_path = format!(
+        "{}/{}.{}",
+        state.config.storage_path,
+        file.id,
+        format_out.extension()
+    );
+    if let Err(err) = tokio::fs::write(&output_path, &compressed).await {
+        mark_file_failed(&state, task_id, file.id, "写入压缩文件失败").await?;
+        let _ = tokio::fs::remove_file(&input_path).await;
+        return Err(AppError::new(ErrorCode::StorageUnavailable, "写入压缩文件失败").with_source(err));
+    }
+
+    if is_task_cancelled(&state, task_id).await? {
+        let _ = tokio::fs::remove_file(&output_path).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,
+        &output_path,
+        original_size as i64,
+        compressed_size as i64,
+        saved_percent,
+        format_in,
+        format_out,
+        charge_units,
+    )
+    .await
+    {
+        if err.code == ErrorCode::QuotaExceeded {
+            let _ = tokio::fs::remove_file(&output_path).await;
+            mark_file_failed(&state, task_id, file.id, &err.message).await?;
+        } else {
+            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(())
+}
+
 async fn finalize_file(
     state: &AppState,
     billing_ctx: &Option<billing::BillingContext>,