From fbc82dfa070db795ce21adab39fd6c0c3d42448c Mon Sep 17 00:00:00 2001 From: 237899745 <237899745@users.noreply.git.workyai.cn> Date: Sun, 26 Jul 2026 06:46:37 +0800 Subject: [PATCH] fix: keep ZIP builds alive after request cancellation --- docker/nginx.conf | 2 + src/api/downloads.rs | 270 +++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 260 insertions(+), 12 deletions(-) diff --git a/docker/nginx.conf b/docker/nginx.conf index 8ba5a42..265b3e4 100644 --- a/docker/nginx.conf +++ b/docker/nginx.conf @@ -40,6 +40,8 @@ http { location /downloads/ { proxy_pass http://imageforge_api; + proxy_read_timeout 300s; + proxy_send_timeout 300s; proxy_set_header Host $host; proxy_set_header X-Real-IP $remote_addr; proxy_set_header X-Forwarded-For $remote_addr; diff --git a/src/api/downloads.rs b/src/api/downloads.rs index e0e76c5..ad44a4f 100644 --- a/src/api/downloads.rs +++ b/src/api/downloads.rs @@ -15,6 +15,7 @@ use sqlx::FromRow; use std::collections::HashMap; use std::net::SocketAddr; use std::path::PathBuf; +use tokio::sync::oneshot; use tokio_util::io::ReaderStream; use uuid::Uuid; @@ -234,7 +235,7 @@ struct TaskZipRow { zip_storage_key: Option, } -#[derive(Debug, FromRow)] +#[derive(Clone, Debug, FromRow)] struct TaskZipFileRow { storage_backend: String, storage_endpoint_id: Option, @@ -259,8 +260,10 @@ enum ZipBuildClaim { Busy, } -const ZIP_BUILD_LEASE_SECONDS: i64 = 15 * 60; -const ZIP_BUILD_WAIT_SECONDS: u64 = 30; +const ZIP_BUILD_LEASE_SECONDS: i64 = 60; +const ZIP_BUILD_HEARTBEAT_SECONDS: u64 = 20; +const ZIP_BUILD_WAIT_SECONDS: u64 = 300; +const ZIP_BUILD_GATE_WAIT_SECONDS: u64 = 60; async fn download_task_zip( State(state): State, @@ -380,7 +383,16 @@ async fn resolve_task_zip( match claim_zip_build(state, task_id).await? { ZipBuildClaim::Cached(object) => return Ok(object), ZipBuildClaim::Acquired { token } => { - return build_claimed_zip(state, task_id, retention_hours, token, &rows).await; + let receiver = spawn_claimed_zip_build( + state.clone(), + task_id, + retention_hours, + token, + rows.clone(), + ); + return receiver.await.map_err(|err| { + AppError::new(ErrorCode::Internal, "后台 ZIP 构建任务异常退出").with_source(err) + })?; } ZipBuildClaim::Busy if tokio::time::Instant::now() < deadline => { tokio::time::sleep(std::time::Duration::from_millis(100)).await; @@ -486,6 +498,15 @@ async fn claim_zip_build(state: &AppState, task_id: Uuid) -> Result Result<(), AppError> { + renew_zip_build_for(state, task_id, token, ZIP_BUILD_LEASE_SECONDS).await +} + +async fn renew_zip_build_for( + state: &AppState, + task_id: Uuid, + token: Uuid, + lease_seconds: i64, +) -> Result<(), AppError> { let updated = sqlx::query( r#" UPDATE tasks @@ -498,7 +519,7 @@ async fn renew_zip_build(state: &AppState, task_id: Uuid, token: Uuid) -> Result ) .bind(task_id) .bind(token) - .bind(ZIP_BUILD_LEASE_SECONDS) + .bind(lease_seconds) .execute(&state.db) .await .map_err(|err| AppError::new(ErrorCode::Internal, "续租 ZIP 构建失败").with_source(err))?; @@ -511,6 +532,98 @@ async fn renew_zip_build(state: &AppState, task_id: Uuid, token: Uuid) -> Result Ok(()) } +struct ZipHeartbeatStop(Option>); + +impl Drop for ZipHeartbeatStop { + fn drop(&mut self) { + if let Some(stop) = self.0.take() { + let _ = stop.send(()); + } + } +} + +fn start_zip_build_heartbeat( + state: AppState, + task_id: Uuid, + token: Uuid, + interval: std::time::Duration, + lease_seconds: i64, +) -> (ZipHeartbeatStop, tokio::task::JoinHandle<()>) { + let (stop_tx, mut stop_rx) = oneshot::channel(); + let handle = tokio::spawn(async move { + let mut ticker = tokio::time::interval(interval); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + ticker.tick().await; + loop { + tokio::select! { + _ = &mut stop_rx => break, + _ = ticker.tick() => { + if let Err(err) = renew_zip_build_for( + &state, + task_id, + token, + lease_seconds, + ) + .await + { + tracing::warn!(task_id = %task_id, zip_build_token = %token, error = %err, "failed to heartbeat ZIP build lease"); + } + } + } + } + }); + (ZipHeartbeatStop(Some(stop_tx)), handle) +} + +fn spawn_claimed_zip_build( + state: AppState, + task_id: Uuid, + retention_hours: i64, + token: Uuid, + rows: Vec, +) -> oneshot::Receiver> { + let (result_tx, result_rx) = oneshot::channel(); + tokio::spawn(async move { + let result = run_claimed_zip_build(state, task_id, retention_hours, token, rows).await; + let _ = result_tx.send(result); + }); + result_rx +} + +async fn run_claimed_zip_build( + state: AppState, + task_id: Uuid, + retention_hours: i64, + token: Uuid, + rows: Vec, +) -> Result { + let (heartbeat_stop, heartbeat) = start_zip_build_heartbeat( + state.clone(), + task_id, + token, + std::time::Duration::from_secs(ZIP_BUILD_HEARTBEAT_SECONDS), + ZIP_BUILD_LEASE_SECONDS, + ); + let build_state = state.clone(); + let build = tokio::spawn(async move { + build_claimed_zip(&build_state, task_id, retention_hours, token, &rows).await + }); + let result = match build.await { + Ok(result) => result, + Err(err) => { + let temp_dir = zip_temp_dir(&state, task_id, token); + let _ = tokio::fs::remove_dir_all(temp_dir).await; + release_zip_build(&state, task_id, token).await; + Err(AppError::new(ErrorCode::Internal, "ZIP 构建任务异常退出").with_source(err)) + } + }; + drop(heartbeat_stop); + if let Err(err) = heartbeat.await { + tracing::warn!(task_id = %task_id, zip_build_token = %token, error = %err, "ZIP lease heartbeat task failed"); + } + result +} + async fn release_zip_build(state: &AppState, task_id: Uuid, token: Uuid) { if let Err(err) = sqlx::query( r#" @@ -539,7 +652,7 @@ async fn build_claimed_zip( rows: &[TaskZipFileRow], ) -> Result { let permit = match tokio::time::timeout( - std::time::Duration::from_secs(ZIP_BUILD_WAIT_SECONDS), + std::time::Duration::from_secs(ZIP_BUILD_GATE_WAIT_SECONDS), state.zip_build_semaphore.clone().acquire_owned(), ) .await @@ -558,10 +671,7 @@ async fn build_claimed_zip( } }; - let temp_dir = PathBuf::from(format!( - "{}/tmp/zips/{task_id}-{token}", - state.config.storage_path - )); + let temp_dir = zip_temp_dir(state, task_id, token); let zip_path = temp_dir.join(format!("task_{task_id}.zip")); let build_result = build_zip_attempt( state, @@ -586,6 +696,13 @@ async fn build_claimed_zip( publish_zip_attempt(state, task_id, token, stored).await } +fn zip_temp_dir(state: &AppState, task_id: Uuid, token: Uuid) -> PathBuf { + PathBuf::from(&state.config.storage_path) + .join("tmp") + .join("zips") + .join(format!("{task_id}-{token}")) +} + async fn build_zip_attempt( state: &AppState, task_id: Uuid, @@ -1159,13 +1276,142 @@ mod tests { .await; assert!(orphan_read.is_err(), "unpublished ZIP object was orphaned"); - for locator in [&locators[0], &takeover_locator] { + let cancelled_task = Uuid::new_v4(); + insert_zip_task( + &pool, + cancelled_task, + &format!("{marker}-cancelled-request"), + &input_path, + 21, + ) + .await; + let gate_one = state + .zip_build_semaphore + .clone() + .acquire_owned() + .await + .expect("acquire first ZIP gate permit"); + let gate_two = state + .zip_build_semaphore + .clone() + .acquire_owned() + .await + .expect("acquire second ZIP gate permit"); + let cancelled_state = state.clone(); + let cancelled_request = + tokio::spawn( + async move { resolve_task_zip(&cancelled_state, cancelled_task, 24).await }, + ); + let mut cancelled_token = None; + for _ in 0..100 { + cancelled_token = sqlx::query_scalar("SELECT zip_build_token FROM tasks WHERE id = $1") + .bind(cancelled_task) + .fetch_one(&pool) + .await + .expect("query cancelled-request ZIP token"); + if cancelled_token.is_some() { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + let cancelled_token = cancelled_token.expect("background ZIP builder never acquired lease"); + cancelled_request.abort(); + assert!( + cancelled_request + .await + .expect_err("aborted ZIP request completed") + .is_cancelled(), + "ZIP request was not cancelled" + ); + drop(gate_one); + drop(gate_two); + + let mut cancelled_locator = None; + for _ in 0..200 { + cancelled_locator = load_published_zip(&state, cancelled_task) + .await + .expect("query background ZIP publication"); + if cancelled_locator.is_some() { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + let cancelled_locator = + cancelled_locator.expect("request cancellation stopped the background ZIP build"); + let cancelled_state_row: (Option, i64) = + sqlx::query_as("SELECT zip_build_token, zip_build_attempt FROM tasks WHERE id = $1") + .bind(cancelled_task) + .fetch_one(&pool) + .await + .expect("query background ZIP terminal state"); + assert_eq!(cancelled_state_row, (None, 1)); + assert!( + !tokio::fs::try_exists(zip_temp_dir(&state, cancelled_task, cancelled_token)) + .await + .expect("check cancelled-request ZIP temp directory"), + "background ZIP temp directory survived publication" + ); + + let heartbeat_task = Uuid::new_v4(); + insert_zip_task( + &pool, + heartbeat_task, + &format!("{marker}-heartbeat"), + &input_path, + 21, + ) + .await; + let heartbeat_token = match claim_zip_build(&state, heartbeat_task) + .await + .expect("claim heartbeat ZIP builder") + { + ZipBuildClaim::Acquired { token } => token, + other => panic!("unexpected heartbeat ZIP claim: {other:?}"), + }; + sqlx::query( + "UPDATE tasks SET zip_build_lease_until = NOW() + INTERVAL '1 second' WHERE id = $1 AND zip_build_token = $2", + ) + .bind(heartbeat_task) + .bind(heartbeat_token) + .execute(&pool) + .await + .expect("shorten heartbeat ZIP lease"); + let (heartbeat_stop, heartbeat) = start_zip_build_heartbeat( + state.clone(), + heartbeat_task, + heartbeat_token, + std::time::Duration::from_millis(100), + 1, + ); + tokio::time::sleep(std::time::Duration::from_millis(1_500)).await; + assert!( + matches!( + claim_zip_build(&state, heartbeat_task) + .await + .expect("probe heartbeat ZIP lease"), + ZipBuildClaim::Busy + ), + "heartbeat did not fence a takeover after the original lease duration" + ); + drop(heartbeat_stop); + heartbeat.await.expect("join ZIP heartbeat task"); + release_zip_build(&state, heartbeat_task, heartbeat_token).await; + + for locator in [&locators[0], &takeover_locator, &cancelled_locator] { storage::delete_object(&state, locator) .await .expect("delete published ZIP test object"); } sqlx::query("DELETE FROM tasks WHERE id = ANY($1)") - .bind(&[task_id, over_budget_task, takeover_task][..]) + .bind( + &[ + task_id, + over_budget_task, + takeover_task, + cancelled_task, + heartbeat_task, + ][..], + ) .execute(&pool) .await .expect("delete ZIP test tasks");