perf: improve worker and storage throughput

src/worker/mod.rs

@@ -11,7 +11,7 @@ use redis::AsyncCommands;
 use sqlx::FromRow;
 use std::net::IpAddr;
 use std::sync::Arc;
-use std::time::Instant;
+use std::time::Duration;
 use tokio::sync::Semaphore;
 use tokio::task::JoinSet;
 use uuid::Uuid;
@@ -21,6 +21,13 @@ 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 QUEUE_BLOCK_MS: usize = 1_000;
+const MAINTENANCE_INTERVAL_SECONDS: u64 = 300;
+const MAINTENANCE_BATCH_SIZE: i64 = 1_000;
+const MAX_MAINTENANCE_BATCHES: usize = 20;
+const INITIAL_RETRY_SECONDS: u64 = 2;
+const MAX_RETRY_SECONDS: u64 = 30;
 
 pub async fn run(state: AppState) -> Result<(), AppError> {
     tracing::info!("Worker started");
@@ -29,21 +36,80 @@ pub async fn run(state: AppState) -> Result<(), AppError> {
 
     let consumer = format!("worker_{}", Uuid::new_v4());
     ensure_group(&state).await?;
+    tokio::spawn(maintenance_loop(state.clone()));
 
-    let mut last_maintenance = Instant::now();
+    let task_concurrency = state.config.worker_task_concurrency.max(1) as usize;
+    let mut inflight = JoinSet::new();
+    let mut poll_backoff = Duration::from_secs(INITIAL_RETRY_SECONDS);
+    tracing::info!(task_concurrency, "Worker task scheduler ready");
 
     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;
+        while let Some(result) = inflight.try_join_next() {
+            log_message_task_result(result);
         }
 
-        if last_maintenance.elapsed().as_secs() >= 300 {
-            if let Err(err) = maintenance(&state).await {
-                tracing::error!(error = ?err, "maintenance failed");
+        let available = task_concurrency.saturating_sub(inflight.len());
+        if available > 0 {
+            match read_messages(&state, &consumer, available).await {
+                Ok(messages) => {
+                    poll_backoff = Duration::from_secs(INITIAL_RETRY_SECONDS);
+                    for message in messages {
+                        let state = state.clone();
+                        let consumer = consumer.clone();
+                        inflight.spawn(async move {
+                            process_message_with_retries(state, consumer, message).await
+                        });
+                    }
+                }
+                Err(err) => {
+                    tracing::error!(
+                        error = ?err,
+                        retry_in_seconds = poll_backoff.as_secs(),
+                        "worker poll error"
+                    );
+                    tokio::time::sleep(poll_backoff).await;
+                    poll_backoff = next_backoff(poll_backoff);
+                    continue;
+                }
             }
-            last_maintenance = Instant::now();
         }
+
+        if inflight.len() >= task_concurrency {
+            if let Some(result) = inflight.join_next().await {
+                log_message_task_result(result);
+            }
+        } else {
+            tokio::select! {
+                result = inflight.join_next(), if !inflight.is_empty() => {
+                    if let Some(result) = result {
+                        log_message_task_result(result);
+                    }
+                }
+                _ = tokio::time::sleep(Duration::from_millis(25)) => {}
+            }
+        }
+    }
+}
+
+async fn maintenance_loop(state: AppState) {
+    let mut interval = tokio::time::interval_at(
+        tokio::time::Instant::now() + Duration::from_secs(MAINTENANCE_INTERVAL_SECONDS),
+        Duration::from_secs(MAINTENANCE_INTERVAL_SECONDS),
+    );
+    interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
+    loop {
+        interval.tick().await;
+        if let Err(err) = maintenance(&state).await {
+            tracing::error!(error = ?err, "maintenance failed");
+        }
+    }
+}
+
+fn log_message_task_result(result: Result<Result<(), AppError>, tokio::task::JoinError>) {
+    match result {
+        Ok(Ok(())) => {}
+        Ok(Err(err)) => tracing::error!(error = ?err, "worker message task stopped"),
+        Err(err) => tracing::error!(error = ?err, "worker message task panicked"),
     }
 }
 
@@ -71,51 +137,144 @@ async fn ensure_group(state: &AppState) -> Result<(), AppError> {
     }
 }
 
-async fn poll_once(state: &AppState, consumer: &str) -> Result<(), AppError> {
+async fn read_messages(
+    state: &AppState,
+    consumer: &str,
+    count: usize,
+) -> Result<Vec<StreamId>, 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;
+        return Ok(vec![msg]);
     }
 
-    // Retry messages already delivered to this consumer before taking new work.
-    let pending_opts = StreamReadOptions::default()
+    let opts = StreamReadOptions::default()
         .group(GROUP_NAME, consumer)
-        .count(1);
-    let mut reply: redis::streams::StreamReadReply = conn
-        .xread_options(&[STREAM_KEY], &["0"], &pending_opts)
+        .count(count)
+        .block(QUEUE_BLOCK_MS);
+    let reply: redis::streams::StreamReadReply = conn
+        .xread_options(&[STREAM_KEY], &[">"], &opts)
         .await
-        .map_err(|err| AppError::new(ErrorCode::Internal, "读取待重试任务失败").with_source(err))?;
+        .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);
+    Ok(reply
+        .keys
+        .into_iter()
+        .flat_map(|key| key.ids.into_iter())
+        .collect())
+}
 
-        reply = conn
-            .xread_options(&[STREAM_KEY], &[">"], &opts)
-            .await
-            .map_err(|err| AppError::new(ErrorCode::Internal, "读取队列失败").with_source(err))?;
-    }
+async fn process_message_with_retries(
+    state: AppState,
+    consumer: String,
+    mut message: StreamId,
+) -> Result<(), AppError> {
+    loop {
+        match handle_message_with_heartbeat(&state, &consumer, &message).await {
+            Ok(()) => return Ok(()),
+            Err(err) => {
+                let mut conn = state.redis.clone();
+                let deliveries = match pending_delivery_count(&mut conn, &message.id).await {
+                    Ok(deliveries) => deliveries,
+                    Err(count_err) => {
+                        tracing::warn!(
+                            message_id = %message.id,
+                            error = ?count_err,
+                            "failed to read delivery count; using initial retry delay"
+                        );
+                        1
+                    }
+                };
+                let delay = delivery_backoff(deliveries);
+                tracing::warn!(
+                    message_id = %message.id,
+                    deliveries,
+                    retry_in_seconds = delay.as_secs(),
+                    error = ?err,
+                    "worker message handling failed; retry scheduled"
+                );
+                tokio::time::sleep(delay).await;
 
-    if reply.keys.is_empty() {
-        return Ok(());
-    }
-
-    for key in reply.keys {
-        for msg in key.ids {
-            handle_message(state, &mut conn, msg).await?;
+                let mut reclaim_backoff = Duration::from_secs(INITIAL_RETRY_SECONDS);
+                loop {
+                    match redeliver_message(&mut conn, &consumer, &message.id).await {
+                        Ok(Some(redelivered)) => {
+                            message = redelivered;
+                            break;
+                        }
+                        Ok(None) => return Ok(()),
+                        Err(reclaim_err) => {
+                            tracing::error!(
+                                message_id = %message.id,
+                                retry_in_seconds = reclaim_backoff.as_secs(),
+                                error = ?reclaim_err,
+                                "failed to reclaim pending worker message"
+                            );
+                            tokio::time::sleep(reclaim_backoff).await;
+                            reclaim_backoff = next_backoff(reclaim_backoff);
+                        }
+                    }
+                }
+            }
         }
     }
+}
 
+async fn handle_message_with_heartbeat(
+    state: &AppState,
+    consumer: &str,
+    message: &StreamId,
+) -> Result<(), AppError> {
+    let mut conn = state.redis.clone();
+    let handling = handle_message(state, &mut conn, message);
+    tokio::pin!(handling);
+    let mut heartbeat = tokio::time::interval_at(
+        tokio::time::Instant::now() + Duration::from_secs(MESSAGE_HEARTBEAT_SECONDS),
+        Duration::from_secs(MESSAGE_HEARTBEAT_SECONDS),
+    );
+    heartbeat.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
+
+    loop {
+        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"
+                    );
+                }
+            }
+        }
+    }
+}
+
+async fn touch_pending_message(
+    state: &AppState,
+    consumer: &str,
+    message_id: &str,
+) -> Result<(), AppError> {
+    let mut conn = state.redis.clone();
+    redis::cmd("XCLAIM")
+        .arg(STREAM_KEY)
+        .arg(GROUP_NAME)
+        .arg(consumer)
+        .arg(0)
+        .arg(message_id)
+        .arg("JUSTID")
+        .query_async::<_, redis::Value>(&mut conn)
+        .await
+        .map_err(|err| {
+            AppError::new(ErrorCode::Internal, "刷新队列任务心跳失败").with_source(err)
+        })?;
     Ok(())
 }
 
 async fn handle_message(
     state: &AppState,
     conn: &mut redis::aio::ConnectionManager,
-    msg: StreamId,
+    msg: &StreamId,
 ) -> Result<(), AppError> {
     let Some(task_id_str) = msg.get::<String>("task_id") else {
         ack_message(conn, &msg.id).await?;
@@ -147,17 +306,25 @@ async fn handle_message(
                 return Ok(());
             }
 
-            tracing::warn!(
-                task_id = %task_id,
-                deliveries,
-                error = %err,
-                "task processing failed; message left pending for retry"
-            );
             Err(err)
         }
     }
 }
 
+async fn redeliver_message(
+    conn: &mut redis::aio::ConnectionManager,
+    consumer: &str,
+    message_id: &str,
+) -> Result<Option<StreamId>, AppError> {
+    let claimed: StreamClaimReply = conn
+        .xclaim(STREAM_KEY, GROUP_NAME, consumer, 0, &[message_id])
+        .await
+        .map_err(|err| {
+            AppError::new(ErrorCode::Internal, "重新认领待重试任务失败").with_source(err)
+        })?;
+    Ok(claimed.ids.into_iter().next())
+}
+
 async fn claim_stale_message(
     conn: &mut redis::aio::ConnectionManager,
     consumer: &str,
@@ -212,6 +379,21 @@ fn should_dead_letter(deliveries: usize) -> bool {
     deliveries >= MAX_DELIVERIES
 }
 
+fn delivery_backoff(deliveries: usize) -> Duration {
+    let exponent = deliveries.saturating_sub(1).min(4) as u32;
+    Duration::from_secs(
+        INITIAL_RETRY_SECONDS
+            .saturating_mul(2_u64.saturating_pow(exponent))
+            .min(MAX_RETRY_SECONDS),
+    )
+}
+
+fn next_backoff(current: Duration) -> Duration {
+    current
+        .saturating_mul(2)
+        .min(Duration::from_secs(MAX_RETRY_SECONDS))
+}
+
 async fn write_dead_letter(
     conn: &mut redis::aio::ConnectionManager,
     message_id: &str,
@@ -1026,22 +1208,38 @@ async fn maintenance(state: &AppState) -> Result<(), AppError> {
 }
 
 async fn settle_finished_anonymous_reservations(state: &AppState) -> Result<(), AppError> {
-    let task_ids: Vec<Uuid> = sqlx::query_scalar(
-        r#"
-        SELECT id
-        FROM tasks
-        WHERE anonymous_units_reserved > 0
-          AND status IN ('completed', 'failed', 'cancelled')
-        ORDER BY completed_at ASC NULLS FIRST
-        LIMIT 200
-        "#,
-    )
-    .fetch_all(&state.db)
-    .await
-    .map_err(|err| AppError::new(ErrorCode::Internal, "查询待结算匿名任务失败").with_source(err))?;
+    for _ in 0..MAX_MAINTENANCE_BATCHES {
+        let task_ids: Vec<Uuid> = sqlx::query_scalar(
+            r#"
+            SELECT id
+            FROM tasks
+            WHERE anonymous_units_reserved > 0
+              AND status IN ('completed', 'failed', 'cancelled')
+            ORDER BY completed_at ASC NULLS FIRST
+            LIMIT $1
+            "#,
+        )
+        .bind(MAINTENANCE_BATCH_SIZE)
+        .fetch_all(&state.db)
+        .await
+        .map_err(|err| {
+            AppError::new(ErrorCode::Internal, "查询待结算匿名任务失败").with_source(err)
+        })?;
 
-    for task_id in task_ids {
-        quota::settle_anonymous_task_reservation(state, task_id).await?;
+        let batch_len = task_ids.len();
+        let mut settled = 0usize;
+        for task_id in task_ids {
+            match quota::settle_anonymous_task_reservation(state, task_id).await {
+                Ok(_) => settled += 1,
+                Err(err) => {
+                    tracing::warn!(task_id = %task_id, error = %err, "anonymous quota settlement deferred")
+                }
+            }
+        }
+        if settled == 0 || batch_len < MAINTENANCE_BATCH_SIZE as usize {
+            break;
+        }
+        tokio::task::yield_now().await;
     }
     Ok(())
 }
@@ -1117,20 +1315,34 @@ async fn cleanup_expired_records(state: &AppState) -> Result<(), AppError> {
 }
 
 async fn cleanup_expired_tasks(state: &AppState) -> Result<(), AppError> {
-    let task_ids: Vec<Uuid> =
-        sqlx::query_scalar("SELECT id FROM tasks WHERE expires_at < NOW() LIMIT 200")
-            .fetch_all(&state.db)
-            .await
-            .unwrap_or_default();
+    for _ in 0..MAX_MAINTENANCE_BATCHES {
+        let task_ids: Vec<Uuid> = sqlx::query_scalar(
+            "SELECT id FROM tasks WHERE expires_at < NOW() ORDER BY expires_at ASC LIMIT $1",
+        )
+        .bind(MAINTENANCE_BATCH_SIZE)
+        .fetch_all(&state.db)
+        .await
+        .map_err(|err| AppError::new(ErrorCode::Internal, "查询过期任务失败").with_source(err))?;
 
-    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");
+        if task_ids.is_empty() {
+            break;
         }
+
+        let batch_len = task_ids.len();
+        let mut cleaned = 0usize;
+        for task_id in task_ids {
+            match cleanup_expired_task(state, task_id).await {
+                Ok(()) => cleaned += 1,
+                Err(err) => {
+                    tracing::warn!(task_id = %task_id, error = %err, "expired task cleanup deferred")
+                }
+            }
+        }
+
+        if cleaned == 0 || batch_len < MAINTENANCE_BATCH_SIZE as usize {
+            break;
+        }
+        tokio::task::yield_now().await;
     }
 
     Ok(())
@@ -1232,4 +1444,21 @@ mod tests {
         assert!(should_dead_letter(3));
         assert!(should_dead_letter(10));
     }
+
+    #[test]
+    fn retry_backoff_is_exponential_and_capped() {
+        assert_eq!(delivery_backoff(1), Duration::from_secs(2));
+        assert_eq!(delivery_backoff(2), Duration::from_secs(4));
+        assert_eq!(delivery_backoff(3), Duration::from_secs(8));
+        assert_eq!(delivery_backoff(10), Duration::from_secs(30));
+        assert_eq!(
+            next_backoff(Duration::from_secs(30)),
+            Duration::from_secs(30)
+        );
+    }
+
+    #[test]
+    fn message_heartbeat_precedes_stale_claim_threshold() {
+        assert!(MESSAGE_HEARTBEAT_SECONDS * 1_000 < STALE_MESSAGE_IDLE_MS as u64);
+    }
 }