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

feat: add configurable S3 object storage

Browse Source
This commit is contained in:
237899745
2026-07-25 13:23:11 +08:00
parent 61fa9cb820
commit d1f093685d
29 changed files with 3703 additions and 284 deletions
Download Patch File Download Diff File Expand all files Collapse all files

272
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::quota;
use crate::services::storage;
use crate::state::AppState;
use redis::streams::StreamReadOptions;
@@ -148,17 +149,33 @@ struct TaskProcRow {
api_key_id: Option<Uuid>,
source: String,
client_ip: Option<String>,
retention_hours: i32,
}
#[derive(Debug, FromRow)]
struct TaskFileProcRow {
id: Uuid,
storage_path: Option<String>,
input_path: Option<String>,
original_format: String,
output_format: String,
status: 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>,
@@ -167,6 +184,7 @@ struct TaskContext {
session_id: Option<String>,
anon_ip: Option<IpAddr>,
is_anonymous: bool,
retention_hours: i32,
}
async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
@@ -183,7 +201,8 @@ async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
session_id,
api_key_id,
source::text AS source,
host(client_ip) AS client_ip
host(client_ip) AS client_ip,
retention_hours
FROM tasks
WHERE id = $1
"#,
@@ -241,7 +260,7 @@ async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
r#"
SELECT
id,
storage_path,
input_path,
original_format,
output_format,
status::text AS status
@@ -273,6 +292,7 @@ async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
session_id: task.session_id.clone(),
anon_ip,
is_anonymous: task.user_id.is_none(),
retention_hours: task.retention_hours,
};
let concurrency = state.config.worker_concurrency.max(1) as usize;
@@ -391,7 +411,7 @@ async fn process_task_file(
return Ok(());
}
let Some(input_path) = file.storage_path.clone() else {
let Some(input_path) = file.input_path.clone() else {
mark_file_failed(&state, task_id, file.id, "原文件不存在").await?;
return Ok(());
};
@@ -462,40 +482,50 @@ async fn process_task_file(
&& max_height.is_none();
let charge_units = !skip_charge && compressed_size < original_size;
let object_key = storage::result_key(
ctx.retention_hours as i64,
task_id,
file.id,
format_out.extension(),
);
let stored = match storage::store_bytes(
&state,
&object_key,
compressed,
format_out.content_type(),
)
.await
{
Ok(value) => value,
Err(err) => {
reset_file_for_retry(&state, file.id, "对象存储暂时不可用").await?;
return Err(err);
}
};
if ctx.is_anonymous && charge_units {
let Some(session_id) = ctx.session_id.as_deref() else {
let _ = storage::delete_object(&state, &stored_locator(&stored)).await;
mark_file_failed(&state, task_id, file.id, "匿名任务缺少 session_id").await?;
let _ = tokio::fs::remove_file(&input_path).await;
return Ok(());
};
let Some(ip) = ctx.anon_ip else {
let _ = storage::delete_object(&state, &stored_locator(&stored)).await;
mark_file_failed(&state, task_id, file.id, "匿名任务缺少 client_ip").await?;
let _ = tokio::fs::remove_file(&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;
mark_file_failed(&state, task_id, file.id, &err.message).await?;
let _ = tokio::fs::remove_file(&input_path).await;
return Ok(());
}
}
let output_path = format!(
"{}/{}.{}",
state.config.storage_path,
file.id,
format_out.extension()
);
if let Err(err) = tokio::fs::write(&output_path, &compressed).await {
mark_file_failed(&state, task_id, file.id, "写入压缩文件失败").await?;
let _ = tokio::fs::remove_file(&input_path).await;
return Err(
AppError::new(ErrorCode::StorageUnavailable, "写入压缩文件失败").with_source(err),
);
}
if is_task_cancelled(&state, task_id).await? {
let _ = tokio::fs::remove_file(&output_path).await;
let _ = storage::delete_object(&state, &stored_locator(&stored)).await;
mark_file_failed(&state, task_id, file.id, "已取消").await?;
let _ = tokio::fs::remove_file(&input_path).await;
return Ok(());
@@ -508,7 +538,7 @@ async fn process_task_file(
&ctx.source,
task_id,
file.id,
&output_path,
&stored,
original_size as i64,
compressed_size as i64,
saved_percent,
@@ -518,12 +548,8 @@ async fn process_task_file(
)
.await
{
if err.code == ErrorCode::QuotaExceeded {
let _ = tokio::fs::remove_file(&output_path).await;
mark_file_failed(&state, task_id, file.id, &err.message).await?;
} else {
mark_file_failed(&state, task_id, file.id, &err.message).await?;
}
let _ = storage::delete_object(&state, &stored_locator(&stored)).await;
mark_file_failed(&state, task_id, file.id, &err.message).await?;
let _ = tokio::fs::remove_file(&input_path).await;
return Ok(());
}
@@ -540,7 +566,7 @@ async fn finalize_file(
source: &str,
task_id: Uuid,
task_file_id: Uuid,
output_path: &str,
stored: &storage::StoredObject,
bytes_in: i64,
bytes_out: i64,
saved_percent: f64,
@@ -577,15 +603,28 @@ async fn finalize_file(
r#"
UPDATE task_files
SET storage_path = $2,
compressed_size = $3,
saved_percent = $4,
storage_backend = $3,
storage_endpoint_id = $4,
storage_key = $5,
storage_etag = $6,
input_path = NULL,
compressed_size = $7,
saved_percent = $8,
status = 'completed',
completed_at = NOW()
WHERE id = $1
"#,
)
.bind(task_file_id)
.bind(output_path)
.bind(if stored.backend == "local" {
Some(stored.key.as_str())
} else {
None
})
.bind(&stored.backend)
.bind(stored.endpoint_id)
.bind(&stored.key)
.bind(&stored.etag)
.bind(bytes_out)
.bind(saved_percent)
.execute(&mut *tx)
@@ -631,6 +670,7 @@ async fn mark_file_failed(
SET status = 'failed',
error_message = $2,
storage_path = NULL,
input_path = NULL,
completed_at = NOW()
WHERE id = $1
AND status NOT IN ('completed', 'failed')
@@ -663,6 +703,22 @@ async fn mark_file_failed(
Ok(())
}
async fn reset_file_for_retry(
state: &AppState,
task_file_id: Uuid,
message: &str,
) -> Result<(), AppError> {
sqlx::query(
"UPDATE task_files SET status = 'pending', error_message = $2 WHERE id = $1 AND status = 'processing'",
)
.bind(task_file_id)
.bind(message)
.execute(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "恢复待重试文件失败").with_source(err))?;
Ok(())
}
async fn finalize_task_status(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
let row: Option<(i32, i32, i32, String)> = sqlx::query_as(
"SELECT total_files, completed_files, failed_files, status::text AS status FROM tasks WHERE id = $1",
@@ -677,7 +733,7 @@ async fn finalize_task_status(state: &AppState, task_id: Uuid) -> Result<(), App
};
if status == "cancelled" {
let paths: Vec<Option<String>> = sqlx::query_scalar(
"SELECT storage_path FROM task_files WHERE task_id = $1 AND status IN ('pending','processing')",
"SELECT input_path FROM task_files WHERE task_id = $1 AND status IN ('pending','processing')",
)
.bind(task_id)
.fetch_all(&state.db)
@@ -688,7 +744,7 @@ async fn finalize_task_status(state: &AppState, task_id: Uuid) -> Result<(), App
}
let _ = sqlx::query(
"UPDATE task_files SET status = 'failed', error_message = '已取消', storage_path = NULL, completed_at = NOW() WHERE task_id = $1 AND status IN ('pending','processing')",
"UPDATE task_files SET status = 'failed', error_message = '已取消', storage_path = NULL, input_path = NULL, completed_at = NOW() WHERE task_id = $1 AND status IN ('pending','processing')",
)
.bind(task_id)
.execute(&state.db)
@@ -810,10 +866,47 @@ async fn charge_one_unit(
async fn maintenance(state: &AppState) -> Result<(), AppError> {
cleanup_expired_tasks(state).await?;
cleanup_stale_zip_temp(state).await?;
cleanup_expired_records(state).await?;
Ok(())
}
async fn cleanup_stale_zip_temp(state: &AppState) -> Result<(), AppError> {
let root = std::path::Path::new(&state.config.storage_path).join("tmp/zips");
let mut entries = match tokio::fs::read_dir(&root).await {
Ok(entries) => entries,
Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(()),
Err(err) => {
return Err(
AppError::new(ErrorCode::StorageUnavailable, "读取 ZIP 临时目录失败")
.with_source(err),
)
}
};
let cutoff = std::time::SystemTime::now()
.checked_sub(std::time::Duration::from_secs(6 * 60 * 60))
.unwrap_or(std::time::UNIX_EPOCH);
while let Some(entry) = entries.next_entry().await.map_err(|err| {
AppError::new(ErrorCode::StorageUnavailable, "遍历 ZIP 临时目录失败").with_source(err)
})? {
let metadata = match entry.metadata().await {
Ok(metadata) => metadata,
Err(_) => continue,
};
if !metadata.is_dir()
|| metadata
.modified()
.map(|time| time >= cutoff)
.unwrap_or(true)
{
continue;
}
let _ = tokio::fs::remove_dir_all(entry.path()).await;
}
Ok(())
}
async fn cleanup_expired_records(state: &AppState) -> Result<(), AppError> {
let _ = sqlx::query("DELETE FROM idempotency_keys WHERE expires_at < NOW()")
.execute(&state.db)
@@ -834,6 +927,17 @@ async fn cleanup_expired_records(state: &AppState) -> Result<(), AppError> {
.execute(&state.db)
.await;
let _ = sqlx::query(
r#"
DELETE FROM storage_endpoints e
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)
"#,
)
.execute(&state.db)
.await;
Ok(())
}
@@ -848,33 +952,87 @@ async fn cleanup_expired_tasks(state: &AppState) -> Result<(), AppError> {
return Ok(());
}
if state.config.storage_type.eq_ignore_ascii_case("local") {
for task_id in &task_ids {
let paths: Vec<Option<String>> =
sqlx::query_scalar("SELECT storage_path FROM task_files WHERE task_id = $1")
.bind(task_id)
.fetch_all(&state.db)
.await
.unwrap_or_default();
for p in paths.into_iter().flatten() {
let _ = tokio::fs::remove_file(p).await;
}
let zip_path = format!("{}/zips/{task_id}.zip", state.config.storage_path);
let _ = tokio::fs::remove_file(zip_path).await;
let orig_dir = format!("{}/orig/{task_id}", state.config.storage_path);
let _ = tokio::fs::remove_dir_all(orig_dir).await;
}
}
for task_id in task_ids {
let _ = sqlx::query("DELETE FROM tasks WHERE id = $1")
.bind(task_id)
.execute(&state.db)
.await;
if let Err(err) = cleanup_expired_task(state, task_id).await {
tracing::warn!(task_id = %task_id, error = %err, "expired task cleanup deferred");
}
}
Ok(())
}
async fn cleanup_expired_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
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;
}
}
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(())
}
fn stored_locator(stored: &storage::StoredObject) -> storage::ObjectLocator {
storage::ObjectLocator {
backend: stored.backend.clone(),
endpoint_id: stored.endpoint_id,
key: stored.key.clone(),
}
}