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

perf: improve compression reliability and deployment safety

Browse Source
This commit is contained in:
237899745
2026-07-25 10:29:49 +08:00
parent 06220ca921
commit 9d7668bdee
34 changed files with 1391 additions and 1042 deletions
Download Patch File Download Diff File Expand all files Collapse all files

185
src/worker/mod.rs
View File

@@ -20,9 +20,7 @@ const GROUP_NAME: &str = "compress_workers";
pub async fn run(state: AppState) -> Result<(), AppError> {
tracing::info!("Worker started");
if let Err(err) = crate::services::bootstrap::ensure_schema(&state).await {
tracing::error!(error = %err, "数据库结构初始化失败");
}
crate::services::bootstrap::ensure_schema(&state).await?;
let consumer = format!("worker_{}", Uuid::new_v4());
ensure_group(&state, &consumer).await?;
@@ -71,15 +69,26 @@ async fn ensure_group(state: &AppState, _consumer: &str) -> Result<(), AppError>
async fn poll_once(state: &AppState, consumer: &str) -> Result<(), AppError> {
let mut conn = state.redis.clone();
let opts = StreamReadOptions::default()
// Retry messages already delivered to this consumer before taking new work.
let pending_opts = StreamReadOptions::default()
.group(GROUP_NAME, consumer)
.count(1)
.block(5000);
let reply: redis::streams::StreamReadReply = conn
.xread_options(&[STREAM_KEY], &[">"], &opts)
.count(1);
let mut reply: redis::streams::StreamReadReply = conn
.xread_options(&[STREAM_KEY], &["0"], &pending_opts)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "读取队列失败").with_source(err))?;
.map_err(|err| AppError::new(ErrorCode::Internal, "读取待重试任务失败").with_source(err))?;
if !reply.keys.iter().any(|key| !key.ids.is_empty()) {
let opts = StreamReadOptions::default()
.group(GROUP_NAME, consumer)
.count(1)
.block(5000);
reply = conn
.xread_options(&[STREAM_KEY], &[">"], &opts)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "读取队列失败").with_source(err))?;
}
if reply.keys.is_empty() {
return Ok(());
@@ -100,9 +109,10 @@ async fn poll_once(state: &AppState, consumer: &str) -> Result<(), AppError> {
}
};
if let Err(err) = process_task(state, task_id).await {
tracing::error!(task_id = %task_id, error = %err, "task processing failed");
}
process_task(state, task_id).await.map_err(|err| {
tracing::error!(task_id = %task_id, error = %err, "task processing failed; message left pending for retry");
err
})?;
ack_message(&mut conn, &msg.id).await?;
}
@@ -111,7 +121,10 @@ async fn poll_once(state: &AppState, consumer: &str) -> Result<(), AppError> {
Ok(())
}
async fn ack_message(conn: &mut redis::aio::ConnectionManager, msg_id: &str) -> Result<(), AppError> {
async fn ack_message(
conn: &mut redis::aio::ConnectionManager,
msg_id: &str,
) -> Result<(), AppError> {
let _: i64 = redis::cmd("XACK")
.arg(STREAM_KEY)
.arg(GROUP_NAME)
@@ -124,16 +137,12 @@ async fn ack_message(conn: &mut redis::aio::ConnectionManager, msg_id: &str) ->
#[derive(Debug, FromRow)]
struct TaskProcRow {
id: Uuid,
status: String,
compression_level: String,
compression_rate: Option<i16>,
max_width: Option<i32>,
max_height: Option<i32>,
preserve_metadata: bool,
total_files: i32,
completed_files: i32,
failed_files: i32,
user_id: Option<Uuid>,
session_id: Option<String>,
api_key_id: Option<Uuid>,
@@ -145,10 +154,8 @@ struct TaskProcRow {
struct TaskFileProcRow {
id: Uuid,
storage_path: Option<String>,
original_name: String,
original_format: String,
output_format: String,
original_size: i64,
status: String,
}
@@ -166,16 +173,12 @@ async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
let mut task: TaskProcRow = sqlx::query_as(
r#"
SELECT
id,
status::text AS status,
compression_level::text AS compression_level,
compression_rate,
max_width,
max_height,
preserve_metadata,
total_files,
completed_files,
failed_files,
user_id,
session_id,
api_key_id,
@@ -195,6 +198,8 @@ async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
return Ok(());
}
let is_retry = task.status == "processing";
let updated = sqlx::query(
r#"
UPDATE tasks
@@ -215,9 +220,17 @@ async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
// Refresh task row after status change
task.status = "processing".to_string();
let compression_rate = task
.compression_rate
.and_then(|v| u8::try_from(v).ok());
if is_retry {
sqlx::query(
"UPDATE task_files SET status = 'pending' WHERE task_id = $1 AND status = 'processing'",
)
.bind(task_id)
.execute(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "恢复待重试文件失败").with_source(err))?;
}
let compression_rate = task.compression_rate.and_then(|v| u8::try_from(v).ok());
let level = compression_rate
.map(compress::rate_to_level)
.unwrap_or(compress::parse_level(&task.compression_level)?);
@@ -229,10 +242,8 @@ async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
SELECT
id,
storage_path,
original_name,
original_format,
output_format,
original_size,
status::text AS status
FROM task_files
WHERE task_id = $1
@@ -281,7 +292,7 @@ async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
join_set.spawn(async move {
let _permit = permit;
if let Err(err) = process_task_file(
let result = process_task_file(
state,
task_id,
file,
@@ -292,19 +303,35 @@ async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
ctx,
billing_ctx,
)
.await
{
.await;
if let Err(err) = &result {
tracing::error!(task_id = %task_id, file_id = %file_id, error = %err, "file processing failed");
}
result
});
}
let mut first_error = None;
while let Some(result) = join_set.join_next().await {
if let Err(err) = result {
tracing::error!(task_id = %task_id, error = %err, "file worker panicked");
match result {
Ok(Ok(())) => {}
Ok(Err(err)) => {
if first_error.is_none() {
first_error = Some(err);
}
}
Err(err) => {
return Err(
AppError::new(ErrorCode::Internal, "文件处理线程异常退出").with_source(err)
);
}
}
}
if let Some(err) = first_error {
return Err(err);
}
finalize_task_status(state, task_id).await?;
Ok(())
}
@@ -332,6 +359,7 @@ async fn is_task_cancelled(state: &AppState, task_id: Uuid) -> Result<bool, AppE
Ok(matches!(status.as_deref(), Some("cancelled")))
}
#[allow(clippy::too_many_arguments)]
async fn process_task_file(
state: AppState,
task_id: Uuid,
@@ -347,11 +375,13 @@ async fn process_task_file(
return Ok(());
}
let updated = sqlx::query("UPDATE task_files SET status = 'processing' WHERE id = $1 AND status = 'pending'")
.bind(file.id)
.execute(&state.db)
.await
.unwrap_or_else(|_| sqlx::postgres::PgQueryResult::default());
let updated = sqlx::query(
"UPDATE task_files SET status = 'processing' WHERE id = $1 AND status = 'pending'",
)
.bind(file.id)
.execute(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "更新文件处理状态失败").with_source(err))?;
if updated.rows_affected() == 0 {
return Ok(());
}
@@ -374,17 +404,32 @@ async fn process_task_file(
}
};
let format_in = parse_image_fmt(&file.original_format)?;
let format_out = parse_image_fmt(&file.output_format)?;
let format_in = match parse_image_fmt(&file.original_format) {
Ok(format) => format,
Err(err) => {
mark_file_failed(&state, task_id, file.id, &err.message).await?;
let _ = tokio::fs::remove_file(&input_path).await;
return Ok(());
}
};
let format_out = match parse_image_fmt(&file.output_format) {
Ok(format) => format,
Err(err) => {
mark_file_failed(&state, task_id, file.id, &err.message).await?;
let _ = tokio::fs::remove_file(&input_path).await;
return Ok(());
}
};
let original_size = input_bytes.len() as u64;
let compressed = match compress::compress_image_bytes(
&state,
&input_bytes,
input_bytes,
format_in,
format_out,
level,
compression_rate,
None, // target_size_bytes: worker 批量任务不支持精确大小
None, // target_size_bytes: worker 批量任务不支持精确大小
max_width,
max_height,
ctx.preserve_metadata,
@@ -405,7 +450,6 @@ async fn process_task_file(
return Ok(());
}
let original_size = input_bytes.len() as u64;
let compressed_size = compressed.len() as u64;
let saved_percent = if original_size == 0 {
0.0
@@ -442,7 +486,9 @@ async fn process_task_file(
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));
return Err(
AppError::new(ErrorCode::StorageUnavailable, "写入压缩文件失败").with_source(err),
);
}
if is_task_cancelled(&state, task_id).await? {
@@ -483,6 +529,7 @@ async fn process_task_file(
Ok(())
}
#[allow(clippy::too_many_arguments)]
async fn finalize_file(
state: &AppState,
billing_ctx: &Option<billing::BillingContext>,
@@ -563,7 +610,12 @@ async fn finalize_file(
Ok(())
}
async fn mark_file_failed(state: &AppState, task_id: Uuid, task_file_id: Uuid, message: &str) -> Result<(), AppError> {
async fn mark_file_failed(
state: &AppState,
task_id: Uuid,
task_file_id: Uuid,
message: &str,
) -> Result<(), AppError> {
let mut tx = state
.db
.begin()
@@ -608,16 +660,6 @@ async fn mark_file_failed(state: &AppState, task_id: Uuid, task_file_id: Uuid, m
Ok(())
}
async fn mark_task_failed(state: &AppState, task_id: Uuid, message: &str) -> Result<(), AppError> {
sqlx::query("UPDATE tasks SET status = 'failed', error_message = $2, completed_at = NOW() WHERE id = $1")
.bind(task_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",
@@ -627,7 +669,9 @@ async fn finalize_task_status(state: &AppState, task_id: Uuid) -> Result<(), App
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "查询任务失败").with_source(err))?;
let Some((total, completed, failed, status)) = row else { return Ok(()); };
let Some((total, completed, failed, status)) = row else {
return Ok(());
};
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')",
@@ -675,6 +719,7 @@ async fn finalize_task_status(state: &AppState, task_id: Uuid) -> Result<(), App
Ok(())
}
#[allow(clippy::too_many_arguments)]
async fn charge_one_unit(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
billing: &billing::BillingContext,
@@ -771,32 +816,36 @@ async fn cleanup_expired_records(state: &AppState) -> Result<(), AppError> {
.execute(&state.db)
.await;
let _ = sqlx::query("DELETE FROM email_verifications WHERE expires_at < NOW() AND verified_at IS NULL")
.execute(&state.db)
.await;
let _ = sqlx::query(
"DELETE FROM email_verifications WHERE expires_at < NOW() AND verified_at IS NULL",
)
.execute(&state.db)
.await;
let _ = sqlx::query("DELETE FROM password_resets WHERE expires_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)
.await;
let _ =
sqlx::query("DELETE FROM webhook_events WHERE received_at < NOW() - INTERVAL '90 days'")
.execute(&state.db)
.await;
Ok(())
}
async fn cleanup_expired_tasks(state: &AppState) -> Result<(), AppError> {
let task_ids: Vec<Uuid> = sqlx::query_scalar("SELECT id FROM tasks WHERE expires_at < NOW() LIMIT 200")
.fetch_all(&state.db)
.await
.unwrap_or_default();
let task_ids: Vec<Uuid> =
sqlx::query_scalar("SELECT id FROM tasks WHERE expires_at < NOW() LIMIT 200")
.fetch_all(&state.db)
.await
.unwrap_or_default();
if task_ids.is_empty() {
return Ok(());
}
if state.config.storage_type.to_ascii_lowercase() == "local" {
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")