237899745/ystp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(storage): make queue and object ownership durable

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

246
src/worker/mod.rs
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)