Compare commits
2 Commits
380e89b058
...
fbc82dfa07
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
fbc82dfa07 | ||
|
|
4d0e8aa70a |
@@ -33,7 +33,7 @@ jobs:
|
||||
IMAGEFORGE_TEST_REDIS_URL: redis://redis:6379/
|
||||
JWT_SECRET: imageforge-ci-jwt-secret
|
||||
API_KEY_PEPPER: imageforge-ci-api-key-pepper
|
||||
EXPECTED_EXTERNAL_TESTS: '9'
|
||||
EXPECTED_EXTERNAL_TESTS: '10'
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v4
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
#!/usr/bin/env bash
|
||||
set -Eeuo pipefail
|
||||
|
||||
EXPECTED_EXTERNAL_TESTS="${EXPECTED_EXTERNAL_TESTS:-9}"
|
||||
EXPECTED_EXTERNAL_TESTS="${EXPECTED_EXTERNAL_TESTS:-10}"
|
||||
WORK_DIR="$(mktemp -d)"
|
||||
POSTGRES_CONTAINER=""
|
||||
REDIS_CONTAINER=""
|
||||
|
||||
@@ -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<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
#[derive(Clone, Debug, FromRow)]
|
||||
struct TaskZipFileRow {
|
||||
storage_backend: String,
|
||||
storage_endpoint_id: Option<Uuid>,
|
||||
@@ -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<AppState>,
|
||||
@@ -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<ZipBuildClai
|
||||
}
|
||||
|
||||
async fn renew_zip_build(state: &AppState, task_id: Uuid, token: Uuid) -> 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<oneshot::Sender<()>>);
|
||||
|
||||
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<TaskZipFileRow>,
|
||||
) -> oneshot::Receiver<Result<storage::ObjectLocator, AppError>> {
|
||||
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<TaskZipFileRow>,
|
||||
) -> Result<storage::ObjectLocator, AppError> {
|
||||
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<storage::ObjectLocator, AppError> {
|
||||
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<Uuid>, 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");
|
||||
|
||||
374
src/api/tasks.rs
374
src/api/tasks.rs
@@ -21,6 +21,7 @@ use serde::{Deserialize, Serialize};
|
||||
use sha2::{Digest, Sha256};
|
||||
use sqlx::FromRow;
|
||||
use std::net::SocketAddr;
|
||||
use std::path::PathBuf;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use uuid::Uuid;
|
||||
|
||||
@@ -115,16 +116,13 @@ async fn create_batch_task(
|
||||
let (files, opts, request_hash) = match parsed {
|
||||
Ok(parsed) => parsed,
|
||||
Err(err) => {
|
||||
let base_dir = format!("{}/orig/{task_id}", state.config.storage_path);
|
||||
let _ = tokio::fs::remove_dir_all(base_dir).await;
|
||||
cleanup_task_input_dir(&state, task_id).await;
|
||||
return Err(err);
|
||||
}
|
||||
};
|
||||
|
||||
if files.is_empty() {
|
||||
cleanup_file_paths(&files).await;
|
||||
let base_dir = format!("{}/orig/{task_id}", state.config.storage_path);
|
||||
let _ = tokio::fs::remove_dir_all(base_dir).await;
|
||||
cleanup_task_input_dir(&state, task_id).await;
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "缺少 files[]"));
|
||||
}
|
||||
|
||||
@@ -141,13 +139,13 @@ async fn create_batch_task(
|
||||
let begin_result = match begin_result {
|
||||
Ok(result) => result,
|
||||
Err(err) => {
|
||||
cleanup_file_paths(&files).await;
|
||||
cleanup_task_input_dir(&state, task_id).await;
|
||||
return Err(err);
|
||||
}
|
||||
};
|
||||
match begin_result {
|
||||
idempotency::BeginResult::Replay { response_body, .. } => {
|
||||
cleanup_file_paths(&files).await;
|
||||
cleanup_task_input_dir(&state, task_id).await;
|
||||
let resp: BatchCreateResponse =
|
||||
serde_json::from_value(response_body).map_err(|err| {
|
||||
AppError::new(ErrorCode::Internal, "幂等结果解析失败").with_source(err)
|
||||
@@ -161,7 +159,7 @@ async fn create_batch_task(
|
||||
));
|
||||
}
|
||||
idempotency::BeginResult::InProgress => {
|
||||
cleanup_file_paths(&files).await;
|
||||
cleanup_task_input_dir(&state, task_id).await;
|
||||
if let Some((_status, body)) =
|
||||
idempotency::wait_for_replay(&state, scope, idem_key, &request_hash, 10_000)
|
||||
.await?
|
||||
@@ -191,6 +189,9 @@ async fn create_batch_task(
|
||||
|
||||
let mut anonymous_reserved_units = 0u32;
|
||||
let mut anonymous_quota_date = None;
|
||||
let mut task_persisted = false;
|
||||
let mut enqueue_failure_finalized = false;
|
||||
let mut cleanup_inputs_on_error = true;
|
||||
let create_result: Result<BatchCreateResponse, AppError> = (async {
|
||||
match &admission.task_owner {
|
||||
TaskOwner::Anonymous { session_id } => {
|
||||
@@ -301,14 +302,21 @@ async fn create_batch_task(
|
||||
tx.commit()
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "提交事务失败").with_source(err))?;
|
||||
task_persisted = true;
|
||||
|
||||
if let Err(err) = enqueue_task(&state, task_id).await {
|
||||
let _ =
|
||||
sqlx::query("UPDATE tasks SET status = 'failed', error_message = $2 WHERE id = $1")
|
||||
.bind(task_id)
|
||||
.bind("队列提交失败")
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
match finalize_enqueue_failure(&state, task_id, "队列提交失败").await {
|
||||
Ok(true) => enqueue_failure_finalized = true,
|
||||
Ok(false) => {
|
||||
// XADD may have succeeded even if the client saw an error. A worker
|
||||
// that already claimed the task owns both the input and settlement.
|
||||
cleanup_inputs_on_error = false;
|
||||
}
|
||||
Err(finalize_err) => {
|
||||
cleanup_inputs_on_error = false;
|
||||
tracing::error!(task_id = %task_id, error = %finalize_err, "failed to finalize task after queue submission error");
|
||||
}
|
||||
}
|
||||
return Err(err);
|
||||
}
|
||||
|
||||
@@ -347,11 +355,19 @@ async fn create_batch_task(
|
||||
Err(err) => {
|
||||
if anonymous_reserved_units > 0 {
|
||||
if let context::Principal::Anonymous { session_id } = &principal {
|
||||
let settlement =
|
||||
quota::settle_anonymous_task_reservation(&state, task_id).await;
|
||||
match settlement {
|
||||
Ok(Some(_)) => {}
|
||||
Ok(None) => {
|
||||
let should_refund_directly = if enqueue_failure_finalized {
|
||||
match quota::settle_anonymous_task_reservation(&state, task_id).await {
|
||||
Ok(Some(_)) => false,
|
||||
Ok(None) => true,
|
||||
Err(settle_err) => {
|
||||
tracing::warn!(task_id = %task_id, error = %settle_err, "failed to settle anonymous batch admission");
|
||||
false
|
||||
}
|
||||
}
|
||||
} else {
|
||||
!task_persisted
|
||||
};
|
||||
if should_refund_directly {
|
||||
if let Some(date) = anonymous_quota_date {
|
||||
if let Err(refund_err) = quota::refund_anonymous_reservation_once(
|
||||
&state,
|
||||
@@ -367,10 +383,6 @@ async fn create_batch_task(
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(settle_err) => {
|
||||
tracing::warn!(task_id = %task_id, error = %settle_err, "failed to settle anonymous batch admission");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if let (Some(scope), Some(idem_key)) = (idempotency_scope, idempotency_key.as_deref()) {
|
||||
@@ -378,7 +390,9 @@ async fn create_batch_task(
|
||||
let _ = idempotency::abort(&state, scope, idem_key, &request_hash).await;
|
||||
}
|
||||
}
|
||||
cleanup_file_paths(&files).await;
|
||||
if cleanup_inputs_on_error {
|
||||
cleanup_task_input_dir(&state, task_id).await;
|
||||
}
|
||||
Err(err)
|
||||
}
|
||||
}
|
||||
@@ -489,6 +503,88 @@ async fn enqueue_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn finalize_enqueue_failure(
|
||||
state: &AppState,
|
||||
task_id: Uuid,
|
||||
error_message: &str,
|
||||
) -> Result<bool, AppError> {
|
||||
let mut tx = state.db.begin().await.map_err(|err| {
|
||||
AppError::new(ErrorCode::Internal, "开启队列失败收口事务失败").with_source(err)
|
||||
})?;
|
||||
let task: Option<(String, i32)> =
|
||||
sqlx::query_as("SELECT status::text, total_files FROM tasks WHERE id = $1 FOR UPDATE")
|
||||
.bind(task_id)
|
||||
.fetch_optional(&mut *tx)
|
||||
.await
|
||||
.map_err(|err| {
|
||||
AppError::new(ErrorCode::Internal, "锁定队列失败任务失败").with_source(err)
|
||||
})?;
|
||||
let Some((status, total_files)) = task else {
|
||||
tx.rollback().await.ok();
|
||||
return Ok(false);
|
||||
};
|
||||
if status != "pending" {
|
||||
tx.rollback().await.ok();
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
let files = sqlx::query(
|
||||
r#"
|
||||
UPDATE task_files
|
||||
SET status = 'failed',
|
||||
error_message = $2,
|
||||
completed_at = NOW(),
|
||||
input_path = NULL,
|
||||
storage_path = NULL,
|
||||
lease_owner = NULL,
|
||||
lease_until = NULL
|
||||
WHERE task_id = $1
|
||||
AND status = 'pending'
|
||||
"#,
|
||||
)
|
||||
.bind(task_id)
|
||||
.bind(error_message)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "终结未入队文件失败").with_source(err))?;
|
||||
if files.rows_affected() != u64::try_from(total_files.max(0)).unwrap_or(0) {
|
||||
tx.rollback().await.ok();
|
||||
return Err(AppError::new(
|
||||
ErrorCode::Internal,
|
||||
"未入队任务的文件状态不一致",
|
||||
));
|
||||
}
|
||||
|
||||
let task = sqlx::query(
|
||||
r#"
|
||||
UPDATE tasks
|
||||
SET status = 'failed',
|
||||
error_message = $2,
|
||||
completed_at = NOW(),
|
||||
completed_files = 0,
|
||||
failed_files = total_files,
|
||||
lease_owner = NULL,
|
||||
lease_until = NULL
|
||||
WHERE id = $1
|
||||
AND status = 'pending'
|
||||
"#,
|
||||
)
|
||||
.bind(task_id)
|
||||
.bind(error_message)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "终结未入队任务失败").with_source(err))?;
|
||||
if task.rows_affected() != 1 {
|
||||
tx.rollback().await.ok();
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
tx.commit().await.map_err(|err| {
|
||||
AppError::new(ErrorCode::Internal, "提交队列失败收口事务失败").with_source(err)
|
||||
})?;
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn parse_batch_request(
|
||||
state: &AppState,
|
||||
task_id: Uuid,
|
||||
@@ -1173,10 +1269,232 @@ async fn cleanup_file_paths(files: &[BatchFileInput]) {
|
||||
.first()
|
||||
.and_then(|file| std::path::Path::new(&file.storage_path).parent())
|
||||
.map(std::path::Path::to_path_buf);
|
||||
for f in files {
|
||||
let _ = tokio::fs::remove_file(&f.storage_path).await;
|
||||
}
|
||||
if let Some(parent) = parent {
|
||||
let _ = tokio::fs::remove_dir(parent).await;
|
||||
let _ = tokio::fs::remove_dir_all(parent).await;
|
||||
}
|
||||
}
|
||||
|
||||
fn task_input_dir(state: &AppState, task_id: Uuid) -> PathBuf {
|
||||
PathBuf::from(&state.config.storage_path)
|
||||
.join("orig")
|
||||
.join(task_id.to_string())
|
||||
}
|
||||
|
||||
async fn cleanup_task_input_dir(state: &AppState, task_id: Uuid) {
|
||||
let path = task_input_dir(state, task_id);
|
||||
match tokio::fs::remove_dir_all(&path).await {
|
||||
Ok(()) => {}
|
||||
Err(err) if err.kind() == std::io::ErrorKind::NotFound => {}
|
||||
Err(err) => {
|
||||
tracing::warn!(task_id = %task_id, path = %path.display(), error = %err, "failed to clean task input directory");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::config::Config;
|
||||
use crate::services::mail::Mailer;
|
||||
use crate::services::settings::RuntimePolicyCache;
|
||||
use crate::services::storage::StorageCache;
|
||||
use crate::worker::TaskProcessOutcome;
|
||||
use sqlx::postgres::PgPoolOptions;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
async fn build_batch_test_state(
|
||||
pool: sqlx::PgPool,
|
||||
database_url: String,
|
||||
redis_url: String,
|
||||
storage_path: String,
|
||||
) -> AppState {
|
||||
let mut config = Config::from_env().expect("load batch test config");
|
||||
config.database_url = database_url;
|
||||
config.redis_url = redis_url;
|
||||
config.storage_path = storage_path;
|
||||
config.mail_enabled = false;
|
||||
config.mail_log_links_when_disabled = false;
|
||||
let redis = redis::Client::open(config.redis_url.clone())
|
||||
.expect("create batch test Redis client")
|
||||
.get_connection_manager()
|
||||
.await
|
||||
.expect("connect batch test Redis");
|
||||
AppState {
|
||||
mailer: Arc::new(Mailer::new(&config).expect("create disabled batch test mailer")),
|
||||
image_processing_semaphore: Arc::new(Semaphore::new(2)),
|
||||
zip_build_semaphore: Arc::new(Semaphore::new(1)),
|
||||
runtime_policy_cache: RuntimePolicyCache::new(),
|
||||
storage_cache: StorageCache::new(),
|
||||
config,
|
||||
db: pool,
|
||||
redis,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[ignore = "requires isolated IMAGEFORGE_TEST_DATABASE_URL and IMAGEFORGE_TEST_REDIS_URL"]
|
||||
async fn enqueue_failure_finalizes_files_and_removes_exact_input_directory() {
|
||||
let database_url = std::env::var("IMAGEFORGE_TEST_DATABASE_URL")
|
||||
.expect("IMAGEFORGE_TEST_DATABASE_URL must be set");
|
||||
assert!(
|
||||
database_url.to_ascii_lowercase().contains("test"),
|
||||
"refusing to run destructive integration test outside a test database"
|
||||
);
|
||||
let redis_url = std::env::var("IMAGEFORGE_TEST_REDIS_URL")
|
||||
.expect("IMAGEFORGE_TEST_REDIS_URL must be set");
|
||||
let pool = PgPoolOptions::new()
|
||||
.max_connections(8)
|
||||
.connect(&database_url)
|
||||
.await
|
||||
.expect("connect batch test database");
|
||||
sqlx::migrate!().run(&pool).await.expect("run migrations");
|
||||
|
||||
let marker = Uuid::new_v4().simple().to_string();
|
||||
let storage_root = std::env::temp_dir().join(format!("imageforge-batch-test-{marker}"));
|
||||
let state = build_batch_test_state(
|
||||
pool.clone(),
|
||||
database_url,
|
||||
redis_url,
|
||||
storage_root.to_string_lossy().to_string(),
|
||||
)
|
||||
.await;
|
||||
let task_id = Uuid::new_v4();
|
||||
let input_dir = task_input_dir(&state, task_id);
|
||||
tokio::fs::create_dir_all(&input_dir)
|
||||
.await
|
||||
.expect("create batch input directory");
|
||||
let first_path = input_dir.join("first.png");
|
||||
let second_path = input_dir.join("second.png");
|
||||
tokio::fs::write(&first_path, b"first")
|
||||
.await
|
||||
.expect("write first input");
|
||||
tokio::fs::write(&second_path, b"second")
|
||||
.await
|
||||
.expect("write second input");
|
||||
tokio::fs::write(input_dir.join("interrupted.upload"), b"partial")
|
||||
.await
|
||||
.expect("write interrupted upload fixture");
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO tasks (
|
||||
id, session_id, client_ip, status, total_files,
|
||||
total_original_size, expires_at, retention_hours
|
||||
) VALUES ($1, $2, '127.0.0.1'::inet, 'pending', 2, 11, NOW() + INTERVAL '1 day', 24)
|
||||
"#,
|
||||
)
|
||||
.bind(task_id)
|
||||
.bind(format!("batch-session-{marker}"))
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("insert pending batch task");
|
||||
for (name, path, size) in [
|
||||
("first.png", &first_path, 5_i64),
|
||||
("second.png", &second_path, 6_i64),
|
||||
] {
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO task_files (
|
||||
id, task_id, original_name, original_format, output_format,
|
||||
original_size, input_path, status
|
||||
) VALUES ($1, $2, $3, 'png', 'png', $4, $5, 'pending')
|
||||
"#,
|
||||
)
|
||||
.bind(Uuid::new_v4())
|
||||
.bind(task_id)
|
||||
.bind(name)
|
||||
.bind(size)
|
||||
.bind(path.to_string_lossy().to_string())
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("insert pending batch file");
|
||||
}
|
||||
|
||||
let mut redis = state.redis.clone();
|
||||
let _: i64 = redis::cmd("DEL")
|
||||
.arg("stream:compress_jobs")
|
||||
.query_async(&mut redis)
|
||||
.await
|
||||
.expect("clear compression stream");
|
||||
let _: () = redis::cmd("SET")
|
||||
.arg("stream:compress_jobs")
|
||||
.arg("wrong-type-fixture")
|
||||
.query_async(&mut redis)
|
||||
.await
|
||||
.expect("install WRONGTYPE fixture");
|
||||
|
||||
let enqueue_error = enqueue_task(&state, task_id)
|
||||
.await
|
||||
.expect_err("XADD unexpectedly accepted a string key");
|
||||
assert_eq!(enqueue_error.code, ErrorCode::Internal);
|
||||
assert!(
|
||||
finalize_enqueue_failure(&state, task_id, "队列提交失败")
|
||||
.await
|
||||
.expect("finalize enqueue failure"),
|
||||
"pending task was not finalized"
|
||||
);
|
||||
cleanup_task_input_dir(&state, task_id).await;
|
||||
|
||||
let task: (String, bool, i32, i32) = sqlx::query_as(
|
||||
r#"
|
||||
SELECT status::text, completed_at IS NOT NULL, completed_files, failed_files
|
||||
FROM tasks WHERE id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(task_id)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("query finalized task");
|
||||
assert_eq!(task, ("failed".to_string(), true, 0, 2));
|
||||
let files: Vec<(String, bool, bool, Option<String>)> = sqlx::query_as(
|
||||
r#"
|
||||
SELECT status::text, completed_at IS NOT NULL, input_path IS NULL, error_message
|
||||
FROM task_files WHERE task_id = $1 ORDER BY original_name
|
||||
"#,
|
||||
)
|
||||
.bind(task_id)
|
||||
.fetch_all(&pool)
|
||||
.await
|
||||
.expect("query finalized task files");
|
||||
assert_eq!(files.len(), 2);
|
||||
assert!(files.iter().all(|row| {
|
||||
row.0 == "failed" && row.1 && row.2 && row.3.as_deref() == Some("队列提交失败")
|
||||
}));
|
||||
assert!(
|
||||
!tokio::fs::try_exists(&input_dir)
|
||||
.await
|
||||
.expect("check input directory"),
|
||||
"task input directory or partial upload survived cleanup"
|
||||
);
|
||||
|
||||
let before = task.clone();
|
||||
let outcome = crate::worker::process_task(&state, task_id, Uuid::new_v4())
|
||||
.await
|
||||
.expect("reprocess terminal task");
|
||||
assert_eq!(outcome, TaskProcessOutcome::Done);
|
||||
let after: (String, bool, i32, i32) = sqlx::query_as(
|
||||
r#"
|
||||
SELECT status::text, completed_at IS NOT NULL, completed_files, failed_files
|
||||
FROM tasks WHERE id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(task_id)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("query task after reprocess");
|
||||
assert_eq!(after, before, "reprocessing changed a terminal task");
|
||||
|
||||
let _: i64 = redis::cmd("DEL")
|
||||
.arg("stream:compress_jobs")
|
||||
.query_async(&mut redis)
|
||||
.await
|
||||
.expect("remove WRONGTYPE fixture");
|
||||
sqlx::query("DELETE FROM tasks WHERE id = $1")
|
||||
.bind(task_id)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("delete batch test task");
|
||||
let _ = tokio::fs::remove_dir_all(&storage_root).await;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -703,12 +703,12 @@ struct TaskContext {
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum TaskProcessOutcome {
|
||||
pub(crate) enum TaskProcessOutcome {
|
||||
Done,
|
||||
LeaseBusy,
|
||||
}
|
||||
|
||||
async fn process_task(
|
||||
pub(crate) async fn process_task(
|
||||
state: &AppState,
|
||||
task_id: Uuid,
|
||||
worker_id: Uuid,
|
||||
|
||||
Reference in New Issue
Block a user