Implement compression quota refunds and admin manual subscription
This commit is contained in:
738
src/worker/mod.rs
Normal file
738
src/worker/mod.rs
Normal file
@@ -0,0 +1,738 @@
|
||||
use crate::error::{AppError, ErrorCode};
|
||||
use crate::services::billing;
|
||||
use crate::services::compress;
|
||||
use crate::services::quota;
|
||||
use crate::state::AppState;
|
||||
|
||||
use redis::streams::StreamReadOptions;
|
||||
use redis::AsyncCommands;
|
||||
use sqlx::FromRow;
|
||||
use std::net::IpAddr;
|
||||
use std::time::Instant;
|
||||
use uuid::Uuid;
|
||||
|
||||
const STREAM_KEY: &str = "stream:compress_jobs";
|
||||
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, "数据库结构初始化失败");
|
||||
}
|
||||
|
||||
let consumer = format!("worker_{}", Uuid::new_v4());
|
||||
ensure_group(&state, &consumer).await?;
|
||||
|
||||
let mut last_maintenance = Instant::now();
|
||||
|
||||
loop {
|
||||
if let Err(err) = poll_once(&state, &consumer).await {
|
||||
tracing::error!(error = ?err, "worker poll error");
|
||||
tokio::time::sleep(std::time::Duration::from_secs(2)).await;
|
||||
}
|
||||
|
||||
if last_maintenance.elapsed().as_secs() >= 300 {
|
||||
if let Err(err) = maintenance(&state).await {
|
||||
tracing::error!(error = ?err, "maintenance failed");
|
||||
}
|
||||
last_maintenance = Instant::now();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn ensure_group(state: &AppState, _consumer: &str) -> Result<(), AppError> {
|
||||
let mut conn = state.redis.clone();
|
||||
|
||||
let res: Result<redis::Value, redis::RedisError> = redis::cmd("XGROUP")
|
||||
.arg("CREATE")
|
||||
.arg(STREAM_KEY)
|
||||
.arg(GROUP_NAME)
|
||||
.arg("0")
|
||||
.arg("MKSTREAM")
|
||||
.query_async(&mut conn)
|
||||
.await;
|
||||
|
||||
match res {
|
||||
Ok(_) => Ok(()),
|
||||
Err(err) => {
|
||||
let msg = err.to_string();
|
||||
if msg.contains("BUSYGROUP") {
|
||||
return Ok(());
|
||||
}
|
||||
Err(AppError::new(ErrorCode::Internal, "初始化队列失败").with_source(err))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn poll_once(state: &AppState, consumer: &str) -> Result<(), AppError> {
|
||||
let mut conn = state.redis.clone();
|
||||
|
||||
let opts = StreamReadOptions::default()
|
||||
.group(GROUP_NAME, consumer)
|
||||
.count(1)
|
||||
.block(5000);
|
||||
|
||||
let reply: redis::streams::StreamReadReply = conn
|
||||
.xread_options(&[STREAM_KEY], &[">"], &opts)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "读取队列失败").with_source(err))?;
|
||||
|
||||
if reply.keys.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
for key in reply.keys {
|
||||
for msg in key.ids {
|
||||
let Some(task_id_str) = msg.get::<String>("task_id") else {
|
||||
ack_message(&mut conn, &msg.id).await?;
|
||||
continue;
|
||||
};
|
||||
|
||||
let task_id = match Uuid::parse_str(&task_id_str) {
|
||||
Ok(v) => v,
|
||||
Err(_) => {
|
||||
ack_message(&mut conn, &msg.id).await?;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
if let Err(err) = process_task(state, task_id).await {
|
||||
tracing::error!(task_id = %task_id, error = %err, "task processing failed");
|
||||
}
|
||||
|
||||
ack_message(&mut conn, &msg.id).await?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
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)
|
||||
.arg(msg_id)
|
||||
.query_async(conn)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "确认队列消息失败").with_source(err))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[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>,
|
||||
source: String,
|
||||
client_ip: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct TaskFileProcRow {
|
||||
id: Uuid,
|
||||
storage_path: Option<String>,
|
||||
original_name: String,
|
||||
original_format: String,
|
||||
output_format: String,
|
||||
original_size: i64,
|
||||
status: String,
|
||||
}
|
||||
|
||||
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,
|
||||
source::text AS source,
|
||||
client_ip::text AS client_ip
|
||||
FROM tasks
|
||||
WHERE id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(task_id)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询任务失败").with_source(err))?
|
||||
.ok_or_else(|| AppError::new(ErrorCode::NotFound, "任务不存在"))?;
|
||||
|
||||
if matches!(task.status.as_str(), "completed" | "failed" | "cancelled") {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let updated = sqlx::query(
|
||||
r#"
|
||||
UPDATE tasks
|
||||
SET status = 'processing', started_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 && task.status == "pending" {
|
||||
// Another worker might have taken it.
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// 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());
|
||||
let level = compression_rate
|
||||
.map(compress::rate_to_level)
|
||||
.unwrap_or(compress::parse_level(&task.compression_level)?);
|
||||
let max_width = task.max_width.and_then(|v| u32::try_from(v).ok());
|
||||
let max_height = task.max_height.and_then(|v| u32::try_from(v).ok());
|
||||
|
||||
let mut files: Vec<TaskFileProcRow> = sqlx::query_as(
|
||||
r#"
|
||||
SELECT
|
||||
id,
|
||||
storage_path,
|
||||
original_name,
|
||||
original_format,
|
||||
output_format,
|
||||
original_size,
|
||||
status::text AS status
|
||||
FROM task_files
|
||||
WHERE task_id = $1
|
||||
ORDER BY created_at ASC
|
||||
"#,
|
||||
)
|
||||
.bind(task_id)
|
||||
.fetch_all(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询任务文件失败").with_source(err))?;
|
||||
|
||||
let billing_ctx = if let Some(user_id) = task.user_id {
|
||||
Some(billing::get_user_billing(state, user_id).await?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let anon_ip: Option<IpAddr> = task
|
||||
.client_ip
|
||||
.as_deref()
|
||||
.and_then(|s| s.parse::<IpAddr>().ok());
|
||||
|
||||
for file in &mut files {
|
||||
// Stop early if cancelled.
|
||||
let status: Option<String> = sqlx::query_scalar("SELECT status::text FROM tasks WHERE id = $1")
|
||||
.bind(task_id)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.unwrap_or(None);
|
||||
if matches!(status.as_deref(), Some("cancelled")) {
|
||||
break;
|
||||
}
|
||||
|
||||
if file.status != "pending" {
|
||||
continue;
|
||||
}
|
||||
|
||||
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());
|
||||
if updated.rows_affected() == 0 {
|
||||
continue;
|
||||
}
|
||||
|
||||
let Some(input_path) = file.storage_path.clone() else {
|
||||
mark_file_failed(state, task_id, file.id, "原文件不存在").await?;
|
||||
continue;
|
||||
};
|
||||
|
||||
let input_bytes = match tokio::fs::read(&input_path).await {
|
||||
Ok(v) => v,
|
||||
Err(_) => {
|
||||
mark_file_failed(state, task_id, file.id, "读取原文件失败").await?;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let format_in = parse_image_fmt(&file.original_format)?;
|
||||
let format_out = parse_image_fmt(&file.output_format)?;
|
||||
|
||||
let compressed = match compress::compress_image_bytes(
|
||||
state,
|
||||
&input_bytes,
|
||||
format_in,
|
||||
format_out,
|
||||
level,
|
||||
compression_rate,
|
||||
max_width,
|
||||
max_height,
|
||||
task.preserve_metadata,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(v) => v,
|
||||
Err(err) => {
|
||||
mark_file_failed(state, task_id, file.id, &err.message).await?;
|
||||
let _ = tokio::fs::remove_file(&input_path).await;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let original_size = input_bytes.len() as u64;
|
||||
let compressed_size = compressed.len() as u64;
|
||||
let saved_percent = if original_size == 0 {
|
||||
0.0
|
||||
} else {
|
||||
(original_size.saturating_sub(compressed_size) as f64) * 100.0 / (original_size as f64)
|
||||
};
|
||||
let charge_units = compressed_size < original_size;
|
||||
|
||||
// Anonymous quota enforcement requires session_id + client_ip.
|
||||
if task.user_id.is_none() && charge_units {
|
||||
let Some(session_id) = task.session_id.as_deref() else {
|
||||
mark_file_failed(state, task_id, file.id, "匿名任务缺少 session_id").await?;
|
||||
let _ = tokio::fs::remove_file(&input_path).await;
|
||||
continue;
|
||||
};
|
||||
let Some(ip) = anon_ip else {
|
||||
mark_file_failed(state, task_id, file.id, "匿名任务缺少 client_ip").await?;
|
||||
let _ = tokio::fs::remove_file(&input_path).await;
|
||||
continue;
|
||||
};
|
||||
if let Err(err) = quota::consume_anonymous_units(state, session_id, ip, 1).await {
|
||||
mark_file_failed(state, task_id, file.id, &err.message).await?;
|
||||
let _ = tokio::fs::remove_file(&input_path).await;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
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 let Err(err) = finalize_file(
|
||||
state,
|
||||
&billing_ctx,
|
||||
task.api_key_id,
|
||||
&task.source,
|
||||
task_id,
|
||||
file.id,
|
||||
&output_path,
|
||||
original_size as i64,
|
||||
compressed_size as i64,
|
||||
saved_percent,
|
||||
format_in,
|
||||
format_out,
|
||||
charge_units,
|
||||
)
|
||||
.await
|
||||
{
|
||||
// If quota exceeded for paid users, don't leave output behind.
|
||||
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 _ = tokio::fs::remove_file(&input_path).await;
|
||||
continue;
|
||||
}
|
||||
|
||||
// Success: remove original.
|
||||
let _ = tokio::fs::remove_file(&input_path).await;
|
||||
}
|
||||
|
||||
finalize_task_status(state, task_id).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn parse_image_fmt(value: &str) -> Result<compress::ImageFmt, AppError> {
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"png" => Ok(compress::ImageFmt::Png),
|
||||
"jpeg" | "jpg" => Ok(compress::ImageFmt::Jpeg),
|
||||
"webp" => Ok(compress::ImageFmt::Webp),
|
||||
"avif" => Ok(compress::ImageFmt::Avif),
|
||||
"gif" => Ok(compress::ImageFmt::Gif),
|
||||
"bmp" => Ok(compress::ImageFmt::Bmp),
|
||||
"tif" | "tiff" => Ok(compress::ImageFmt::Tiff),
|
||||
"ico" => Ok(compress::ImageFmt::Ico),
|
||||
_ => Err(AppError::new(ErrorCode::InvalidRequest, "未知图片格式")),
|
||||
}
|
||||
}
|
||||
|
||||
async fn finalize_file(
|
||||
state: &AppState,
|
||||
billing_ctx: &Option<billing::BillingContext>,
|
||||
api_key_id: Option<Uuid>,
|
||||
source: &str,
|
||||
task_id: Uuid,
|
||||
task_file_id: Uuid,
|
||||
output_path: &str,
|
||||
bytes_in: i64,
|
||||
bytes_out: i64,
|
||||
saved_percent: f64,
|
||||
format_in: compress::ImageFmt,
|
||||
format_out: compress::ImageFmt,
|
||||
charge_units: bool,
|
||||
) -> Result<(), AppError> {
|
||||
let mut tx = state
|
||||
.db
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "开启事务失败").with_source(err))?;
|
||||
|
||||
// Paid users: charge before marking file completed (atomic w/ status update).
|
||||
if charge_units {
|
||||
if let Some(billing) = billing_ctx {
|
||||
charge_one_unit(
|
||||
&mut tx,
|
||||
billing,
|
||||
api_key_id,
|
||||
source,
|
||||
task_id,
|
||||
task_file_id,
|
||||
format_in,
|
||||
format_out,
|
||||
bytes_in as u64,
|
||||
bytes_out as u64,
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE task_files
|
||||
SET storage_path = $2,
|
||||
compressed_size = $3,
|
||||
saved_percent = $4,
|
||||
status = 'completed',
|
||||
completed_at = NOW()
|
||||
WHERE id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(task_file_id)
|
||||
.bind(output_path)
|
||||
.bind(bytes_out)
|
||||
.bind(saved_percent)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "更新文件失败").with_source(err))?;
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE tasks
|
||||
SET completed_files = completed_files + 1,
|
||||
total_compressed_size = total_compressed_size + $2
|
||||
WHERE id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(task_id)
|
||||
.bind(bytes_out)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "更新任务统计失败").with_source(err))?;
|
||||
|
||||
tx.commit()
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "提交事务失败").with_source(err))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn mark_file_failed(state: &AppState, task_id: Uuid, task_file_id: Uuid, message: &str) -> Result<(), AppError> {
|
||||
let mut tx = state
|
||||
.db
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "开启事务失败").with_source(err))?;
|
||||
|
||||
let updated = sqlx::query(
|
||||
r#"
|
||||
UPDATE task_files
|
||||
SET status = 'failed',
|
||||
error_message = $2,
|
||||
storage_path = NULL,
|
||||
completed_at = NOW()
|
||||
WHERE id = $1
|
||||
AND status NOT IN ('completed', 'failed')
|
||||
"#,
|
||||
)
|
||||
.bind(task_file_id)
|
||||
.bind(message)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "更新文件失败").with_source(err))?;
|
||||
|
||||
if updated.rows_affected() > 0 {
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE tasks
|
||||
SET failed_files = failed_files + 1
|
||||
WHERE id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(task_id)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "更新任务统计失败").with_source(err))?;
|
||||
}
|
||||
|
||||
tx.commit()
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "提交事务失败").with_source(err))?;
|
||||
|
||||
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",
|
||||
)
|
||||
.bind(task_id)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询任务失败").with_source(err))?;
|
||||
|
||||
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')",
|
||||
)
|
||||
.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 _ = 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')",
|
||||
)
|
||||
.bind(task_id)
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
|
||||
let _ = sqlx::query(
|
||||
"UPDATE tasks SET failed_files = GREATEST(total_files - completed_files, 0), completed_at = NOW() WHERE id = $1 AND completed_at IS NULL",
|
||||
)
|
||||
.bind(task_id)
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if completed + failed >= total && total > 0 {
|
||||
let final_status = if completed == 0 && failed == total {
|
||||
"failed"
|
||||
} else {
|
||||
"completed"
|
||||
};
|
||||
sqlx::query(
|
||||
"UPDATE tasks SET status = $2::task_status, completed_at = NOW() WHERE id = $1",
|
||||
)
|
||||
.bind(task_id)
|
||||
.bind(final_status)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "更新任务状态失败").with_source(err))?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn charge_one_unit(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
billing: &billing::BillingContext,
|
||||
api_key_id: Option<Uuid>,
|
||||
source: &str,
|
||||
task_id: Uuid,
|
||||
task_file_id: Uuid,
|
||||
format_in: compress::ImageFmt,
|
||||
format_out: compress::ImageFmt,
|
||||
bytes_in: u64,
|
||||
bytes_out: u64,
|
||||
) -> Result<(), AppError> {
|
||||
// Ensure usage period row exists.
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO usage_periods (user_id, subscription_id, period_start, period_end)
|
||||
VALUES ($1, $2, $3, $4)
|
||||
ON CONFLICT (user_id, period_start, period_end) DO NOTHING
|
||||
"#,
|
||||
)
|
||||
.bind(billing.user_id)
|
||||
.bind(billing.subscription_id)
|
||||
.bind(billing.period_start)
|
||||
.bind(billing.period_end)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "初始化用量周期失败").with_source(err))?;
|
||||
|
||||
let updated: Option<i32> = sqlx::query_scalar(
|
||||
r#"
|
||||
UPDATE usage_periods
|
||||
SET used_units = used_units + 1,
|
||||
bytes_in = bytes_in + $1,
|
||||
bytes_out = bytes_out + $2,
|
||||
updated_at = NOW()
|
||||
WHERE user_id = $3
|
||||
AND period_start = $4
|
||||
AND period_end = $5
|
||||
AND used_units + 1 <= $6 + bonus_units
|
||||
RETURNING used_units
|
||||
"#,
|
||||
)
|
||||
.bind(bytes_in as i64)
|
||||
.bind(bytes_out as i64)
|
||||
.bind(billing.user_id)
|
||||
.bind(billing.period_start)
|
||||
.bind(billing.period_end)
|
||||
.bind(billing.plan.included_units_per_period)
|
||||
.fetch_optional(&mut **tx)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "扣减配额失败").with_source(err))?;
|
||||
|
||||
if updated.is_none() {
|
||||
return Err(AppError::new(ErrorCode::QuotaExceeded, "当期配额已用完"));
|
||||
}
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO usage_events (
|
||||
user_id, api_key_id, source,
|
||||
task_id, task_file_id,
|
||||
units, bytes_in, bytes_out, format_in, format_out
|
||||
) VALUES (
|
||||
$1, $2, $3::task_source,
|
||||
$4, $5,
|
||||
1, $6, $7, $8, $9
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.bind(billing.user_id)
|
||||
.bind(api_key_id)
|
||||
.bind(source)
|
||||
.bind(task_id)
|
||||
.bind(task_file_id)
|
||||
.bind(bytes_in as i64)
|
||||
.bind(bytes_out as i64)
|
||||
.bind(format_in.as_str())
|
||||
.bind(format_out.as_str())
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "写入用量明细失败").with_source(err))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn maintenance(state: &AppState) -> Result<(), AppError> {
|
||||
cleanup_expired_tasks(state).await?;
|
||||
cleanup_expired_records(state).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)
|
||||
.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;
|
||||
|
||||
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();
|
||||
|
||||
if task_ids.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if state.config.storage_type.to_ascii_lowercase() == "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;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user