fix(storage): make queue and object ownership durable

This commit is contained in:
237899745
2026-07-26 10:54:13 +08:00
parent f2d490edce
commit 408e09cda8
14 changed files with 2855 additions and 608 deletions

View File

@@ -8,6 +8,7 @@ use crate::services::compress;
use crate::services::compress::{CompressionLevel, ImageFmt};
use crate::services::filename;
use crate::services::idempotency;
use crate::services::object_lifecycle;
use crate::services::quota;
use crate::services::storage;
use crate::state::AppState;
@@ -21,6 +22,7 @@ use chrono::{DateTime, Duration, Utc};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use sqlx::FromRow;
use std::future::Future;
use std::net::{IpAddr, SocketAddr};
use uuid::Uuid;
@@ -30,6 +32,14 @@ pub fn router() -> Router<AppState> {
.route("/compress/direct", post(compress_direct))
}
fn spawn_detached_operation<F, T>(future: F) -> tokio::task::JoinHandle<T>
where
F: Future<Output = T> + Send + 'static,
T: Send + 'static,
{
tokio::spawn(future)
}
#[derive(Debug, Serialize, Deserialize)]
struct BillingView {
units_charged: i32,
@@ -214,6 +224,7 @@ async fn compress_json(
let quota_ctx = admission.quota_ctx;
let mut idem_acquired = false;
let mut idem_owner = None;
if let (Some(scope), Some(idem_key), Some(request_hash)) = (
idempotency_scope,
idempotency_key.as_deref(),
@@ -262,15 +273,35 @@ async fn compress_json(
"请求正在处理中,请稍后重试",
));
}
idempotency::BeginResult::Acquired => {
idempotency::BeginResult::Acquired { owner } => {
idem_acquired = true;
idem_owner = Some(owner);
}
}
}
let task_id = Uuid::new_v4();
let mut anonymous_reservation_date = None;
let op: Result<CompressResponse, AppError> = (async {
let operation_state = state.clone();
let operation_principal = principal.clone();
let operation_quota_ctx = quota_ctx.clone();
let operation_idempotency_key = idempotency_key.clone();
let operation_request_hash = request_hash.clone();
let operation_idem_owner = idem_owner;
let operation = spawn_detached_operation(async move {
let state = operation_state;
let principal = operation_principal;
let quota_ctx = operation_quota_ctx;
let idempotency_key = operation_idempotency_key;
let request_hash = operation_request_hash;
let _idempotency_heartbeat = start_idempotency_heartbeat(
&state,
idempotency_scope,
idempotency_key.as_deref(),
request_hash.as_deref(),
operation_idem_owner,
);
let mut anonymous_reservation_date = None;
let op: Result<CompressResponse, AppError> = (async {
match &quota_ctx {
QuotaContext::User(billing) => ensure_quota_available(&state, billing, 1).await?,
QuotaContext::ApiKey(billing, _) => ensure_quota_available(&state, billing, 1).await?,
@@ -318,11 +349,41 @@ async fn compress_json(
let retention_hours = retention.num_hours();
let object_key =
storage::result_key(retention_hours, task_id, file_id, format_out.extension());
let stored =
storage::store_bytes(&state, &object_key, compressed, format_out.content_type())
.await?;
let tracked = object_lifecycle::store_tracked_bytes(
&state,
task_id,
Some(file_id),
"result",
&object_key,
compressed,
format_out.content_type(),
)
.await?;
let expires_at = Utc::now() + retention;
let response = CompressResponse {
task_id,
file_id,
format_in: format_in.as_str().to_string(),
format_out: format_out.as_str().to_string(),
original_size,
compressed_size,
saved_bytes,
saved_percent,
download_url: format!("/downloads/{file_id}"),
expires_at,
billing: BillingView {
units_charged: metered_units(charge_units),
},
};
let idem_completion = build_idempotency_completion(
idem_acquired,
operation_idem_owner,
idempotency_scope,
idempotency_key.as_deref(),
request_hash.as_deref(),
&response,
)?;
if let Err(err) = record_task_and_metering(
&state,
@@ -330,7 +391,7 @@ async fn compress_json(
ip,
task_id,
file_id,
&stored,
&tracked,
&req.file_name,
req.max_width,
req.max_height,
@@ -345,19 +406,31 @@ async fn compress_json(
retention_hours,
&quota_ctx,
charge_units,
idem_completion.as_ref(),
)
.await
{
let _ = storage::delete_object(
&state,
&storage::ObjectLocator {
backend: stored.backend.clone(),
endpoint_id: stored.endpoint_id,
key: stored.key.clone(),
},
)
.await;
return Err(err);
match sync_result_was_committed(&state, task_id, file_id, &tracked).await {
Ok(true) => {
tracing::warn!(task_id = %task_id, file_id = %file_id, error = %err, "sync result commit response was lost; recovered committed publication");
}
Ok(false) => {
if let Err(cleanup_err) = object_lifecycle::schedule_tracked_delete(
&state,
&tracked,
Some(&err),
)
.await
{
tracing::error!(task_id = %task_id, storage_object_id = %tracked.lifecycle_id, error = %cleanup_err, "failed to persist rejected sync result cleanup");
}
return Err(err);
}
Err(probe_err) => {
tracing::error!(task_id = %task_id, storage_object_id = %tracked.lifecycle_id, error = %probe_err, original_error = %err, "sync result commit state is unknown; staging lease will reconcile object");
return Err(err);
}
}
}
if anonymous_reservation_date.is_some() {
@@ -371,51 +444,23 @@ async fn compress_json(
}
}
Ok(CompressResponse {
task_id,
file_id,
format_in: format_in.as_str().to_string(),
format_out: format_out.as_str().to_string(),
original_size,
compressed_size,
saved_bytes,
saved_percent,
download_url: format!("/downloads/{file_id}"),
expires_at,
billing: BillingView {
units_charged: metered_units(charge_units),
},
Ok(response)
})
})
.await;
.await;
(op, anonymous_reservation_date)
});
let (op, anonymous_reservation_date) = operation.await.map_err(|err| {
AppError::new(ErrorCode::Internal, "同步压缩后台任务异常退出").with_source(err)
})?;
match op {
Ok(resp) => {
if let (Some(scope), Some(idem_key), Some(request_hash)) = (
idempotency_scope,
idempotency_key.as_deref(),
request_hash.as_deref(),
) {
if idem_acquired {
let _ = idempotency::complete(
&state,
scope,
idem_key,
request_hash,
200,
serde_json::to_value(&resp).unwrap_or(serde_json::Value::Null),
)
.await;
}
}
Ok((
jar,
Json(Envelope {
success: true,
data: resp,
}),
))
}
Ok(resp) => Ok((
jar,
Json(Envelope {
success: true,
data: resp,
}),
)),
Err(err) => {
if anonymous_reservation_date.is_some() {
if let Err(refund_err) =
@@ -430,7 +475,10 @@ async fn compress_json(
request_hash.as_deref(),
) {
if idem_acquired {
let _ = idempotency::abort(&state, scope, idem_key, request_hash).await;
if let Some(owner) = idem_owner {
let _ =
idempotency::abort(&state, scope, idem_key, request_hash, owner).await;
}
}
}
Err(err)
@@ -525,6 +573,7 @@ async fn compress_direct(
let quota_ctx = admission.quota_ctx;
let mut idem_acquired = false;
let mut idem_owner = None;
if let (Some(scope), Some(idem_key), Some(request_hash)) = (
idempotency_scope,
idempotency_key.as_deref(),
@@ -566,13 +615,33 @@ async fn compress_direct(
"请求正在处理中,请稍后重试",
));
}
idempotency::BeginResult::Acquired => {
idempotency::BeginResult::Acquired { owner } => {
idem_acquired = true;
idem_owner = Some(owner);
}
}
}
let op: Result<(axum::response::Response, DirectIdempotencyData), AppError> = (async {
let operation_state = state.clone();
let operation_principal = principal.clone();
let operation_quota_ctx = quota_ctx.clone();
let operation_idempotency_key = idempotency_key.clone();
let operation_request_hash = request_hash.clone();
let operation_idem_owner = idem_owner;
let operation = spawn_detached_operation(async move {
let state = operation_state;
let principal = operation_principal;
let quota_ctx = operation_quota_ctx;
let idempotency_key = operation_idempotency_key;
let request_hash = operation_request_hash;
let _idempotency_heartbeat = start_idempotency_heartbeat(
&state,
idempotency_scope,
idempotency_key.as_deref(),
request_hash.as_deref(),
operation_idem_owner,
);
let op: Result<(axum::response::Response, DirectIdempotencyData), AppError> = (async {
match &quota_ctx {
QuotaContext::User(billing) => ensure_quota_available(&state, billing, 1).await?,
QuotaContext::ApiKey(billing, _) => ensure_quota_available(&state, billing, 1).await?,
@@ -617,8 +686,11 @@ async fn compress_direct(
let retention_hours = retention.num_hours();
let object_key =
storage::result_key(retention_hours, task_id, file_id, format_out.extension());
let stored = storage::store_bytes(
let tracked = object_lifecycle::store_tracked_bytes(
&state,
task_id,
Some(file_id),
"result",
&object_key,
compressed.clone(),
format_out.content_type(),
@@ -626,6 +698,23 @@ async fn compress_direct(
.await?;
let expires_at = Utc::now() + retention;
let idem_data = DirectIdempotencyData {
file_id,
format_out: format_out.as_str().to_string(),
original_size,
compressed_size,
saved_bytes,
saved_percent,
units_charged: metered_units(charge_units),
};
let idem_completion = build_idempotency_completion(
idem_acquired,
operation_idem_owner,
idempotency_scope,
idempotency_key.as_deref(),
request_hash.as_deref(),
&idem_data,
)?;
if let Err(err) = record_task_and_metering(
&state,
@@ -633,7 +722,7 @@ async fn compress_direct(
ip,
task_id,
file_id,
&stored,
&tracked,
&req.file_name,
req.max_width,
req.max_height,
@@ -648,56 +737,45 @@ async fn compress_direct(
retention_hours,
&quota_ctx,
charge_units,
idem_completion.as_ref(),
)
.await
{
let _ = storage::delete_object(
&state,
&storage::ObjectLocator {
backend: stored.backend.clone(),
endpoint_id: stored.endpoint_id,
key: stored.key.clone(),
},
)
.await;
return Err(err);
}
let idem_data = DirectIdempotencyData {
file_id,
format_out: format_out.as_str().to_string(),
original_size,
compressed_size,
saved_bytes,
saved_percent,
units_charged: metered_units(charge_units),
};
let response = direct_response(compressed, format_out, &idem_data);
Ok((response, idem_data))
})
.await;
match op {
Ok((response, idem_data)) => {
if let (Some(scope), Some(idem_key), Some(request_hash)) = (
idempotency_scope,
idempotency_key.as_deref(),
request_hash.as_deref(),
) {
if idem_acquired {
let _ = idempotency::complete(
match sync_result_was_committed(&state, task_id, file_id, &tracked).await {
Ok(true) => {
tracing::warn!(task_id = %task_id, file_id = %file_id, error = %err, "direct result commit response was lost; recovered committed publication");
}
Ok(false) => {
if let Err(cleanup_err) = object_lifecycle::schedule_tracked_delete(
&state,
scope,
idem_key,
request_hash,
200,
serde_json::to_value(&idem_data).unwrap_or(serde_json::Value::Null),
&tracked,
Some(&err),
)
.await;
.await
{
tracing::error!(task_id = %task_id, storage_object_id = %tracked.lifecycle_id, error = %cleanup_err, "failed to persist rejected direct result cleanup");
}
return Err(err);
}
Err(probe_err) => {
tracing::error!(task_id = %task_id, storage_object_id = %tracked.lifecycle_id, error = %probe_err, original_error = %err, "direct result commit state is unknown; staging lease will reconcile object");
return Err(err);
}
}
Ok((jar, response))
}
let response = direct_response(compressed, format_out, &idem_data);
Ok((response, idem_data))
})
.await;
op
});
let op = operation.await.map_err(|err| {
AppError::new(ErrorCode::Internal, "直接压缩后台任务异常退出").with_source(err)
})?;
match op {
Ok((response, _idem_data)) => Ok((jar, response)),
Err(err) => {
if let (Some(scope), Some(idem_key), Some(request_hash)) = (
idempotency_scope,
@@ -705,7 +783,10 @@ async fn compress_direct(
request_hash.as_deref(),
) {
if idem_acquired {
let _ = idempotency::abort(&state, scope, idem_key, request_hash).await;
if let Some(owner) = idem_owner {
let _ =
idempotency::abort(&state, scope, idem_key, request_hash, owner).await;
}
}
}
Err(err)
@@ -713,6 +794,45 @@ async fn compress_direct(
}
}
async fn sync_result_was_committed(
state: &AppState,
task_id: Uuid,
file_id: Uuid,
tracked: &object_lifecycle::TrackedStoredObject,
) -> Result<bool, AppError> {
sqlx::query_scalar(
r#"
SELECT EXISTS(
SELECT 1
FROM tasks AS task
JOIN task_files AS file ON file.task_id = task.id
JOIN storage_objects AS object ON object.id = $3
WHERE task.id = $1
AND task.status = 'completed'
AND file.id = $2
AND file.status = 'completed'
AND file.storage_backend = $4
AND file.storage_endpoint_id IS NOT DISTINCT FROM $5
AND COALESCE(file.storage_key, file.storage_path) = $6
AND object.state = 'published'
AND object.task_id = task.id
AND object.task_file_id = file.id
)
"#,
)
.bind(task_id)
.bind(file_id)
.bind(tracked.lifecycle_id)
.bind(&tracked.stored.backend)
.bind(tracked.stored.endpoint_id)
.bind(&tracked.stored.key)
.fetch_one(&state.db)
.await
.map_err(|err| {
AppError::new(ErrorCode::StorageUnavailable, "核验同步压缩提交结果失败").with_source(err)
})
}
#[derive(Debug, FromRow)]
struct DirectReplayRow {
storage_backend: String,
@@ -742,6 +862,7 @@ async fn load_direct_replay_bytes(
FROM task_files f
JOIN tasks t ON t.id = f.task_id
WHERE f.id = $1 AND t.user_id = $2
AND t.deletion_started_at IS NULL
"#,
)
.bind(file_id)
@@ -762,6 +883,7 @@ async fn load_direct_replay_bytes(
FROM task_files f
JOIN tasks t ON t.id = f.task_id
WHERE f.id = $1 AND t.api_key_id = $2
AND t.deletion_started_at IS NULL
"#,
)
.bind(file_id)
@@ -946,6 +1068,63 @@ enum QuotaContext {
ApiKey(BillingContext, Uuid),
}
struct IdempotencyCompletion {
scope: idempotency::Scope,
owner: Uuid,
key: String,
request_hash: String,
response_body: serde_json::Value,
}
fn build_idempotency_completion<T: Serialize>(
acquired: bool,
owner: Option<Uuid>,
scope: Option<idempotency::Scope>,
key: Option<&str>,
request_hash: Option<&str>,
response: &T,
) -> Result<Option<IdempotencyCompletion>, AppError> {
if !acquired {
return Ok(None);
}
let (owner, scope, key, request_hash) = match (owner, scope, key, request_hash) {
(Some(owner), Some(scope), Some(key), Some(request_hash)) => {
(owner, scope, key, request_hash)
}
_ => return Err(AppError::new(ErrorCode::Internal, "幂等请求上下文不完整")),
};
Ok(Some(IdempotencyCompletion {
scope,
owner,
key: key.to_string(),
request_hash: request_hash.to_string(),
response_body: serde_json::to_value(response).map_err(|err| {
AppError::new(ErrorCode::Internal, "序列化幂等响应失败").with_source(err)
})?,
}))
}
fn start_idempotency_heartbeat(
state: &AppState,
scope: Option<idempotency::Scope>,
key: Option<&str>,
request_hash: Option<&str>,
owner: Option<Uuid>,
) -> Option<idempotency::LeaseHeartbeat> {
match (scope, key, request_hash, owner) {
(Some(scope), Some(key), Some(request_hash), Some(owner)) => {
Some(idempotency::start_lease_heartbeat(
state.clone(),
scope,
key.to_string(),
request_hash.to_string(),
owner,
))
}
_ => None,
}
}
struct SingleAdmission {
retention: Duration,
quota_ctx: QuotaContext,
@@ -1044,7 +1223,7 @@ async fn record_task_and_metering(
client_ip: IpAddr,
task_id: Uuid,
file_id: Uuid,
stored: &storage::StoredObject,
tracked: &object_lifecycle::TrackedStoredObject,
original_name: &str,
max_width: Option<u32>,
max_height: Option<u32>,
@@ -1059,7 +1238,9 @@ async fn record_task_and_metering(
retention_hours: i64,
quota_ctx: &QuotaContext,
charge_units: bool,
idempotency_completion: Option<&IdempotencyCompletion>,
) -> Result<(), AppError> {
let stored = &tracked.stored;
let (user_id, session_id, api_key_id, source) = match principal {
context::Principal::Anonymous { session_id } => {
(None, Some(session_id.clone()), None, "web")
@@ -1191,6 +1372,20 @@ async fn record_task_and_metering(
}
}
object_lifecycle::publish_in_tx(&mut tx, tracked).await?;
if let Some(idempotency_completion) = idempotency_completion {
idempotency::complete_in_tx(
&mut tx,
idempotency_completion.scope,
&idempotency_completion.key,
&idempotency_completion.request_hash,
idempotency_completion.owner,
200,
idempotency_completion.response_body.clone(),
)
.await?;
}
tx.commit()
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "提交事务失败").with_source(err))?;
@@ -1244,6 +1439,8 @@ async fn charge_one_unit(
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
#[test]
fn direct_response_has_consistent_compression_headers() {
@@ -1296,4 +1493,35 @@ mod tests {
assert_eq!(metered_units(charged), expected_units);
}
}
#[tokio::test]
async fn detached_sync_operation_survives_waiter_abort() {
let started = Arc::new(tokio::sync::Notify::new());
let release = Arc::new(tokio::sync::Notify::new());
let completed = Arc::new(AtomicBool::new(false));
let request_started = started.clone();
let request_release = release.clone();
let request_completed = completed.clone();
let request = tokio::spawn(async move {
spawn_detached_operation(async move {
request_started.notify_one();
request_release.notified().await;
request_completed.store(true, Ordering::SeqCst);
})
.await
.expect("detached operation panicked");
});
started.notified().await;
request.abort();
request.await.expect_err("request waiter was not aborted");
release.notify_one();
tokio::time::timeout(std::time::Duration::from_secs(1), async {
while !completed.load(Ordering::SeqCst) {
tokio::task::yield_now().await;
}
})
.await
.expect("detached operation was canceled with its request waiter");
}
}

View File

@@ -1,5 +1,6 @@
use crate::api::context;
use crate::error::{AppError, ErrorCode};
use crate::services::object_lifecycle;
use crate::services::storage;
use crate::state::AppState;
@@ -63,7 +64,7 @@ async fn download_file(
t.expires_at
FROM task_files f
JOIN tasks t ON t.id = f.task_id
WHERE f.id = $1
WHERE f.id = $1 AND t.deletion_started_at IS NULL
"#,
)
.bind(file_id)
@@ -289,7 +290,7 @@ async fn download_task_zip(
zip_storage_endpoint_id,
zip_storage_key
FROM tasks
WHERE id = $1
WHERE id = $1 AND deletion_started_at IS NULL
"#,
)
.bind(task_id)
@@ -451,6 +452,7 @@ async fn claim_zip_build(state: &AppState, task_id: Uuid) -> Result<ZipBuildClai
zip_build_attempt = zip_build_attempt + 1
WHERE id = $1
AND zip_storage_key IS NULL
AND deletion_started_at IS NULL
AND completed_at IS NOT NULL
AND expires_at > NOW()
AND (
@@ -474,7 +476,7 @@ async fn claim_zip_build(state: &AppState, task_id: Uuid) -> Result<ZipBuildClai
r#"
SELECT zip_storage_backend, zip_storage_endpoint_id, zip_storage_key, expires_at
FROM tasks
WHERE id = $1
WHERE id = $1 AND deletion_started_at IS NULL
"#,
)
.bind(task_id)
@@ -515,6 +517,7 @@ async fn renew_zip_build_for(
AND zip_build_token = $2
AND zip_storage_key IS NULL
AND expires_at > NOW()
AND deletion_started_at IS NULL
"#,
)
.bind(task_id)
@@ -711,7 +714,7 @@ async fn build_zip_attempt(
rows: &[TaskZipFileRow],
temp_dir: &std::path::Path,
zip_path: &std::path::Path,
) -> Result<storage::StoredObject, AppError> {
) -> Result<object_lifecycle::TrackedStoredObject, AppError> {
tokio::fs::create_dir_all(temp_dir).await.map_err(|err| {
AppError::new(ErrorCode::StorageUnavailable, "创建 ZIP 临时目录失败").with_source(err)
})?;
@@ -766,15 +769,27 @@ async fn build_zip_attempt(
renew_zip_build(state, task_id, token).await?;
let object_key = storage::archive_attempt_key(retention_hours, task_id, token);
storage::store_file(state, &object_key, zip_path, "application/zip").await
object_lifecycle::store_tracked_file(
state,
task_id,
"zip_attempt",
&object_key,
zip_path,
"application/zip",
)
.await
}
async fn publish_zip_attempt(
state: &AppState,
task_id: Uuid,
token: Uuid,
stored: storage::StoredObject,
tracked: object_lifecycle::TrackedStoredObject,
) -> Result<storage::ObjectLocator, AppError> {
let stored = &tracked.stored;
let mut tx = state.db.begin().await.map_err(|err| {
AppError::new(ErrorCode::Internal, "开启 ZIP 发布事务失败").with_source(err)
})?;
let published = sqlx::query(
r#"
UPDATE tasks
@@ -789,6 +804,7 @@ async fn publish_zip_attempt(
AND zip_build_token = $2
AND zip_storage_key IS NULL
AND expires_at > NOW()
AND deletion_started_at IS NULL
"#,
)
.bind(task_id)
@@ -798,24 +814,36 @@ async fn publish_zip_attempt(
.bind(&stored.key)
.bind(&stored.etag)
.bind(stored.size as i64)
.execute(&state.db)
.execute(&mut *tx)
.await;
match published {
Ok(result) if result.rows_affected() == 1 => Ok(storage::ObjectLocator {
backend: stored.backend,
endpoint_id: stored.endpoint_id,
key: stored.key,
}),
Ok(result) if result.rows_affected() == 1 => {
if let Err(err) = object_lifecycle::publish_in_tx(&mut tx, &tracked).await {
tx.rollback().await.ok();
delete_unpublished_zip(state, task_id, token, &tracked).await;
return Err(err);
}
tx.commit().await.map_err(|err| {
AppError::new(ErrorCode::Internal, "提交 ZIP 发布事务失败").with_source(err)
})?;
Ok(storage::ObjectLocator {
backend: tracked.stored.backend,
endpoint_id: tracked.stored.endpoint_id,
key: tracked.stored.key,
})
}
Ok(_) => {
delete_unpublished_zip(state, task_id, token, &stored).await;
tx.rollback().await.ok();
delete_unpublished_zip(state, task_id, token, &tracked).await;
let current = load_published_zip(state, task_id).await?;
current.ok_or_else(|| {
AppError::new(ErrorCode::StorageUnavailable, "ZIP 发布租约已失效,请重试")
})
}
Err(err) => {
delete_unpublished_zip(state, task_id, token, &stored).await;
tx.rollback().await.ok();
delete_unpublished_zip(state, task_id, token, &tracked).await;
Err(AppError::new(ErrorCode::Internal, "记录 ZIP 对象失败").with_source(err))
}
}
@@ -825,28 +853,13 @@ async fn delete_unpublished_zip(
state: &AppState,
task_id: Uuid,
token: Uuid,
stored: &storage::StoredObject,
tracked: &object_lifecycle::TrackedStoredObject,
) {
let object = storage::ObjectLocator {
backend: stored.backend.clone(),
endpoint_id: stored.endpoint_id,
key: stored.key.clone(),
};
let mut last_error = None;
for attempt in 1..=3_u64 {
match storage::delete_object(state, &object).await {
Ok(()) => {
last_error = None;
break;
}
Err(err) => {
last_error = Some(err);
tokio::time::sleep(std::time::Duration::from_millis(100 * attempt)).await;
}
}
}
if let Some(err) = last_error {
tracing::error!(task_id = %task_id, zip_build_token = %token, object_key = %stored.key, error = %err, "failed to delete unpublished ZIP attempt after retries");
if let Err(err) = object_lifecycle::schedule_tracked_delete(state, tracked, None).await {
tracing::error!(task_id = %task_id, zip_build_token = %token, object_key = %tracked.stored.key, error = %err, "failed to persist unpublished ZIP deletion");
} else if let Err(err) = object_lifecycle::cleanup_ready_objects(state, 1, Some(task_id)).await
{
tracing::warn!(task_id = %task_id, zip_build_token = %token, object_key = %tracked.stored.key, error = %err, "unpublished ZIP deletion deferred");
}
release_zip_build(state, task_id, token).await;
}
@@ -856,7 +869,7 @@ async fn load_published_zip(
task_id: Uuid,
) -> Result<Option<storage::ObjectLocator>, AppError> {
let row: Option<(Option<String>, Option<Uuid>, Option<String>)> = sqlx::query_as(
"SELECT zip_storage_backend, zip_storage_endpoint_id, zip_storage_key FROM tasks WHERE id = $1",
"SELECT zip_storage_backend, zip_storage_endpoint_id, zip_storage_key FROM tasks WHERE id = $1 AND deletion_started_at IS NULL",
)
.bind(task_id)
.fetch_optional(&state.db)
@@ -1247,8 +1260,10 @@ mod tests {
tokio::fs::write(&unpublished_path, b"unpublished-zip")
.await
.expect("write unpublished ZIP fixture");
let unpublished = storage::store_file(
let unpublished = object_lifecycle::store_tracked_file(
&state,
deleted_task,
"zip_attempt",
&storage::archive_attempt_key(24, deleted_task, deleted_token),
&unpublished_path,
"application/zip",
@@ -1268,13 +1283,20 @@ mod tests {
let orphan_read = storage::read_bytes(
&state,
&storage::ObjectLocator {
backend: unpublished.backend,
endpoint_id: unpublished.endpoint_id,
key: unpublished.key,
backend: unpublished.stored.backend,
endpoint_id: unpublished.stored.endpoint_id,
key: unpublished.stored.key,
},
)
.await;
assert!(orphan_read.is_err(), "unpublished ZIP object was orphaned");
let unpublished_state: String =
sqlx::query_scalar("SELECT state FROM storage_objects WHERE id = $1")
.bind(unpublished.lifecycle_id)
.fetch_one(&pool)
.await
.expect("query unpublished ZIP lifecycle state");
assert_eq!(unpublished_state, "deleted");
let cancelled_task = Uuid::new_v4();
insert_zip_task(
@@ -1415,6 +1437,20 @@ mod tests {
.execute(&pool)
.await
.expect("delete ZIP test tasks");
sqlx::query("DELETE FROM storage_objects WHERE task_id = ANY($1)")
.bind(
&[
task_id,
over_budget_task,
takeover_task,
deleted_task,
cancelled_task,
heartbeat_task,
][..],
)
.execute(&pool)
.await
.expect("delete ZIP lifecycle test rows");
if let Some(endpoint_id) = endpoint_id {
sqlx::query("DELETE FROM storage_endpoints WHERE id = $1")
.bind(endpoint_id)

View File

@@ -17,18 +17,33 @@ const SCRAPE_TIMEOUT: Duration = Duration::from_secs(2);
pub async fn metrics(State(state): State<AppState>) -> impl IntoResponse {
let database = tokio::time::timeout(
SCRAPE_TIMEOUT,
sqlx::query_scalar::<_, i64>(
"SELECT COUNT(*) FROM tasks WHERE status IN ('pending', 'processing')",
sqlx::query_as::<_, (i64, i64, i64, i64, i64)>(
r#"
SELECT
(SELECT COUNT(*) FROM tasks WHERE status IN ('pending', 'processing')),
(SELECT COUNT(*) FROM task_queue_outbox WHERE status IN ('pending', 'delivering')),
(SELECT COUNT(*) FROM task_queue_outbox WHERE status = 'dead'),
(SELECT COUNT(*) FROM storage_objects WHERE state = 'delete_pending'),
(SELECT COUNT(*) FROM storage_objects WHERE state = 'staging')
"#,
)
.fetch_one(&state.db),
);
let redis = tokio::time::timeout(SCRAPE_TIMEOUT, redis_queue_stats(state.redis.clone()));
let (database, redis) = tokio::join!(database, redis);
let (database_up, active_tasks) = match database {
Ok(Ok(value)) => (1, value),
_ => (0, 0),
};
let (database_up, active_tasks, outbox_pending, outbox_dead, delete_pending, staging) =
match database {
Ok(Ok((active, outbox_pending, outbox_dead, delete_pending, staging))) => (
1,
active,
outbox_pending,
outbox_dead,
delete_pending,
staging,
),
_ => (0, 0, 0, 0, 0, 0),
};
let (redis_up, queue_length, pending, dead_length, cluster) = match redis {
Ok(Ok((queue_length, pending, dead_length, cluster))) => {
(1, queue_length, pending, dead_length, cluster)
@@ -51,6 +66,28 @@ pub async fn metrics(State(state): State<AppState>) -> impl IntoResponse {
output.push_str("# HELP imageforge_active_tasks Current pending or processing tasks.\n");
output.push_str("# TYPE imageforge_active_tasks gauge\n");
let _ = writeln!(output, "imageforge_active_tasks {active_tasks}");
output.push_str("# HELP imageforge_task_outbox Current durable task delivery states.\n");
output.push_str("# TYPE imageforge_task_outbox gauge\n");
let _ = writeln!(
output,
"imageforge_task_outbox{{state=\"pending\"}} {outbox_pending}"
);
let _ = writeln!(
output,
"imageforge_task_outbox{{state=\"dead\"}} {outbox_dead}"
);
output.push_str(
"# HELP imageforge_storage_object_lifecycle Current durable object cleanup states.\n",
);
output.push_str("# TYPE imageforge_storage_object_lifecycle gauge\n");
let _ = writeln!(
output,
"imageforge_storage_object_lifecycle{{state=\"delete_pending\"}} {delete_pending}"
);
let _ = writeln!(
output,
"imageforge_storage_object_lifecycle{{state=\"staging\"}} {staging}"
);
output.push_str("# HELP imageforge_queue_messages Current Redis stream message counts.\n");
output.push_str("# TYPE imageforge_queue_messages gauge\n");
let _ = writeln!(

View File

@@ -53,12 +53,19 @@ pub async fn run(state: AppState) -> Result<(), AppError> {
tracing::info!(addr = %addr, "API server listening");
let reconciliation_task = tokio::spawn(webhooks::reconciliation_loop(state.clone()));
let queue_dispatch_task =
tokio::spawn(crate::services::task_queue::dispatch_loop(state.clone()));
let object_lifecycle_task = tokio::spawn(crate::services::object_lifecycle::maintenance_loop(
state.clone(),
));
let serve_result = axum::serve(
listener,
app.into_make_service_with_connect_info::<SocketAddr>(),
)
.await;
reconciliation_task.abort();
queue_dispatch_task.abort();
object_lifecycle_task.abort();
serve_result
.map_err(|err| AppError::new(ErrorCode::Internal, "HTTP 服务异常退出").with_source(err))
}

View File

@@ -8,8 +8,9 @@ use crate::services::compress;
use crate::services::compress::{CompressionLevel, ImageFmt};
use crate::services::filename;
use crate::services::idempotency;
use crate::services::object_lifecycle;
use crate::services::quota;
use crate::services::storage;
use crate::services::task_queue;
use crate::state::AppState;
use axum::extract::{ConnectInfo, Multipart, Path, State};
@@ -127,6 +128,7 @@ async fn create_batch_task(
}
let mut idem_acquired = false;
let mut idem_owner = None;
if let (Some(scope), Some(idem_key)) = (idempotency_scope, idempotency_key.as_deref()) {
let begin_result = idempotency::begin(
&state,
@@ -181,8 +183,9 @@ async fn create_batch_task(
"请求正在处理中,请稍后重试",
));
}
idempotency::BeginResult::Acquired => {
idempotency::BeginResult::Acquired { owner } => {
idem_acquired = true;
idem_owner = Some(owner);
}
}
}
@@ -190,8 +193,6 @@ 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 } => {
@@ -299,75 +300,67 @@ async fn create_batch_task(
})?;
}
sqlx::query("INSERT INTO task_queue_outbox (task_id) VALUES ($1)")
.bind(task_id)
.execute(&mut *tx)
.await
.map_err(|err| {
AppError::new(ErrorCode::Internal, "创建任务队列 outbox 失败").with_source(err)
})?;
let response = BatchCreateResponse {
task_id,
total_files: files.len() as i32,
status: "pending".to_string(),
status_url: format!("/api/v1/compress/tasks/{task_id}"),
};
if let (Some(scope), Some(idem_key)) = (idempotency_scope, idempotency_key.as_deref()) {
if idem_acquired {
idempotency::complete_in_tx(
&mut tx,
scope,
idem_key,
&request_hash,
idem_owner.ok_or_else(|| {
AppError::new(ErrorCode::Internal, "幂等操作租约缺失")
})?,
200,
serde_json::to_value(&response).map_err(|err| {
AppError::new(ErrorCode::Internal, "序列化幂等响应失败").with_source(err)
})?,
)
.await?;
}
}
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 {
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");
}
let dispatch_state = state.clone();
tokio::spawn(async move {
if let Err(err) = task_queue::dispatch_task(&dispatch_state, task_id).await {
tracing::warn!(task_id = %task_id, error = %err, "immediate task queue dispatch deferred to outbox loop");
}
return Err(err);
}
});
Ok(BatchCreateResponse {
task_id,
total_files: files.len() as i32,
status: "pending".to_string(),
status_url: format!("/api/v1/compress/tasks/{task_id}"),
})
Ok(response)
})
.await;
match create_result {
Ok(resp) => {
if let (Some(scope), Some(idem_key)) = (idempotency_scope, idempotency_key.as_deref()) {
if idem_acquired {
let _ = idempotency::complete(
&state,
scope,
idem_key,
&request_hash,
200,
serde_json::to_value(&resp).unwrap_or(serde_json::Value::Null),
)
.await;
}
}
Ok((
jar,
Json(Envelope {
success: true,
data: resp,
}),
))
}
Ok(resp) => Ok((
jar,
Json(Envelope {
success: true,
data: resp,
}),
)),
Err(err) => {
if anonymous_reserved_units > 0 {
if let context::Principal::Anonymous { session_id } = &principal {
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 !task_persisted {
if let Some(date) = anonymous_quota_date {
if let Err(refund_err) = quota::refund_anonymous_reservation_once(
&state,
@@ -387,10 +380,13 @@ async fn create_batch_task(
}
if let (Some(scope), Some(idem_key)) = (idempotency_scope, idempotency_key.as_deref()) {
if idem_acquired {
let _ = idempotency::abort(&state, scope, idem_key, &request_hash).await;
if let Some(owner) = idem_owner {
let _ =
idempotency::abort(&state, scope, idem_key, &request_hash, owner).await;
}
}
}
if cleanup_inputs_on_error {
if !task_persisted {
cleanup_task_input_dir(&state, task_id).await;
}
Err(err)
@@ -484,107 +480,6 @@ fn plan_upload_limits(plan: &Plan) -> Result<BatchUploadLimits, AppError> {
})
}
async fn enqueue_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
let mut conn = state.redis.clone();
let now = Utc::now().to_rfc3339();
redis::cmd("XADD")
.arg("stream:compress_jobs")
.arg("MAXLEN")
.arg("~")
.arg(100_000)
.arg("*")
.arg("task_id")
.arg(task_id.to_string())
.arg("created_at")
.arg(now)
.query_async::<_, redis::Value>(&mut conn)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "写入队列失败").with_source(err))?;
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,
@@ -948,7 +843,7 @@ async fn get_task(
user_id,
session_id
FROM tasks
WHERE id = $1
WHERE id = $1 AND deletion_started_at IS NULL
"#,
)
.bind(task_id)
@@ -1051,7 +946,7 @@ async fn cancel_task(
context::require_api_permission(&principal, &["compress", "batch_compress"])?;
let task = sqlx::query_as::<_, TaskRow>(
"SELECT status::text AS status, total_files, completed_files, failed_files, created_at, completed_at, expires_at, user_id, session_id FROM tasks WHERE id = $1",
"SELECT status::text AS status, total_files, completed_files, failed_files, created_at, completed_at, expires_at, user_id, session_id FROM tasks WHERE id = $1 AND deletion_started_at IS NULL",
)
.bind(task_id)
.fetch_optional(&state.db)
@@ -1077,7 +972,7 @@ async fn cancel_task(
}
let updated = sqlx::query(
"UPDATE tasks SET status = 'cancelled', completed_at = NOW() WHERE id = $1 AND status IN ('pending', 'processing')",
"UPDATE tasks SET status = 'cancelled', completed_at = NOW() WHERE id = $1 AND deletion_started_at IS NULL AND status IN ('pending', 'processing')",
)
.bind(task_id)
.execute(&state.db)
@@ -1116,13 +1011,16 @@ async fn delete_task(
let (jar, principal) = context::authenticate(&state, jar, &headers, ip).await?;
context::require_api_permission(&principal, &["compress", "batch_compress"])?;
let mut tx = state.db.begin().await.map_err(|err| {
AppError::new(ErrorCode::Internal, "开启任务删除事务失败").with_source(err)
})?;
let task = sqlx::query_as::<_, TaskRow>(
"SELECT status::text AS status, total_files, completed_files, failed_files, created_at, completed_at, expires_at, user_id, session_id FROM tasks WHERE id = $1",
"SELECT status::text AS status, total_files, completed_files, failed_files, created_at, completed_at, expires_at, user_id, session_id FROM tasks WHERE id = $1 FOR UPDATE",
)
.bind(task_id)
.fetch_optional(&state.db)
.fetch_optional(&mut *tx)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "查询任务失败").with_source(err))?
.map_err(|err| AppError::new(ErrorCode::Internal, "锁定待删除任务失败").with_source(err))?
.ok_or_else(|| AppError::new(ErrorCode::NotFound, "任务不存在"))?;
authorize_task(
@@ -1138,87 +1036,18 @@ async fn delete_task(
));
}
if task.status == "pending" {
let updated = sqlx::query(
"UPDATE tasks SET status = 'cancelled', completed_at = NOW() WHERE id = $1 AND status = 'pending'",
)
.bind(task_id)
.execute(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "锁定待删除任务失败").with_source(err))?;
if updated.rows_affected() == 0 {
return Err(AppError::new(
ErrorCode::InvalidRequest,
"任务状态已变化,请刷新后重试",
));
}
if !object_lifecycle::mark_task_deleting(&mut tx, task_id, "user", false).await? {
return Err(AppError::new(
ErrorCode::InvalidRequest,
"任务状态已变化,请刷新后重试",
));
}
tx.commit().await.map_err(|err| {
AppError::new(ErrorCode::Internal, "提交任务删除状态失败").with_source(err)
})?;
quota::settle_anonymous_task_reservation(&state, task_id).await?;
let files = sqlx::query_as::<_, TaskStorageRow>(
r#"
SELECT storage_backend, storage_endpoint_id,
COALESCE(storage_key, storage_path) AS storage_key,
input_path
FROM task_files
WHERE task_id = $1
"#,
)
.bind(task_id)
.fetch_all(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "查询文件失败").with_source(err))?;
for file in files {
if let Some(key) = file.storage_key {
storage::delete_object(
&state,
&storage::ObjectLocator {
backend: file.storage_backend,
endpoint_id: file.storage_endpoint_id,
key,
},
)
.await?;
}
if let Some(input_path) = file.input_path {
let _ = tokio::fs::remove_file(input_path).await;
}
}
let zip = sqlx::query_as::<_, TaskZipStorageRow>(
"SELECT zip_storage_backend, zip_storage_endpoint_id, zip_storage_key FROM tasks WHERE id = $1",
)
.bind(task_id)
.fetch_one(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "查询 ZIP 文件失败").with_source(err))?;
if let (Some(backend), Some(key)) = (zip.zip_storage_backend, zip.zip_storage_key) {
storage::delete_object(
&state,
&storage::ObjectLocator {
backend,
endpoint_id: zip.zip_storage_endpoint_id,
key,
},
)
.await?;
}
let legacy_zip_path = format!("{}/zips/{task_id}.zip", state.config.storage_path);
let _ = tokio::fs::remove_file(legacy_zip_path).await;
let orig_dir = format!("{}/orig/{task_id}", state.config.storage_path);
let _ = tokio::fs::remove_dir_all(orig_dir).await;
let deleted = sqlx::query("DELETE FROM tasks WHERE id = $1")
.bind(task_id)
.execute(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "删除任务失败").with_source(err))?;
if deleted.rows_affected() == 0 {
return Err(AppError::new(ErrorCode::NotFound, "任务不存在"));
if let Err(err) = object_lifecycle::finalize_task_deletion(&state, task_id).await {
tracing::warn!(task_id = %task_id, error = %err, "task deletion persisted and will be retried by lifecycle maintenance");
}
Ok((
@@ -1230,21 +1059,6 @@ async fn delete_task(
))
}
#[derive(Debug, FromRow)]
struct TaskStorageRow {
storage_backend: String,
storage_endpoint_id: Option<Uuid>,
storage_key: Option<String>,
input_path: Option<String>,
}
#[derive(Debug, FromRow)]
struct TaskZipStorageRow {
zip_storage_backend: Option<String>,
zip_storage_endpoint_id: Option<Uuid>,
zip_storage_key: Option<String>,
}
fn authorize_task(
principal: &context::Principal,
user_id: Option<Uuid>,
@@ -1334,7 +1148,7 @@ mod tests {
#[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() {
async fn outbox_retries_ambiguous_delivery_and_dead_letters_pending_task() {
let database_url = std::env::var("IMAGEFORGE_TEST_DATABASE_URL")
.expect("IMAGEFORGE_TEST_DATABASE_URL must be set");
assert!(
@@ -1410,6 +1224,11 @@ mod tests {
.await
.expect("insert pending batch file");
}
sqlx::query("INSERT INTO task_queue_outbox (task_id) VALUES ($1)")
.bind(task_id)
.execute(&pool)
.await
.expect("insert task outbox");
let mut redis = state.redis.clone();
let _: i64 = redis::cmd("DEL")
@@ -1424,17 +1243,48 @@ mod tests {
.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, "队列提交失败")
task_queue::dispatch_task(&state, task_id)
.await
.expect("finalize enqueue failure"),
"pending task was not finalized"
.expect("dispatch outbox through WRONGTYPE"),
"outbox was not claimed"
);
let after_first_failure: (String, String, i32) = sqlx::query_as(
r#"
SELECT task.status::text, outbox.status, outbox.attempts
FROM tasks AS task
JOIN task_queue_outbox AS outbox ON outbox.task_id = task.id
WHERE task.id = $1
"#,
)
.bind(task_id)
.fetch_one(&pool)
.await
.expect("query retryable outbox state");
assert_eq!(
after_first_failure,
("pending".to_string(), "pending".to_string(), 1)
);
assert!(
tokio::fs::try_exists(&input_dir)
.await
.expect("check retained input directory"),
"a retryable Redis error deleted task inputs"
);
sqlx::query(
"UPDATE task_queue_outbox SET attempts = 19, next_attempt_at = NOW() WHERE task_id = $1",
)
.bind(task_id)
.execute(&pool)
.await
.expect("advance outbox to final attempt");
assert!(
task_queue::dispatch_task(&state, task_id)
.await
.expect("dispatch final outbox attempt"),
"final outbox attempt was not claimed"
);
cleanup_task_input_dir(&state, task_id).await;
let task: (String, bool, i32, i32) = sqlx::query_as(
r#"
@@ -1459,8 +1309,22 @@ mod tests {
.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("队列提交失败")
row.0 == "failed"
&& row.1
&& row.2
&& row
.3
.as_deref()
.is_some_and(|message| message.starts_with("队列持续不可用:"))
}));
let outbox_status: (String, i32, bool) = sqlx::query_as(
"SELECT status, attempts, last_error IS NOT NULL FROM task_queue_outbox WHERE task_id = $1",
)
.bind(task_id)
.fetch_one(&pool)
.await
.expect("query dead outbox");
assert_eq!(outbox_status, ("dead".to_string(), 20, true));
assert!(
!tokio::fs::try_exists(&input_dir)
.await
@@ -1490,11 +1354,225 @@ mod tests {
.query_async(&mut redis)
.await
.expect("remove WRONGTYPE fixture");
sqlx::query("DELETE FROM tasks WHERE id = $1")
.bind(task_id)
let ambiguous_task = Uuid::new_v4();
sqlx::query(
r#"
INSERT INTO tasks (
id, session_id, status, total_files, expires_at, retention_hours
) VALUES ($1, $2, 'pending', 0, NOW() + INTERVAL '1 day', 24)
"#,
)
.bind(ambiguous_task)
.bind(format!("ambiguous-{marker}"))
.execute(&pool)
.await
.expect("insert ambiguous delivery task");
sqlx::query("INSERT INTO task_queue_outbox (task_id) VALUES ($1)")
.bind(ambiguous_task)
.execute(&pool)
.await
.expect("delete batch test task");
.expect("insert ambiguous delivery outbox");
task_queue::enqueue_task(&state, ambiguous_task)
.await
.expect("simulate XADD success with lost reply");
assert!(task_queue::dispatch_task(&state, ambiguous_task)
.await
.expect("retry ambiguous delivery"));
let ambiguous_stream_len: i64 = redis::cmd("XLEN")
.arg("stream:compress_jobs")
.query_async(&mut redis)
.await
.expect("count duplicate ambiguous messages");
assert_eq!(ambiguous_stream_len, 2);
sqlx::query("UPDATE tasks SET status = 'completed', completed_at = NOW() WHERE id = $1")
.bind(ambiguous_task)
.execute(&pool)
.await
.expect("simulate first duplicate message completing the task");
assert_eq!(
crate::worker::process_task(&state, ambiguous_task, Uuid::new_v4())
.await
.expect("process ambiguous delivery task"),
TaskProcessOutcome::Done
);
assert_eq!(
crate::worker::process_task(&state, ambiguous_task, Uuid::new_v4())
.await
.expect("reprocess duplicate ambiguous delivery task"),
TaskProcessOutcome::Done
);
let ambiguous_status: String =
sqlx::query_scalar("SELECT status::text FROM tasks WHERE id = $1")
.bind(ambiguous_task)
.fetch_one(&pool)
.await
.expect("query ambiguous task status");
assert_eq!(ambiguous_status, "completed");
let concurrent_task = Uuid::new_v4();
sqlx::query(
"INSERT INTO tasks (id, session_id, status, total_files, expires_at, retention_hours) VALUES ($1, $2, 'pending', 0, NOW() + INTERVAL '1 day', 24)",
)
.bind(concurrent_task)
.bind(format!("concurrent-{marker}"))
.execute(&pool)
.await
.expect("insert concurrent dispatcher task");
sqlx::query("INSERT INTO task_queue_outbox (task_id) VALUES ($1)")
.bind(concurrent_task)
.execute(&pool)
.await
.expect("insert concurrent dispatcher outbox");
let barrier = Arc::new(tokio::sync::Barrier::new(20));
let mut dispatchers = Vec::new();
for _ in 0..20 {
let state = state.clone();
let barrier = barrier.clone();
dispatchers.push(tokio::spawn(async move {
barrier.wait().await;
task_queue::dispatch_task(&state, concurrent_task).await
}));
}
let mut claimed = 0;
for dispatcher in dispatchers {
if dispatcher
.await
.expect("dispatcher task panicked")
.expect("concurrent dispatcher failed")
{
claimed += 1;
}
}
assert_eq!(claimed, 1, "multiple dispatchers owned one outbox row");
let takeover_task = Uuid::new_v4();
sqlx::query(
"INSERT INTO tasks (id, session_id, status, total_files, expires_at, retention_hours) VALUES ($1, $2, 'pending', 0, NOW() + INTERVAL '1 day', 24)",
)
.bind(takeover_task)
.bind(format!("takeover-{marker}"))
.execute(&pool)
.await
.expect("insert dispatcher takeover task");
sqlx::query(
r#"
INSERT INTO task_queue_outbox (
task_id, status, attempts, lease_owner, lease_until
) VALUES ($1, 'delivering', 1, $2, NOW() - INTERVAL '1 second')
"#,
)
.bind(takeover_task)
.bind(Uuid::new_v4())
.execute(&pool)
.await
.expect("insert expired dispatcher lease");
assert!(task_queue::dispatch_task(&state, takeover_task)
.await
.expect("take over expired dispatcher lease"));
let takeover_state: (String, i32) =
sqlx::query_as("SELECT status, attempts FROM task_queue_outbox WHERE task_id = $1")
.bind(takeover_task)
.fetch_one(&pool)
.await
.expect("query dispatcher takeover state");
assert_eq!(takeover_state, ("delivered".to_string(), 2));
let idem_user: Uuid = sqlx::query_scalar(
r#"
INSERT INTO users (email, username, password_hash, email_verified_at)
VALUES ($1, $2, 'test', NOW())
RETURNING id
"#,
)
.bind(format!("outbox-idem-{marker}@example.test"))
.bind(format!("outbox-idem-{marker}"))
.fetch_one(&pool)
.await
.expect("insert outbox idempotency user");
let idem_key = format!("outbox-{marker}");
let idem_hash = "b".repeat(64);
let idem_owner = match idempotency::begin(
&state,
idempotency::Scope::User(idem_user),
&idem_key,
&idem_hash,
24,
)
.await
.expect("acquire outbox idempotency key")
{
idempotency::BeginResult::Acquired { owner } => owner,
other => panic!("unexpected outbox idempotency begin result: {other:?}"),
};
let idem_task = Uuid::new_v4();
let idem_response = serde_json::json!({"task_id": idem_task});
let mut idem_tx = pool
.begin()
.await
.expect("begin atomic task/outbox response tx");
sqlx::query(
"INSERT INTO tasks (id, user_id, status, total_files, expires_at, retention_hours) VALUES ($1, $2, 'pending', 0, NOW() + INTERVAL '1 day', 24)",
)
.bind(idem_task)
.bind(idem_user)
.execute(&mut *idem_tx)
.await
.expect("insert idempotent task");
sqlx::query("INSERT INTO task_queue_outbox (task_id) VALUES ($1)")
.bind(idem_task)
.execute(&mut *idem_tx)
.await
.expect("insert idempotent task outbox");
idempotency::complete_in_tx(
&mut idem_tx,
idempotency::Scope::User(idem_user),
&idem_key,
&idem_hash,
idem_owner,
200,
idem_response.clone(),
)
.await
.expect("persist atomic idempotent task response");
idem_tx.commit().await.expect("commit idempotent task");
match idempotency::begin(
&state,
idempotency::Scope::User(idem_user),
&idem_key,
&idem_hash,
24,
)
.await
.expect("replay idempotent task")
{
idempotency::BeginResult::Replay { response_body } => {
assert_eq!(response_body, idem_response)
}
other => panic!("same idempotency key did not replay task: {other:?}"),
}
sqlx::query("DELETE FROM tasks WHERE id = ANY($1)")
.bind(vec![
task_id,
ambiguous_task,
concurrent_task,
takeover_task,
idem_task,
])
.execute(&pool)
.await
.expect("delete batch outbox test tasks");
sqlx::query("DELETE FROM users WHERE id = $1")
.bind(idem_user)
.execute(&pool)
.await
.expect("delete outbox idempotency user");
let _: i64 = redis::cmd("DEL")
.arg("stream:compress_jobs")
.query_async(&mut redis)
.await
.expect("clean compression stream");
let _ = tokio::fs::remove_dir_all(&storage_root).await;
}
}

View File

@@ -6,6 +6,8 @@ use serde_json::Value as JsonValue;
use sqlx::FromRow;
use uuid::Uuid;
const OPERATION_LEASE_MINUTES: i64 = 30;
#[derive(Debug, Clone, Copy)]
pub enum Scope {
User(Uuid),
@@ -14,7 +16,7 @@ pub enum Scope {
#[derive(Debug)]
pub enum BeginResult {
Acquired,
Acquired { owner: Uuid },
Replay { response_body: JsonValue },
InProgress,
}
@@ -54,6 +56,8 @@ pub async fn begin(
let now = Utc::now();
let expires_at = now + Duration::hours(ttl_hours.max(1));
let owner = Uuid::new_v4();
let lease_until = now + Duration::minutes(OPERATION_LEASE_MINUTES);
cleanup_expired_for_key(state, scope, idempotency_key, now).await?;
@@ -64,11 +68,11 @@ pub async fn begin(
INSERT INTO idempotency_keys (
user_id, idempotency_key, request_hash,
response_status, response_body,
expires_at
expires_at, lease_owner, lease_until
) VALUES (
$1, $2, $3,
0, NULL,
$4
$4, $5, $6
)
ON CONFLICT DO NOTHING
"#,
@@ -77,6 +81,8 @@ pub async fn begin(
.bind(idempotency_key)
.bind(request_hash)
.bind(expires_at)
.bind(owner)
.bind(lease_until)
.execute(&state.db)
.await
}
@@ -86,11 +92,11 @@ pub async fn begin(
INSERT INTO idempotency_keys (
api_key_id, idempotency_key, request_hash,
response_status, response_body,
expires_at
expires_at, lease_owner, lease_until
) VALUES (
$1, $2, $3,
0, NULL,
$4
$4, $5, $6
)
ON CONFLICT DO NOTHING
"#,
@@ -99,6 +105,8 @@ pub async fn begin(
.bind(idempotency_key)
.bind(request_hash)
.bind(expires_at)
.bind(owner)
.bind(lease_until)
.execute(&state.db)
.await
}
@@ -106,12 +114,15 @@ pub async fn begin(
.map_err(|err| AppError::new(ErrorCode::Internal, "写入幂等记录失败").with_source(err))?;
if inserted.rows_affected() > 0 {
return Ok(BeginResult::Acquired);
return Ok(BeginResult::Acquired { owner });
}
let row = get_row(state, scope, idempotency_key, now).await?;
let Some(row) = row else {
return Ok(BeginResult::Acquired);
return Err(AppError::new(
ErrorCode::StorageUnavailable,
"幂等记录状态已变化,请重试",
));
};
if row.request_hash != request_hash {
@@ -122,6 +133,19 @@ pub async fn begin(
}
if row.response_status == 0 || row.response_body.is_none() {
if take_over_stale_operation(
state,
scope,
idempotency_key,
request_hash,
owner,
lease_until,
now,
)
.await?
{
return Ok(BeginResult::Acquired { owner });
}
return Ok(BeginResult::InProgress);
}
@@ -130,6 +154,152 @@ pub async fn begin(
})
}
async fn take_over_stale_operation(
state: &AppState,
scope: Scope,
idempotency_key: &str,
request_hash: &str,
owner: Uuid,
lease_until: DateTime<Utc>,
now: DateTime<Utc>,
) -> Result<bool, AppError> {
let updated = match scope {
Scope::User(user_id) => {
sqlx::query(
r#"
UPDATE idempotency_keys
SET lease_owner = $4, lease_until = $5
WHERE user_id = $1 AND idempotency_key = $2
AND request_hash = $3 AND response_status = 0
AND (lease_until IS NULL OR lease_until <= $6)
"#,
)
.bind(user_id)
.bind(idempotency_key)
.bind(request_hash)
.bind(owner)
.bind(lease_until)
.bind(now)
.execute(&state.db)
.await
}
Scope::ApiKey(api_key_id) => {
sqlx::query(
r#"
UPDATE idempotency_keys
SET lease_owner = $4, lease_until = $5
WHERE api_key_id = $1 AND idempotency_key = $2
AND request_hash = $3 AND response_status = 0
AND (lease_until IS NULL OR lease_until <= $6)
"#,
)
.bind(api_key_id)
.bind(idempotency_key)
.bind(request_hash)
.bind(owner)
.bind(lease_until)
.bind(now)
.execute(&state.db)
.await
}
}
.map_err(|err| AppError::new(ErrorCode::Internal, "接管过期幂等操作失败").with_source(err))?;
Ok(updated.rows_affected() == 1)
}
pub struct LeaseHeartbeat(Option<tokio::sync::oneshot::Sender<()>>);
impl Drop for LeaseHeartbeat {
fn drop(&mut self) {
if let Some(stop) = self.0.take() {
let _ = stop.send(());
}
}
}
pub fn start_lease_heartbeat(
state: AppState,
scope: Scope,
idempotency_key: String,
request_hash: String,
owner: Uuid,
) -> LeaseHeartbeat {
let (stop_tx, mut stop_rx) = tokio::sync::oneshot::channel();
tokio::spawn(async move {
let mut interval = tokio::time::interval(std::time::Duration::from_secs(60));
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
interval.tick().await;
loop {
tokio::select! {
_ = &mut stop_rx => break,
_ = interval.tick() => {
match renew_lease(
&state,
scope,
&idempotency_key,
&request_hash,
owner,
).await {
Ok(true) => {}
Ok(false) => break,
Err(err) => tracing::warn!(error = %err, "failed to renew idempotency operation lease"),
}
}
}
}
});
LeaseHeartbeat(Some(stop_tx))
}
async fn renew_lease(
state: &AppState,
scope: Scope,
idempotency_key: &str,
request_hash: &str,
owner: Uuid,
) -> Result<bool, AppError> {
let updated = match scope {
Scope::User(user_id) => {
sqlx::query(
r#"
UPDATE idempotency_keys
SET lease_until = NOW() + ($5 * INTERVAL '1 minute')
WHERE user_id = $1 AND idempotency_key = $2
AND request_hash = $3 AND lease_owner = $4
AND response_status = 0
"#,
)
.bind(user_id)
.bind(idempotency_key)
.bind(request_hash)
.bind(owner)
.bind(OPERATION_LEASE_MINUTES)
.execute(&state.db)
.await
}
Scope::ApiKey(api_key_id) => {
sqlx::query(
r#"
UPDATE idempotency_keys
SET lease_until = NOW() + ($5 * INTERVAL '1 minute')
WHERE api_key_id = $1 AND idempotency_key = $2
AND request_hash = $3 AND lease_owner = $4
AND response_status = 0
"#,
)
.bind(api_key_id)
.bind(idempotency_key)
.bind(request_hash)
.bind(owner)
.bind(OPERATION_LEASE_MINUTES)
.execute(&state.db)
.await
}
}
.map_err(|err| AppError::new(ErrorCode::Internal, "续租幂等操作失败").with_source(err))?;
Ok(updated.rows_affected() == 1)
}
pub async fn wait_for_replay(
state: &AppState,
scope: Scope,
@@ -165,11 +335,12 @@ pub async fn wait_for_replay(
}
}
pub async fn complete(
state: &AppState,
pub async fn complete_in_tx(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
scope: Scope,
idempotency_key: &str,
request_hash: &str,
owner: Uuid,
response_status: i32,
response_body: JsonValue,
) -> Result<(), AppError> {
@@ -179,10 +350,13 @@ pub async fn complete(
r#"
UPDATE idempotency_keys
SET response_status = $4,
response_body = $5
response_body = $5,
lease_owner = NULL,
lease_until = NULL
WHERE user_id = $1
AND idempotency_key = $2
AND request_hash = $3
AND lease_owner = $6
AND response_status = 0
"#,
)
@@ -191,7 +365,8 @@ pub async fn complete(
.bind(request_hash)
.bind(response_status)
.bind(response_body)
.execute(&state.db)
.bind(owner)
.execute(&mut **tx)
.await
}
Scope::ApiKey(api_key_id) => {
@@ -199,10 +374,13 @@ pub async fn complete(
r#"
UPDATE idempotency_keys
SET response_status = $4,
response_body = $5
response_body = $5,
lease_owner = NULL,
lease_until = NULL
WHERE api_key_id = $1
AND idempotency_key = $2
AND request_hash = $3
AND lease_owner = $6
AND response_status = 0
"#,
)
@@ -211,16 +389,19 @@ pub async fn complete(
.bind(request_hash)
.bind(response_status)
.bind(response_body)
.execute(&state.db)
.bind(owner)
.execute(&mut **tx)
.await
}
}
.map_err(|err| AppError::new(ErrorCode::Internal, "写入幂等结果失败").with_source(err))?;
.map_err(|err| AppError::new(ErrorCode::Internal, "事务内写入幂等结果失败").with_source(err))?;
if updated.rows_affected() == 0 {
tracing::warn!("idempotency record not updated (maybe already completed?)");
if updated.rows_affected() != 1 {
return Err(AppError::new(
ErrorCode::IdempotencyConflict,
"幂等请求所有权已变化,请重试",
));
}
Ok(())
}
@@ -229,25 +410,28 @@ pub async fn abort(
scope: Scope,
idempotency_key: &str,
request_hash: &str,
owner: Uuid,
) -> Result<(), AppError> {
match scope {
Scope::User(user_id) => {
let _ = sqlx::query(
"DELETE FROM idempotency_keys WHERE user_id = $1 AND idempotency_key = $2 AND request_hash = $3 AND response_status = 0",
"DELETE FROM idempotency_keys WHERE user_id = $1 AND idempotency_key = $2 AND request_hash = $3 AND lease_owner = $4 AND response_status = 0",
)
.bind(user_id)
.bind(idempotency_key)
.bind(request_hash)
.bind(owner)
.execute(&state.db)
.await;
}
Scope::ApiKey(api_key_id) => {
let _ = sqlx::query(
"DELETE FROM idempotency_keys WHERE api_key_id = $1 AND idempotency_key = $2 AND request_hash = $3 AND response_status = 0",
"DELETE FROM idempotency_keys WHERE api_key_id = $1 AND idempotency_key = $2 AND request_hash = $3 AND lease_owner = $4 AND response_status = 0",
)
.bind(api_key_id)
.bind(idempotency_key)
.bind(request_hash)
.bind(owner)
.execute(&state.db)
.await;
}
@@ -334,3 +518,149 @@ async fn get_row(
Ok(row)
}
#[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 sqlx::postgres::PgPoolOptions;
use std::sync::Arc;
use tokio::sync::Semaphore;
async fn test_state(database_url: String, redis_url: String) -> AppState {
let mut config = Config::from_env().expect("load idempotency test config");
config.database_url = database_url.clone();
config.redis_url = redis_url;
config.mail_enabled = false;
config.mail_log_links_when_disabled = false;
let db = PgPoolOptions::new()
.max_connections(8)
.connect(&database_url)
.await
.expect("connect idempotency test database");
sqlx::migrate!().run(&db).await.expect("run migrations");
let redis = redis::Client::open(config.redis_url.clone())
.expect("create idempotency test Redis client")
.get_connection_manager()
.await
.expect("connect idempotency test Redis");
AppState {
mailer: Arc::new(Mailer::new(&config).expect("create disabled test mailer")),
image_processing_semaphore: Arc::new(Semaphore::new(1)),
zip_build_semaphore: Arc::new(Semaphore::new(1)),
runtime_policy_cache: RuntimePolicyCache::new(),
storage_cache: StorageCache::new(),
config,
db,
redis,
}
}
#[tokio::test]
#[ignore = "requires IMAGEFORGE_TEST_DATABASE_URL and IMAGEFORGE_TEST_REDIS_URL"]
async fn stale_operation_is_fenced_and_replay_remains_atomic() {
let database_url = std::env::var("IMAGEFORGE_TEST_DATABASE_URL")
.expect("IMAGEFORGE_TEST_DATABASE_URL is required");
let redis_url = std::env::var("IMAGEFORGE_TEST_REDIS_URL")
.expect("IMAGEFORGE_TEST_REDIS_URL is required");
let state = test_state(database_url, redis_url).await;
let marker = Uuid::new_v4().simple().to_string();
let user_id: Uuid = sqlx::query_scalar(
r#"
INSERT INTO users (email, username, password_hash, email_verified_at)
VALUES ($1, $2, 'test', NOW())
RETURNING id
"#,
)
.bind(format!("idem-{marker}@example.test"))
.bind(format!("idem-{marker}"))
.fetch_one(&state.db)
.await
.expect("insert idempotency test user");
let scope = Scope::User(user_id);
let key = format!("idem-{marker}");
let request_hash = "a".repeat(64);
let owner_one = match begin(&state, scope, &key, &request_hash, 24)
.await
.expect("acquire first operation")
{
BeginResult::Acquired { owner } => owner,
other => panic!("unexpected first begin result: {other:?}"),
};
assert!(matches!(
begin(&state, scope, &key, &request_hash, 24)
.await
.expect("probe live operation"),
BeginResult::InProgress
));
sqlx::query(
"UPDATE idempotency_keys SET lease_until = NOW() - INTERVAL '1 second' WHERE user_id = $1 AND idempotency_key = $2",
)
.bind(user_id)
.bind(&key)
.execute(&state.db)
.await
.expect("expire first operation lease");
let owner_two = match begin(&state, scope, &key, &request_hash, 24)
.await
.expect("take over stale operation")
{
BeginResult::Acquired { owner } => owner,
other => panic!("unexpected takeover result: {other:?}"),
};
assert_ne!(owner_one, owner_two);
let mut stale_tx = state.db.begin().await.expect("begin stale completion tx");
let stale_error = complete_in_tx(
&mut stale_tx,
scope,
&key,
&request_hash,
owner_one,
200,
serde_json::json!({"owner": "stale"}),
)
.await
.expect_err("stale operation completed after takeover");
assert_eq!(stale_error.code, ErrorCode::IdempotencyConflict);
stale_tx
.rollback()
.await
.expect("rollback stale completion");
let expected = serde_json::json!({"owner": "current"});
let mut current_tx = state.db.begin().await.expect("begin current completion tx");
complete_in_tx(
&mut current_tx,
scope,
&key,
&request_hash,
owner_two,
200,
expected.clone(),
)
.await
.expect("complete current operation");
current_tx
.commit()
.await
.expect("commit current completion");
match begin(&state, scope, &key, &request_hash, 24)
.await
.expect("replay completed operation")
{
BeginResult::Replay { response_body } => assert_eq!(response_body, expected),
other => panic!("unexpected replay result: {other:?}"),
}
sqlx::query("DELETE FROM users WHERE id = $1")
.bind(user_id)
.execute(&state.db)
.await
.expect("clean idempotency test user");
}
}

View File

@@ -6,7 +6,9 @@ pub mod filename;
pub mod idempotency;
pub mod mail;
pub mod metrics;
pub mod object_lifecycle;
pub mod quota;
pub mod rate_limit;
pub mod settings;
pub mod storage;
pub mod task_queue;

File diff suppressed because it is too large Load Diff

View File

@@ -330,27 +330,7 @@ fn retention_prefix(hours: i64) -> String {
}
}
pub async fn store_bytes<B>(
state: &AppState,
key: &str,
bytes: B,
content_type: &str,
) -> Result<StoredObject, AppError>
where
B: Into<Bytes>,
{
let bytes = bytes.into();
if let Some(endpoint) = active_endpoint(state).await? {
match store_bytes_s3(state, &endpoint, key, bytes.clone(), content_type).await {
Ok(stored) => return Ok(stored),
Err(err) => log_local_fallback(state, &endpoint, key, &err),
}
}
store_bytes_local(state, key, bytes.as_ref()).await
}
async fn store_bytes_s3(
pub(crate) async fn store_bytes_s3(
state: &AppState,
endpoint: &StorageEndpoint,
key: &str,
@@ -378,7 +358,7 @@ async fn store_bytes_s3(
})
}
async fn store_bytes_local(
pub(crate) async fn store_bytes_local(
state: &AppState,
key: &str,
bytes: &[u8],
@@ -402,27 +382,7 @@ async fn store_bytes_local(
})
}
pub async fn store_file(
state: &AppState,
key: &str,
path: &Path,
content_type: &str,
) -> Result<StoredObject, AppError> {
let metadata = tokio::fs::metadata(path).await.map_err(|err| {
AppError::new(ErrorCode::StorageUnavailable, "读取待上传文件失败").with_source(err)
})?;
if let Some(endpoint) = active_endpoint(state).await? {
match store_file_s3(state, &endpoint, key, path, content_type, metadata.len()).await {
Ok(stored) => return Ok(stored),
Err(err) => log_local_fallback(state, &endpoint, key, &err),
}
}
store_file_local(state, key, path, metadata.len()).await
}
async fn store_file_s3(
pub(crate) async fn store_file_s3(
state: &AppState,
endpoint: &StorageEndpoint,
key: &str,
@@ -459,7 +419,7 @@ async fn store_file_s3(
})
}
async fn store_file_local(
pub(crate) async fn store_file_local(
state: &AppState,
key: &str,
path: &Path,
@@ -484,7 +444,12 @@ async fn store_file_local(
})
}
fn log_local_fallback(state: &AppState, endpoint: &StorageEndpoint, key: &str, err: &AppError) {
pub(crate) fn log_local_fallback(
state: &AppState,
endpoint: &StorageEndpoint,
key: &str,
err: &AppError,
) {
crate::services::metrics::record_storage_fallback(state);
tracing::warn!(
storage_endpoint_id = %endpoint.id,
@@ -816,7 +781,7 @@ async fn endpoint_for_object(
get_endpoint(state, endpoint_id).await
}
fn local_path(state: &AppState, key: &str) -> Result<PathBuf, AppError> {
pub(crate) fn local_path(state: &AppState, key: &str) -> Result<PathBuf, AppError> {
if key.is_empty()
|| key.starts_with('/')
|| key.starts_with('\\')

367
src/services/task_queue.rs Normal file
View File

@@ -0,0 +1,367 @@
use crate::error::{AppError, ErrorCode};
use crate::services::{metrics, object_lifecycle, quota};
use crate::state::AppState;
use chrono::Utc;
use sqlx::FromRow;
use std::time::Duration;
use uuid::Uuid;
const DISPATCH_INTERVAL: Duration = Duration::from_secs(1);
const DISPATCH_BATCH_SIZE: i64 = 50;
const DELIVERY_LEASE_SECONDS: i64 = 30;
const MAX_DELIVERY_ATTEMPTS: i32 = 20;
const MAX_RETRY_SECONDS: u64 = 60;
#[derive(Debug, FromRow)]
struct OutboxClaim {
task_id: Uuid,
attempts: i32,
}
pub async fn dispatch_loop(state: AppState) {
let mut interval = tokio::time::interval(DISPATCH_INTERVAL);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
loop {
interval.tick().await;
if let Err(err) = dispatch_ready(&state, DISPATCH_BATCH_SIZE).await {
tracing::error!(error = %err, "task queue outbox dispatch iteration failed");
}
}
}
pub async fn dispatch_ready(state: &AppState, limit: i64) -> Result<usize, AppError> {
reconcile_non_pending_tasks(state).await?;
let lease_owner = Uuid::new_v4();
let claims = claim_ready(state, lease_owner, limit.max(1), None).await?;
let count = claims.len();
for claim in claims {
dispatch_claim(state, lease_owner, claim).await?;
}
Ok(count)
}
pub async fn dispatch_task(state: &AppState, task_id: Uuid) -> Result<bool, AppError> {
let lease_owner = Uuid::new_v4();
let mut claims = claim_ready(state, lease_owner, 1, Some(task_id)).await?;
let Some(claim) = claims.pop() else {
return Ok(false);
};
dispatch_claim(state, lease_owner, claim).await?;
Ok(true)
}
async fn claim_ready(
state: &AppState,
lease_owner: Uuid,
limit: i64,
task_id: Option<Uuid>,
) -> Result<Vec<OutboxClaim>, AppError> {
sqlx::query_as::<_, OutboxClaim>(
r#"
WITH candidate AS (
SELECT outbox.task_id
FROM task_queue_outbox AS outbox
JOIN tasks AS task ON task.id = outbox.task_id
WHERE outbox.status IN ('pending', 'delivering')
AND outbox.next_attempt_at <= NOW()
AND (outbox.lease_until IS NULL OR outbox.lease_until <= NOW())
AND task.status = 'pending'
AND task.deletion_started_at IS NULL
AND ($3::uuid IS NULL OR outbox.task_id = $3)
ORDER BY outbox.next_attempt_at, outbox.created_at
FOR UPDATE OF outbox SKIP LOCKED
LIMIT $2
)
UPDATE task_queue_outbox AS outbox
SET status = 'delivering',
attempts = outbox.attempts + 1,
lease_owner = $1,
lease_until = NOW() + ($4 * INTERVAL '1 second'),
updated_at = NOW()
FROM candidate
WHERE outbox.task_id = candidate.task_id
RETURNING outbox.task_id, outbox.attempts
"#,
)
.bind(lease_owner)
.bind(limit)
.bind(task_id)
.bind(DELIVERY_LEASE_SECONDS)
.fetch_all(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "领取任务队列 outbox 失败").with_source(err))
}
async fn dispatch_claim(
state: &AppState,
lease_owner: Uuid,
claim: OutboxClaim,
) -> Result<(), AppError> {
match enqueue_task(state, claim.task_id).await {
Ok(()) => mark_delivered(state, claim.task_id, lease_owner).await,
Err(err) if claim.attempts >= MAX_DELIVERY_ATTEMPTS => {
dead_letter_pending_task(state, claim.task_id, lease_owner, &err).await
}
Err(err) => release_for_retry(state, claim, lease_owner, &err).await,
}
}
pub(crate) async fn enqueue_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
let mut connection = state.redis.clone();
redis::cmd("XADD")
.arg(metrics::QUEUE_STREAM_KEY)
.arg("MAXLEN")
.arg("~")
.arg(100_000)
.arg("*")
.arg("task_id")
.arg(task_id.to_string())
.arg("created_at")
.arg(Utc::now().to_rfc3339())
.query_async::<_, redis::Value>(&mut connection)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "写入队列失败").with_source(err))?;
Ok(())
}
async fn mark_delivered(
state: &AppState,
task_id: Uuid,
lease_owner: Uuid,
) -> Result<(), AppError> {
sqlx::query(
r#"
UPDATE task_queue_outbox
SET status = 'delivered',
delivered_at = COALESCE(delivered_at, NOW()),
lease_owner = NULL,
lease_until = NULL,
last_error = NULL,
updated_at = NOW()
WHERE task_id = $1
AND status = 'delivering'
AND lease_owner = $2
"#,
)
.bind(task_id)
.bind(lease_owner)
.execute(&state.db)
.await
.map_err(|err| {
AppError::new(ErrorCode::Internal, "完成任务队列 outbox 失败").with_source(err)
})?;
Ok(())
}
async fn release_for_retry(
state: &AppState,
claim: OutboxClaim,
lease_owner: Uuid,
error: &AppError,
) -> Result<(), AppError> {
let delay = retry_delay(claim.attempts);
let message = truncate_error(error);
sqlx::query(
r#"
UPDATE task_queue_outbox
SET status = 'pending',
next_attempt_at = NOW() + ($3 * INTERVAL '1 second'),
lease_owner = NULL,
lease_until = NULL,
last_error = $4,
updated_at = NOW()
WHERE task_id = $1
AND status = 'delivering'
AND lease_owner = $2
"#,
)
.bind(claim.task_id)
.bind(lease_owner)
.bind(delay.as_secs() as i64)
.bind(&message)
.execute(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "记录任务队列重试失败").with_source(err))?;
tracing::warn!(task_id = %claim.task_id, attempts = claim.attempts, retry_seconds = delay.as_secs(), error = %error, "task queue delivery deferred");
Ok(())
}
async fn dead_letter_pending_task(
state: &AppState,
task_id: Uuid,
lease_owner: Uuid,
error: &AppError,
) -> Result<(), AppError> {
let message = format!("队列持续不可用:{}", truncate_error(error));
let input_dir = std::path::PathBuf::from(&state.config.storage_path)
.join("orig")
.join(task_id.to_string())
.to_string_lossy()
.to_string();
let mut tx = state.db.begin().await.map_err(|err| {
AppError::new(ErrorCode::Internal, "开启 outbox 死信事务失败").with_source(err)
})?;
let task: Option<String> =
sqlx::query_scalar("SELECT status::text FROM tasks WHERE id = $1 FOR UPDATE")
.bind(task_id)
.fetch_optional(&mut *tx)
.await
.map_err(|err| {
AppError::new(ErrorCode::Internal, "锁定 outbox 死信任务失败").with_source(err)
})?;
let Some(status) = task else {
tx.rollback().await.ok();
return Ok(());
};
let owned: Option<Uuid> = sqlx::query_scalar(
r#"
SELECT task_id
FROM task_queue_outbox
WHERE task_id = $1
AND status = 'delivering'
AND lease_owner = $2
FOR UPDATE
"#,
)
.bind(task_id)
.bind(lease_owner)
.fetch_optional(&mut *tx)
.await
.map_err(|err| {
AppError::new(ErrorCode::Internal, "校验 outbox 死信租约失败").with_source(err)
})?;
if owned.is_none() {
tx.rollback().await.ok();
return Ok(());
}
let failed = status == "pending";
if failed {
sqlx::query(
r#"
UPDATE task_files
SET status = 'failed',
error_message = $2,
completed_at = NOW(),
input_path = NULL,
lease_owner = NULL,
lease_until = NULL
WHERE task_id = $1
AND status IN ('pending', 'processing')
"#,
)
.bind(task_id)
.bind(&message)
.execute(&mut *tx)
.await
.map_err(|err| {
AppError::new(ErrorCode::Internal, "收口 outbox 死信文件失败").with_source(err)
})?;
sqlx::query(
r#"
INSERT INTO storage_objects (
task_id, object_kind, state, backend, object_key
) VALUES ($1, 'input_dir', 'delete_pending', 'local_dir', $2)
ON CONFLICT DO NOTHING
"#,
)
.bind(task_id)
.bind(&input_dir)
.execute(&mut *tx)
.await
.map_err(|err| {
AppError::new(ErrorCode::Internal, "安排 outbox 死信输入清理失败").with_source(err)
})?;
sqlx::query(
r#"
UPDATE tasks
SET status = 'failed',
completed_files = 0,
failed_files = total_files,
error_message = $2,
completed_at = NOW(),
lease_owner = NULL,
lease_until = NULL
WHERE id = $1 AND status = 'pending'
"#,
)
.bind(task_id)
.bind(&message)
.execute(&mut *tx)
.await
.map_err(|err| {
AppError::new(ErrorCode::Internal, "收口 outbox 死信任务失败").with_source(err)
})?;
}
sqlx::query(
r#"
UPDATE task_queue_outbox
SET status = CASE WHEN $3 THEN 'dead' ELSE 'delivered' END,
delivered_at = CASE WHEN $3 THEN delivered_at ELSE COALESCE(delivered_at, NOW()) END,
lease_owner = NULL,
lease_until = NULL,
last_error = $4,
updated_at = NOW()
WHERE task_id = $1
AND lease_owner = $2
"#,
)
.bind(task_id)
.bind(lease_owner)
.bind(failed)
.bind(&message)
.execute(&mut *tx)
.await
.map_err(|err| {
AppError::new(ErrorCode::Internal, "提交 outbox 死信状态失败").with_source(err)
})?;
tx.commit().await.map_err(|err| {
AppError::new(ErrorCode::Internal, "提交 outbox 死信事务失败").with_source(err)
})?;
if failed {
metrics::record_dead_letter(state);
if let Err(err) = object_lifecycle::cleanup_ready_objects(state, 10, Some(task_id)).await {
tracing::warn!(task_id = %task_id, error = %err, "dead outbox input cleanup deferred");
}
quota::settle_anonymous_task_reservation(state, task_id).await?;
}
Ok(())
}
async fn reconcile_non_pending_tasks(state: &AppState) -> Result<(), AppError> {
sqlx::query(
r#"
UPDATE task_queue_outbox AS outbox
SET status = 'delivered',
delivered_at = COALESCE(outbox.delivered_at, NOW()),
lease_owner = NULL,
lease_until = NULL,
updated_at = NOW()
FROM tasks AS task
WHERE task.id = outbox.task_id
AND task.status <> 'pending'
AND outbox.status IN ('pending', 'delivering')
"#,
)
.execute(&state.db)
.await
.map_err(|err| {
AppError::new(ErrorCode::Internal, "对账任务队列 outbox 失败").with_source(err)
})?;
Ok(())
}
fn retry_delay(attempts: i32) -> Duration {
let exponent = attempts.saturating_sub(1).min(6) as u32;
Duration::from_secs(2_u64.saturating_pow(exponent).min(MAX_RETRY_SECONDS))
}
fn truncate_error(error: &AppError) -> String {
format!("{}: {}", error.code.as_str(), error.message)
.chars()
.take(2_000)
.collect()
}

View File

@@ -2,6 +2,7 @@ use crate::error::{AppError, ErrorCode};
use crate::services::billing;
use crate::services::compress;
use crate::services::metrics;
use crate::services::object_lifecycle;
use crate::services::quota;
use crate::services::storage;
use crate::state::AppState;
@@ -40,6 +41,10 @@ pub async fn run(state: AppState) -> Result<(), AppError> {
let consumer = format!("worker_{worker_id}");
ensure_group(&state).await?;
tokio::spawn(maintenance_loop(state.clone()));
tokio::spawn(crate::services::task_queue::dispatch_loop(state.clone()));
tokio::spawn(crate::services::object_lifecycle::maintenance_loop(
state.clone(),
));
let task_concurrency = state.config.worker_task_concurrency.max(1) as usize;
let mut inflight = JoinSet::new();
@@ -675,21 +680,6 @@ struct TaskFileProcRow {
output_format: String,
}
#[derive(Debug, FromRow)]
struct CleanupFileRow {
storage_backend: String,
storage_endpoint_id: Option<Uuid>,
storage_key: Option<String>,
input_path: Option<String>,
}
#[derive(Debug, FromRow)]
struct CleanupZipRow {
zip_storage_backend: Option<String>,
zip_storage_endpoint_id: Option<Uuid>,
zip_storage_key: Option<String>,
}
#[derive(Clone)]
struct TaskContext {
api_key_id: Option<Uuid>,
@@ -729,6 +719,7 @@ pub(crate) async fn process_task(
lease_owner = $2,
lease_until = NOW() + $3 * INTERVAL '1 second'
WHERE id = $1
AND deletion_started_at IS NULL
AND (
status = 'pending'
OR (
@@ -940,6 +931,7 @@ async fn file_attempt_is_current(state: &AppState, fence: &FileFence) -> Result<
JOIN task_files f ON f.task_id = t.id
WHERE t.id = $1
AND t.status = 'processing'
AND t.deletion_started_at IS NULL
AND t.processing_attempt = $2
AND t.lease_owner = $5
AND t.lease_until > NOW()
@@ -1103,10 +1095,13 @@ async fn process_task_file(
file_attempt,
format_out.extension(),
);
let stored = match storage::store_bytes(
let tracked = match object_lifecycle::store_tracked_bytes(
&state,
task_id,
Some(file.id),
"result",
&object_key,
compressed,
compressed.into(),
format_out.content_type(),
)
.await
@@ -1122,19 +1117,19 @@ async fn process_task_file(
if ctx.is_anonymous && charge_units && !ctx.anonymous_quota_reserved {
let Some(session_id) = ctx.session_id.as_deref() else {
let _ = storage::delete_object(&state, &stored_locator(&stored)).await;
discard_tracked_result(&state, &tracked, None).await;
mark_file_failed_and_cleanup(&state, &fence, "匿名任务缺少 session_id", &input_path)
.await?;
return Ok(());
};
let Some(ip) = ctx.anon_ip else {
let _ = storage::delete_object(&state, &stored_locator(&stored)).await;
discard_tracked_result(&state, &tracked, None).await;
mark_file_failed_and_cleanup(&state, &fence, "匿名任务缺少 client_ip", &input_path)
.await?;
return Ok(());
};
if let Err(err) = quota::consume_anonymous_units(&state, session_id, ip, 1).await {
let _ = storage::delete_object(&state, &stored_locator(&stored)).await;
discard_tracked_result(&state, &tracked, Some(&err)).await;
mark_file_failed_and_cleanup(&state, &fence, &err.message, &input_path).await?;
return Ok(());
}
@@ -1146,7 +1141,7 @@ async fn process_task_file(
ctx.api_key_id,
&ctx.source,
&fence,
&stored,
&tracked,
original_size as i64,
compressed_size as i64,
saved_percent,
@@ -1160,16 +1155,85 @@ async fn process_task_file(
let _ = tokio::fs::remove_file(&input_path).await;
}
Ok(FinalizeFileOutcome::LeaseLost) => {
let _ = storage::delete_object(&state, &stored_locator(&stored)).await;
}
Err(err) => {
let _ = storage::delete_object(&state, &stored_locator(&stored)).await;
mark_file_failed_and_cleanup(&state, &fence, &err.message, &input_path).await?;
discard_tracked_result(&state, &tracked, None).await;
}
Err(err) => match worker_result_was_committed(&state, &fence, &tracked).await {
Ok(true) => {
tracing::warn!(task_id = %task_id, file_id = %fence.file_id, error = %err, "worker result commit response was lost; recovered committed publication");
let _ = tokio::fs::remove_file(&input_path).await;
}
Ok(false) => {
discard_tracked_result(&state, &tracked, Some(&err)).await;
mark_file_failed_and_cleanup(&state, &fence, &err.message, &input_path).await?;
}
Err(probe_err) => {
tracing::error!(task_id = %task_id, file_id = %fence.file_id, error = %probe_err, original_error = %err, "worker result commit state is unknown; staging lease will reconcile object");
return Err(err);
}
},
}
Ok(())
}
async fn discard_tracked_result(
state: &AppState,
tracked: &object_lifecycle::TrackedStoredObject,
error: Option<&AppError>,
) {
if let Err(schedule_err) =
object_lifecycle::schedule_tracked_delete(state, tracked, error).await
{
tracing::error!(storage_object_id = %tracked.lifecycle_id, error = %schedule_err, "failed to persist discarded worker object cleanup");
return;
}
if let Err(cleanup_err) =
object_lifecycle::cleanup_ready_objects(state, 1, Some(tracked.task_id)).await
{
tracing::warn!(storage_object_id = %tracked.lifecycle_id, error = %cleanup_err, "discarded worker object cleanup deferred");
}
}
async fn worker_result_was_committed(
state: &AppState,
fence: &FileFence,
tracked: &object_lifecycle::TrackedStoredObject,
) -> Result<bool, AppError> {
sqlx::query_scalar(
r#"
SELECT EXISTS(
SELECT 1
FROM tasks AS task
JOIN task_files AS file ON file.task_id = task.id
JOIN storage_objects AS object ON object.id = $3
WHERE task.id = $1
AND file.id = $2
AND file.status = 'completed'
AND file.storage_backend = $4
AND file.storage_endpoint_id IS NOT DISTINCT FROM $5
AND COALESCE(file.storage_key, file.storage_path) = $6
AND object.state = 'published'
AND object.task_id = task.id
AND object.task_file_id = file.id
)
"#,
)
.bind(fence.task_id)
.bind(fence.file_id)
.bind(tracked.lifecycle_id)
.bind(&tracked.stored.backend)
.bind(tracked.stored.endpoint_id)
.bind(&tracked.stored.key)
.fetch_one(&state.db)
.await
.map_err(|err| {
AppError::new(
ErrorCode::StorageUnavailable,
"核验 Worker 结果提交状态失败",
)
.with_source(err)
})
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum FinalizeFileOutcome {
Committed,
@@ -1183,7 +1247,7 @@ async fn finalize_file(
api_key_id: Option<Uuid>,
source: &str,
fence: &FileFence,
stored: &storage::StoredObject,
tracked: &object_lifecycle::TrackedStoredObject,
bytes_in: i64,
bytes_out: i64,
saved_percent: f64,
@@ -1191,6 +1255,7 @@ async fn finalize_file(
format_out: compress::ImageFmt,
charge_units: bool,
) -> Result<FinalizeFileOutcome, AppError> {
let stored = &tracked.stored;
let mut tx = state
.db
.begin()
@@ -1332,6 +1397,8 @@ async fn finalize_file(
return Ok(FinalizeFileOutcome::LeaseLost);
}
object_lifecycle::publish_in_tx(&mut tx, tracked).await?;
tx.commit()
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "提交事务失败").with_source(err))?;
@@ -1840,6 +1907,12 @@ async fn cleanup_expired_records(state: &AppState) -> Result<(), AppError> {
.execute(&state.db)
.await;
let _ = sqlx::query(
"DELETE FROM storage_objects WHERE state = 'deleted' AND deleted_at < NOW() - INTERVAL '7 days'",
)
.execute(&state.db)
.await;
let _ =
sqlx::query("DELETE FROM webhook_events WHERE received_at < NOW() - INTERVAL '90 days'")
.execute(&state.db)
@@ -1851,6 +1924,7 @@ async fn cleanup_expired_records(state: &AppState) -> Result<(), AppError> {
WHERE e.deleted_at < NOW() - INTERVAL '30 days'
AND NOT EXISTS (SELECT 1 FROM task_files f WHERE f.storage_endpoint_id = e.id)
AND NOT EXISTS (SELECT 1 FROM tasks t WHERE t.zip_storage_endpoint_id = e.id)
AND NOT EXISTS (SELECT 1 FROM storage_objects o WHERE o.storage_endpoint_id = e.id)
"#,
)
.execute(&state.db)
@@ -1894,82 +1968,13 @@ async fn cleanup_expired_tasks(state: &AppState) -> Result<(), AppError> {
}
async fn cleanup_expired_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
sqlx::query(
"UPDATE tasks SET status = 'cancelled', completed_at = COALESCE(completed_at, NOW()) WHERE id = $1 AND expires_at < NOW() AND status IN ('pending', 'processing')",
)
.bind(task_id)
.execute(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "终止过期任务失败").with_source(err))?;
quota::settle_anonymous_task_reservation(state, task_id).await?;
let files: Vec<CleanupFileRow> = sqlx::query_as(
r#"
SELECT storage_backend, storage_endpoint_id,
COALESCE(storage_key, storage_path) AS storage_key,
input_path
FROM task_files
WHERE task_id = $1
"#,
)
.bind(task_id)
.fetch_all(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "查询过期任务文件失败").with_source(err))?;
for file in files {
if let Some(key) = file.storage_key {
storage::delete_object(
state,
&storage::ObjectLocator {
backend: file.storage_backend,
endpoint_id: file.storage_endpoint_id,
key,
},
)
.await?;
}
if let Some(input_path) = file.input_path {
let _ = tokio::fs::remove_file(input_path).await;
}
if object_lifecycle::mark_expired_task(state, task_id).await? {
object_lifecycle::finalize_task_deletion(state, task_id).await?;
}
let zip: Option<CleanupZipRow> = sqlx::query_as(
"SELECT zip_storage_backend, zip_storage_endpoint_id, zip_storage_key FROM tasks WHERE id = $1",
)
.bind(task_id)
.fetch_optional(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "查询过期 ZIP 失败").with_source(err))?;
if let Some(zip) = zip {
if let (Some(backend), Some(key)) = (zip.zip_storage_backend, zip.zip_storage_key) {
storage::delete_object(
state,
&storage::ObjectLocator {
backend,
endpoint_id: zip.zip_storage_endpoint_id,
key,
},
)
.await?;
}
}
let legacy_zip_path = format!("{}/zips/{task_id}.zip", state.config.storage_path);
let _ = tokio::fs::remove_file(legacy_zip_path).await;
let orig_dir = format!("{}/orig/{task_id}", state.config.storage_path);
let _ = tokio::fs::remove_dir_all(orig_dir).await;
sqlx::query("DELETE FROM tasks WHERE id = $1 AND expires_at < NOW()")
.bind(task_id)
.execute(&state.db)
.await
.map_err(|err| {
AppError::new(ErrorCode::Internal, "删除过期任务记录失败").with_source(err)
})?;
Ok(())
}
#[cfg(test)]
fn stored_locator(stored: &storage::StoredObject) -> storage::ObjectLocator {
storage::ObjectLocator {
backend: stored.backend.clone(),
@@ -2148,16 +2153,22 @@ mod tests {
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");
let stale_object = storage::store_bytes(
let stale_object = object_lifecycle::store_tracked_bytes(
&state,
task_id,
Some(file_id),
"result",
&stale_key,
Bytes::from_static(b"stale-attempt"),
"image/png",
)
.await
.expect("store stale attempt object");
let winning_object = storage::store_bytes(
let winning_object = object_lifecycle::store_tracked_bytes(
&state,
task_id,
Some(file_id),
"result",
&winning_key,
Bytes::from_static(b"winning-attempt"),
"image/png",
@@ -2165,8 +2176,8 @@ mod tests {
.await
.expect("store winning attempt object");
if let Ok(expected_backend) = std::env::var("IMAGEFORGE_TEST_EXPECT_STORAGE_BACKEND") {
assert_eq!(stale_object.backend, expected_backend);
assert_eq!(winning_object.backend, expected_backend);
assert_eq!(stale_object.stored.backend, expected_backend);
assert_eq!(winning_object.stored.backend, expected_backend);
}
let period_start = Utc::now() - chrono::Duration::hours(1);
@@ -2259,14 +2270,14 @@ mod tests {
assert_eq!(stale_result, FinalizeFileOutcome::LeaseLost);
assert_eq!(winning_result, FinalizeFileOutcome::Committed);
storage::delete_object(&state, &stored_locator(&stale_object))
.await
.expect("delete stale attempt object");
assert!(storage::read_bytes(&state, &stored_locator(&stale_object))
.await
.is_err());
discard_tracked_result(&state, &stale_object, None).await;
assert!(
storage::read_bytes(&state, &stored_locator(&stale_object.stored))
.await
.is_err()
);
assert_eq!(
storage::read_bytes(&state, &stored_locator(&winning_object))
storage::read_bytes(&state, &stored_locator(&winning_object.stored))
.await
.expect("read winning object"),
b"winning-attempt"
@@ -2292,7 +2303,11 @@ mod tests {
.expect("query test file");
assert_eq!(
file,
("completed".to_string(), winning_object.key.clone(), 40)
(
"completed".to_string(),
winning_object.stored.key.clone(),
40
)
);
let usage_event_count: i64 =
sqlx::query_scalar("SELECT COUNT(*) FROM usage_events WHERE task_file_id = $1")
@@ -2312,9 +2327,7 @@ mod tests {
.expect("query used units");
assert_eq!(used_units, 1);
storage::delete_object(&state, &stored_locator(&winning_object))
.await
.expect("delete winning object");
discard_tracked_result(&state, &winning_object, None).await;
sqlx::query("DELETE FROM usage_events WHERE task_id = $1")
.bind(task_id)
.execute(&pool)
@@ -2325,6 +2338,11 @@ mod tests {
.execute(&pool)
.await
.expect("delete test task");
sqlx::query("DELETE FROM storage_objects WHERE task_id = $1")
.bind(task_id)
.execute(&pool)
.await
.expect("delete test storage lifecycle rows");
sqlx::query("DELETE FROM usage_periods WHERE user_id = $1")
.bind(user_id)
.execute(&pool)