fix: settle quota and bound compression work
Some checks failed
CI / verify (push) Has been cancelled

This commit is contained in:
237899745
2026-07-25 21:40:57 +08:00
parent 64b1169e8c
commit b67275b3e8
12 changed files with 1150 additions and 657 deletions

4
.cargo/audit.toml Normal file
View File

@@ -0,0 +1,4 @@
# SQLx locks optional MySQL dependencies even though this project builds only
# the PostgreSQL driver. The vulnerable RSA code is absent from `cargo tree`.
[advisories]
ignore = ["RUSTSEC-2023-0071"]

View File

@@ -30,6 +30,12 @@ jobs:
- name: Run Rust tests
run: cargo test --all-targets
- name: Install cargo-audit
run: cargo install cargo-audit --locked --version 0.22.2
- name: Audit Rust dependencies
run: cargo audit
- name: Install Node.js
uses: actions/setup-node@v4
with:

1037
Cargo.lock generated

File diff suppressed because it is too large Load Diff

View File

@@ -11,11 +11,11 @@ axum = { version = "0.8", features = ["multipart"] }
tokio = { version = "1", features = ["full"] }
tower-http = { version = "0.6", features = ["cors", "trace", "compression-full", "fs"] }
axum-extra = { version = "0.12", features = ["cookie"] }
time = "0.3"
time = "0.3.47"
serde = { version = "1", features = ["derive"] }
serde_json = "1"
bytes = "1"
bytes = "1.11.1"
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
@@ -25,7 +25,7 @@ chrono = { version = "0.4", features = ["serde"] }
dotenvy = "0.15"
sqlx = { version = "0.7", features = ["runtime-tokio-rustls", "postgres", "uuid", "chrono", "json"] }
sqlx = { version = "0.8.6", default-features = false, features = ["runtime-tokio-rustls", "postgres", "uuid", "chrono", "json", "macros", "migrate"] }
redis = { version = "0.24.1", features = ["tokio-comp", "connection-manager", "streams"] }
# Auth / security
@@ -35,12 +35,12 @@ hex = "0.4"
hmac = "0.12"
jsonwebtoken = "9"
percent-encoding = "2"
rand = "0.8"
rand = "0.8.7"
sha2 = "0.10"
aes-gcm = "0.10"
# S3-compatible object storage (Garage, MinIO, AWS S3, etc.)
aws-sdk-s3 = { version = "1", default-features = false, features = ["rustls", "default-https-client", "rt-tokio"] }
aws-sdk-s3 = { version = "1.137.0", default-features = false, features = ["default-https-client", "rt-tokio"] }
# Images
# Keep image-rs limited to formats exposed by the API. AVIF decoding uses
@@ -58,7 +58,7 @@ reqwest = { version = "0.12", default-features = false, features = ["rustls-tls"
url = "2"
# Mail
lettre = { version = "0.11", default-features = false, features = ["tokio1", "tokio1-rustls-tls", "builder", "smtp-transport"] }
lettre = { version = "0.11.22", default-features = false, features = ["tokio1", "tokio1-rustls-tls", "builder", "smtp-transport"] }
# ZIP download (batch)
tokio-util = { version = "0.7", features = ["io"] }

View File

@@ -248,7 +248,7 @@ webp = "0.3" # WebP 编解码
ravif = "0.11" # AVIF 编码
# 数据库
sqlx = { version = "0.7", features = ["runtime-tokio", "postgres", "uuid", "chrono"] }
sqlx = { version = "0.8.6", default-features = false, features = ["runtime-tokio-rustls", "postgres", "uuid", "chrono", "json", "macros", "migrate"] }
# Redis
redis = { version = "0.24", features = ["tokio-comp"] }

View File

@@ -0,0 +1,11 @@
ALTER TABLE tasks
ADD COLUMN IF NOT EXISTS anonymous_quota_date DATE;
UPDATE tasks
SET anonymous_quota_date = (created_at AT TIME ZONE 'Asia/Shanghai')::date
WHERE anonymous_units_reserved > 0
AND anonymous_quota_date IS NULL;
CREATE INDEX IF NOT EXISTS idx_tasks_anonymous_reservation_pending
ON tasks(id)
WHERE anonymous_units_reserved > 0;

View File

@@ -3,6 +3,7 @@ use crate::api::envelope::Envelope;
use crate::error::{AppError, ErrorCode};
use crate::services::billing;
use crate::services::mail;
use crate::services::quota;
use crate::services::settings;
use crate::services::settings::{
AuthConfigStored, MailConfigStored, MailCustomSmtp, StripeConfigStored,
@@ -543,6 +544,7 @@ async fn cancel_task(
};
if matches!(status.as_str(), "completed" | "failed" | "cancelled") {
quota::settle_anonymous_task_reservation(&state, task_id).await?;
return Ok(Json(Envelope {
success: true,
data: MessageResponse {
@@ -595,6 +597,8 @@ async fn cancel_task(
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "提交事务失败").with_source(err))?;
quota::settle_anonymous_task_reservation(&state, task_id).await?;
Ok(Json(Envelope {
success: true,
data: MessageResponse {

View File

@@ -22,6 +22,7 @@ use crate::state::AppState;
use axum::extract::DefaultBodyLimit;
use axum::Router;
use std::net::SocketAddr;
use tower_http::compression::CompressionLayer;
use tower_http::services::{ServeDir, ServeFile};
pub async fn run(state: AppState) -> Result<(), AppError> {
@@ -33,7 +34,9 @@ pub async fn run(state: AppState) -> Result<(), AppError> {
let static_service =
ServeDir::new("static").not_found_service(ServeFile::new("static/index.html"));
let v1 = v1_router().layer(DefaultBodyLimit::max(100 * 1024 * 1024));
let v1 = v1_router()
.layer(CompressionLayer::new())
.layer(DefaultBodyLimit::max(100 * 1024 * 1024));
let app = Router::new()
.route("/health", axum::routing::get(health::health))

View File

@@ -190,13 +190,15 @@ async fn create_batch_task(
}
let mut anonymous_reserved_units = 0u32;
let mut anonymous_quota_date = None;
let create_result: Result<BatchCreateResponse, AppError> = (async {
match &admission.task_owner {
TaskOwner::Anonymous { session_id } => {
let units = u32::try_from(files.len()).map_err(|_| {
AppError::new(ErrorCode::InvalidRequest, "批量文件数量超出限制")
})?;
quota::consume_anonymous_units(&state, session_id, ip, units).await?;
anonymous_quota_date =
Some(quota::reserve_anonymous_units(&state, session_id, ip, units).await?);
anonymous_reserved_units = units;
}
TaskOwner::User { .. } | TaskOwner::ApiKey { .. } => {
@@ -234,13 +236,13 @@ async fn create_batch_task(
compression_rate,
total_files, completed_files, failed_files,
total_original_size, total_compressed_size,
expires_at, retention_hours, anonymous_units_reserved
expires_at, retention_hours, anonymous_units_reserved, anonymous_quota_date
) VALUES (
$1, $2, $3, $4, $5::inet, $6::task_source, 'pending',
$7::compression_level, $8, $9, $10, $11, $12,
$13, 0, 0,
$14, 0,
$15, $16, $17
$15, $16, $17, $18
)
"#,
)
@@ -261,6 +263,7 @@ async fn create_batch_task(
.bind(expires_at)
.bind(retention_hours as i32)
.bind(anonymous_reserved_units as i32)
.bind(anonymous_quota_date)
.execute(&mut *tx)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "创建任务失败").with_source(err))?;
@@ -344,13 +347,30 @@ async fn create_batch_task(
Err(err) => {
if anonymous_reserved_units > 0 {
if let context::Principal::Anonymous { session_id } = &principal {
let _ = quota::refund_anonymous_units(
&state,
session_id,
ip,
anonymous_reserved_units,
)
.await;
let settlement =
quota::settle_anonymous_task_reservation(&state, task_id).await;
match settlement {
Ok(Some(_)) => {}
Ok(None) => {
if let Some(date) = anonymous_quota_date {
if let Err(refund_err) = quota::refund_anonymous_reservation_once(
&state,
task_id,
session_id,
ip,
date,
anonymous_reserved_units,
)
.await
{
tracing::warn!(task_id = %task_id, error = %refund_err, "failed to refund anonymous batch admission");
}
}
}
Err(settle_err) => {
tracing::warn!(task_id = %task_id, error = %settle_err, "failed to settle anonymous batch admission");
}
}
}
}
if let (Some(scope), Some(idem_key)) = (idempotency_scope, idempotency_key.as_deref()) {
@@ -950,6 +970,7 @@ async fn cancel_task(
)?;
if matches!(task.status.as_str(), "completed" | "failed" | "cancelled") {
quota::settle_anonymous_task_reservation(&state, task_id).await?;
return Ok((
jar,
Json(Envelope {
@@ -971,6 +992,8 @@ async fn cancel_task(
return Err(AppError::new(ErrorCode::InvalidRequest, "任务状态不可取消"));
}
quota::settle_anonymous_task_reservation(&state, task_id).await?;
Ok((
jar,
Json(Envelope {
@@ -1019,6 +1042,24 @@ 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,
"任务状态已变化,请刷新后重试",
));
}
}
quota::settle_anonymous_task_reservation(&state, task_id).await?;
let files = sqlx::query_as::<_, TaskStorageRow>(
r#"
SELECT storage_backend, storage_endpoint_id,

View File

@@ -16,18 +16,79 @@ use img_parts::{Bytes as ImgBytes, DynImage, ImageEXIF, ImageICC};
use oxipng::StripChunks;
use rgb::FromSlice;
use std::io::Cursor;
use std::time::Instant;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
const TARGET_MIN_LONG_EDGE: u32 = 640;
const TARGET_MIN_SCALE: f64 = 0.55;
const TARGET_SCALE_SEARCH_ATTEMPTS: usize = 6;
const TARGET_QUALITY_INTERPOLATION_ATTEMPTS: usize = 2;
const COMPRESSION_TIME_BUDGET: Duration = Duration::from_secs(30);
const COMPRESSION_TIMEOUT_GRACE: Duration = Duration::from_secs(1);
const JPEG_TARGET_MIN_QUALITY: u8 = 25;
const WEBP_TARGET_MIN_QUALITY: u8 = 30;
const AVIF_TARGET_MIN_QUALITY: u8 = 38;
const METADATA_TARGET_OVERHEAD: u64 = 1024;
#[derive(Clone)]
struct CompressionDeadline {
deadline: Option<Instant>,
cancelled: Arc<AtomicBool>,
}
impl CompressionDeadline {
fn new(budget: Duration) -> Self {
Self {
deadline: Instant::now().checked_add(budget),
cancelled: Arc::new(AtomicBool::new(false)),
}
}
#[cfg(test)]
fn unlimited() -> Self {
Self {
deadline: None,
cancelled: Arc::new(AtomicBool::new(false)),
}
}
fn cancel(&self) {
self.cancelled.store(true, Ordering::Relaxed);
}
fn expired(&self) -> bool {
self.cancelled.load(Ordering::Relaxed)
|| self
.deadline
.is_some_and(|deadline| Instant::now() >= deadline)
}
fn remaining(&self) -> Duration {
if self.cancelled.load(Ordering::Relaxed) {
return Duration::ZERO;
}
self.deadline
.map(|deadline| deadline.saturating_duration_since(Instant::now()))
.unwrap_or(Duration::MAX)
}
fn check(&self) -> Result<(), AppError> {
if self.expired() {
return Err(compression_timeout_error());
}
Ok(())
}
}
fn compression_timeout_error() -> AppError {
AppError::new(
ErrorCode::CompressionFailed,
"图片处理超过 30 秒,请降低分辨率或改用 JPEG/WebP 格式",
)
}
#[derive(Debug, Clone, Copy)]
pub enum CompressionLevel {
High,
@@ -294,16 +355,27 @@ pub async fn compress_image_bytes(
.await?
.file_limits
.max_image_pixels;
let permit = state
.image_processing_semaphore
.clone()
.acquire_owned()
.await
.map_err(|err| {
AppError::new(ErrorCode::Internal, "图片处理并发控制器已关闭").with_source(err)
})?;
let deadline = CompressionDeadline::new(COMPRESSION_TIME_BUDGET);
let permit = match tokio::time::timeout(
deadline.remaining(),
state.image_processing_semaphore.clone().acquire_owned(),
)
.await
{
Ok(Ok(permit)) => permit,
Ok(Err(err)) => {
return Err(
AppError::new(ErrorCode::Internal, "图片处理并发控制器已关闭").with_source(err),
)
}
Err(_) => {
crate::services::metrics::record_compression(state, started.elapsed(), bytes_in, None);
return Err(compression_timeout_error());
}
};
let result = match tokio::task::spawn_blocking(move || {
let blocking_deadline = deadline.clone();
let handle = tokio::task::spawn_blocking(move || {
let _permit = permit;
compress_image_bytes_sync(
input,
@@ -316,14 +388,25 @@ pub async fn compress_image_bytes(
max_height,
preserve_metadata,
max_image_pixels,
&blocking_deadline,
)
})
});
let result = match tokio::time::timeout(
deadline
.remaining()
.saturating_add(COMPRESSION_TIMEOUT_GRACE),
handle,
)
.await
{
Ok(result) => result,
Err(err) => Err(
Ok(Ok(result)) => result,
Ok(Err(err)) => Err(
AppError::new(ErrorCode::CompressionFailed, "图片处理任务异常退出").with_source(err),
),
Err(_) => {
deadline.cancel();
Err(compression_timeout_error())
}
};
crate::services::metrics::record_compression(
state,
@@ -346,7 +429,9 @@ fn compress_image_bytes_sync(
max_height: Option<u32>,
preserve_metadata: bool,
max_image_pixels: u64,
deadline: &CompressionDeadline,
) -> Result<Vec<u8>, AppError> {
deadline.check()?;
let original_size = input.len() as u64;
#[cfg(not(target_os = "linux"))]
@@ -358,6 +443,7 @@ fn compress_image_bytes_sync(
}
let orientation = inspect_image(&input, max_image_pixels)?;
deadline.check()?;
if is_animated_image(&input, format_in)? {
return Err(AppError::new(
@@ -365,6 +451,7 @@ fn compress_image_bytes_sync(
format!("暂不支持动画 {}", format_in.as_str().to_ascii_uppercase()),
));
}
deadline.check()?;
let retention_rate = effective_rate(compression_rate, level);
// 优先使用直接指定的目标大小,其次根据百分比计算
@@ -446,24 +533,26 @@ fn compress_image_bytes_sync(
})?
} else {
let mut image = decode_image(&input, format_in)?;
deadline.check()?;
image.apply_orientation(orientation);
let (image, did_resize) = resize_if_needed(image, max_width, max_height);
transformed |= did_resize;
deadline.check()?;
match format_out {
ImageFmt::Png => encode_png(image, strength_rate, preserve_metadata)?,
ImageFmt::Jpeg => match encoding_target_size {
Some(target) => encode_jpeg_target(image, target)?,
Some(target) => encode_jpeg_target(image, target, deadline)?,
None => encode_jpeg(image, strength_rate)?,
},
ImageFmt::Webp => match encoding_target_size {
Some(target) => encode_webp_target(image, target)?,
Some(target) => encode_webp_target(image, target, deadline)?,
None => encode_webp(image, strength_rate)?,
},
ImageFmt::Avif => match encoding_target_size {
Some(target) => encode_avif_target(image, target)?,
Some(target) => encode_avif_target(image, target, deadline)?,
None => encode_avif(image, strength_rate)?,
},
ImageFmt::Gif => encode_gif(image, strength_rate)?,
@@ -472,7 +561,6 @@ fn compress_image_bytes_sync(
ImageFmt::Ico => encode_ico(image)?,
}
};
if preserve_metadata {
output = apply_metadata(output, icc_profile, exif)?;
}
@@ -775,28 +863,61 @@ fn encode_avif_raw(raw: &[u8], w: u32, h: u32, quality: u8) -> Result<Vec<u8>, A
Ok(encoded.avif_file)
}
fn encode_jpeg_target(image: DynamicImage, target_size: u64) -> Result<Vec<u8>, AppError> {
encode_with_auto_resize(image, target_size, JPEG_TARGET_MIN_QUALITY, 95, |img, q| {
let rgb = jpeg_rgb(img);
let (w, h) = rgb.dimensions();
encode_jpeg_raw(rgb.as_raw(), w, h, q)
})
fn encode_jpeg_target(
image: DynamicImage,
target_size: u64,
deadline: &CompressionDeadline,
) -> Result<Vec<u8>, AppError> {
encode_with_auto_resize(
image,
target_size,
JPEG_TARGET_MIN_QUALITY,
95,
deadline,
|img, q| {
let rgb = jpeg_rgb(img);
let (w, h) = rgb.dimensions();
encode_jpeg_raw(rgb.as_raw(), w, h, q)
},
)
}
fn encode_webp_target(image: DynamicImage, target_size: u64) -> Result<Vec<u8>, AppError> {
encode_with_auto_resize(image, target_size, WEBP_TARGET_MIN_QUALITY, 95, |img, q| {
let rgba = img.to_rgba8();
let (w, h) = rgba.dimensions();
encode_webp_raw(rgba.as_raw(), w, h, q)
})
fn encode_webp_target(
image: DynamicImage,
target_size: u64,
deadline: &CompressionDeadline,
) -> Result<Vec<u8>, AppError> {
encode_with_auto_resize(
image,
target_size,
WEBP_TARGET_MIN_QUALITY,
95,
deadline,
|img, q| {
let rgba = img.to_rgba8();
let (w, h) = rgba.dimensions();
encode_webp_raw(rgba.as_raw(), w, h, q)
},
)
}
fn encode_avif_target(image: DynamicImage, target_size: u64) -> Result<Vec<u8>, AppError> {
encode_with_auto_resize(image, target_size, AVIF_TARGET_MIN_QUALITY, 95, |img, q| {
let rgba = img.to_rgba8();
let (w, h) = rgba.dimensions();
encode_avif_raw(rgba.as_raw(), w, h, q)
})
fn encode_avif_target(
image: DynamicImage,
target_size: u64,
deadline: &CompressionDeadline,
) -> Result<Vec<u8>, AppError> {
encode_with_auto_resize(
image,
target_size,
AVIF_TARGET_MIN_QUALITY,
95,
deadline,
|img, q| {
let rgba = img.to_rgba8();
let (w, h) = rgba.dimensions();
encode_avif_raw(rgba.as_raw(), w, h, q)
},
)
}
/// 目标体积压缩(质量优先 + 有边界的降尺寸)
@@ -811,18 +932,27 @@ fn encode_with_auto_resize<F>(
target_size: u64,
min_q: u8,
max_q: u8,
deadline: &CompressionDeadline,
mut encode_fn: F,
) -> Result<Vec<u8>, AppError>
where
F: FnMut(&DynamicImage, u8) -> Result<Vec<u8>, AppError>,
{
deadline.check()?;
let (orig_w, orig_h) = image.dimensions();
let long_edge = orig_w.max(orig_h);
let long_edge_floor = TARGET_MIN_LONG_EDGE.min(long_edge) as f64 / long_edge as f64;
let min_scale = TARGET_MIN_SCALE.max(long_edge_floor).min(1.0);
let (full_result, full_dimensions) =
encode_target_at_scale(&image, 1.0, min_q, max_q, target_size, &mut encode_fn)?;
let (full_result, full_dimensions) = encode_target_at_scale(
&image,
1.0,
min_q,
max_q,
target_size,
deadline,
&mut encode_fn,
)?;
if full_result.len() as u64 <= target_size {
return Ok(full_result);
}
@@ -834,8 +964,16 @@ where
));
}
let (min_result, min_dimensions) =
encode_target_at_scale(&image, min_scale, min_q, max_q, target_size, &mut encode_fn)?;
deadline.check()?;
let (min_result, min_dimensions) = encode_target_at_scale(
&image,
min_scale,
min_q,
max_q,
target_size,
deadline,
&mut encode_fn,
)?;
if min_result.len() as u64 > target_size {
return Err(target_unreachable_error(
target_size,
@@ -849,7 +987,10 @@ where
let mut over_scale = 1.0;
let mut over_dimensions = full_dimensions;
for _ in 0..TARGET_SCALE_SEARCH_ATTEMPTS {
for _ in 0..scale_search_attempts(orig_w, orig_h) {
if deadline.expired() {
return Ok(best_under);
}
let candidate_scale = (under_scale + over_scale) / 2.0;
let candidate_dimensions = dimensions_at_scale(orig_w, orig_h, candidate_scale);
@@ -868,6 +1009,7 @@ where
min_q,
max_q,
target_size,
deadline,
&mut encode_fn,
)?;
if candidate_result.len() as u64 <= target_size {
@@ -883,6 +1025,14 @@ where
Ok(best_under)
}
fn scale_search_attempts(width: u32, height: u32) -> usize {
match u64::from(width).saturating_mul(u64::from(height)) {
12_000_000.. => 3,
4_000_000.. => 4,
_ => TARGET_SCALE_SEARCH_ATTEMPTS,
}
}
fn dimensions_at_scale(orig_w: u32, orig_h: u32, scale: f64) -> (u32, u32) {
(
((orig_w as f64 * scale).round() as u32).clamp(1, orig_w),
@@ -896,11 +1046,13 @@ fn encode_target_at_scale<F>(
min_q: u8,
max_q: u8,
target_size: u64,
deadline: &CompressionDeadline,
encode_fn: &mut F,
) -> Result<(Vec<u8>, (u32, u32)), AppError>
where
F: FnMut(&DynamicImage, u8) -> Result<Vec<u8>, AppError>,
{
deadline.check()?;
let (orig_w, orig_h) = image.dimensions();
let dimensions = dimensions_at_scale(orig_w, orig_h, scale);
let resized = (dimensions != (orig_w, orig_h)).then(|| {
@@ -910,8 +1062,16 @@ where
image::imageops::FilterType::Lanczos3,
)
});
deadline.check()?;
let candidate = resized.as_ref().unwrap_or(image);
let result = encode_target_quality_with_image(candidate, min_q, max_q, target_size, encode_fn)?;
let result = encode_target_quality_with_image(
candidate,
min_q,
max_q,
target_size,
deadline,
encode_fn,
)?;
Ok((result, dimensions))
}
@@ -921,11 +1081,13 @@ fn encode_target_quality_with_image<F>(
min_q: u8,
max_q: u8,
target_size: u64,
deadline: &CompressionDeadline,
encode_fn: &mut F,
) -> Result<Vec<u8>, AppError>
where
F: FnMut(&DynamicImage, u8) -> Result<Vec<u8>, AppError>,
{
deadline.check()?;
// Start with the highest quality. If it already fits, no lower-quality
// encodes can improve the result.
let max_quality = encode_fn(image, max_q)?;
@@ -934,6 +1096,7 @@ where
return Ok(max_quality);
}
deadline.check()?;
let min_quality = encode_fn(image, min_q)?;
let min_size = min_quality.len() as u64;
if min_size > target_size {
@@ -950,6 +1113,9 @@ where
// jumps near the target first; the binary phase still proves the exact
// highest fitting integer quality, so output quality is not approximated.
for _ in 0..TARGET_QUALITY_INTERPOLATION_ATTEMPTS {
if deadline.expired() {
return Ok(best_under);
}
if under_q.saturating_add(1) >= over_q {
break;
}
@@ -979,6 +1145,9 @@ where
}
while under_q.saturating_add(1) < over_q {
if deadline.expired() {
return Ok(best_under);
}
let mid = under_q + (over_q - under_q) / 2;
let bytes = encode_fn(image, mid)?;
let size = bytes.len() as u64;
@@ -1372,6 +1541,7 @@ mod tests {
None,
false,
1_000_000,
&CompressionDeadline::unlimited(),
)
.unwrap();
assert_eq!(detect_format(&output).unwrap(), ImageFmt::Webp);
@@ -1406,10 +1576,17 @@ mod tests {
let calls = Cell::new(0);
let image = DynamicImage::new_rgb8(800, 600);
let result = encode_with_auto_resize(image, 100, 40, 95, |_image, quality| {
calls.set(calls.get() + 1);
Ok(vec![0; quality as usize])
})
let result = encode_with_auto_resize(
image,
100,
40,
95,
&CompressionDeadline::unlimited(),
|_image, quality| {
calls.set(calls.get() + 1);
Ok(vec![0; quality as usize])
},
)
.unwrap();
assert_eq!(result.len(), 95);
@@ -1422,12 +1599,18 @@ mod tests {
let calls = RefCell::new(Vec::new());
let image = DynamicImage::new_rgb8(32, 32);
let result =
encode_target_quality_with_image(&image, 25, 95, 50_000, &mut |_image, quality| {
let result = encode_target_quality_with_image(
&image,
25,
95,
50_000,
&CompressionDeadline::unlimited(),
&mut |_image, quality| {
calls.borrow_mut().push(quality);
Ok(vec![0; usize::from(quality) * 1_000])
})
.unwrap();
},
)
.unwrap();
assert_eq!(result.len(), 50_000);
assert_eq!(*calls.borrow(), vec![95, 25, 50, 51]);
@@ -1436,10 +1619,17 @@ mod tests {
#[test]
fn target_encoder_can_reduce_a_landscape_at_the_long_edge_floor() {
let image = DynamicImage::new_rgb8(960, 640);
let result = encode_with_auto_resize(image, 40_000, 40, 40, |image, _quality| {
let (width, height) = image.dimensions();
Ok(vec![0; (width as usize * height as usize) / 10])
})
let result = encode_with_auto_resize(
image,
40_000,
40,
40,
&CompressionDeadline::unlimited(),
|image, _quality| {
let (width, height) = image.dimensions();
Ok(vec![0; (width as usize * height as usize) / 10])
},
)
.unwrap();
assert!(result.len() <= 40_000);
@@ -1455,6 +1645,7 @@ mod tests {
1_000,
40,
40,
&CompressionDeadline::unlimited(),
|_image, _quality| {
calls.set(calls.get() + 1);
Ok(vec![0; 50_000])
@@ -1467,6 +1658,74 @@ mod tests {
assert_eq!(calls.get(), 2);
}
#[test]
fn target_encoder_rejects_work_after_the_deadline() {
use std::cell::Cell;
let calls = Cell::new(0);
let deadline = CompressionDeadline::unlimited();
deadline.cancel();
let error = encode_with_auto_resize(
DynamicImage::new_rgb8(960, 640),
40_000,
40,
40,
&deadline,
|_image, _quality| {
calls.set(calls.get() + 1);
Ok(vec![0; 1])
},
)
.unwrap_err();
assert_eq!(error.code, ErrorCode::CompressionFailed);
assert_eq!(calls.get(), 0);
}
#[test]
fn compression_deadline_reports_no_remaining_time_after_cancellation() {
let deadline = CompressionDeadline::unlimited();
assert_eq!(deadline.remaining(), Duration::MAX);
deadline.cancel();
assert_eq!(deadline.remaining(), Duration::ZERO);
}
#[test]
fn target_encoder_returns_the_best_candidate_when_search_is_cancelled() {
use std::cell::Cell;
let calls = Cell::new(0);
let deadline = CompressionDeadline::unlimited();
let controller = deadline.clone();
let result = encode_with_auto_resize(
DynamicImage::new_rgb8(960, 640),
40_000,
40,
40,
&deadline,
|image, _quality| {
calls.set(calls.get() + 1);
if calls.get() == 2 {
controller.cancel();
}
let (width, height) = image.dimensions();
Ok(vec![0; (width as usize * height as usize) / 10])
},
)
.unwrap();
assert!(result.len() <= 40_000);
assert_eq!(calls.get(), 2);
}
#[test]
fn large_images_use_a_smaller_scale_search_budget() {
assert_eq!(scale_search_attempts(4_000, 3_000), 3);
assert_eq!(scale_search_attempts(2_000, 2_000), 4);
assert_eq!(scale_search_attempts(1_999, 2_000), 6);
}
#[test]
fn png_fast_path_enforces_the_pixel_limit() {
let input = encode_png(DynamicImage::new_rgb8(20, 20), 100, false).unwrap();
@@ -1481,6 +1740,7 @@ mod tests {
None,
true,
399,
&CompressionDeadline::unlimited(),
)
.unwrap_err();
@@ -1506,6 +1766,7 @@ mod tests {
None,
false,
1_000_000,
&CompressionDeadline::unlimited(),
)
.unwrap();
let decoded = image::load_from_memory(&output).unwrap();
@@ -1564,6 +1825,7 @@ mod tests {
None,
false,
1_000_000,
&CompressionDeadline::unlimited(),
)
.unwrap();

View File

@@ -2,7 +2,7 @@ use crate::error::{AppError, ErrorCode};
use crate::services::billing::BillingContext;
use crate::state::AppState;
use chrono::{DateTime, Duration, Utc};
use chrono::{DateTime, Duration, NaiveDate, Utc};
use sqlx::{FromRow, Postgres, Transaction};
use std::net::IpAddr;
use uuid::Uuid;
@@ -249,13 +249,24 @@ pub async fn consume_anonymous_units(
ip: IpAddr,
units: u32,
) -> Result<(), AppError> {
reserve_anonymous_units(state, session_id, ip, units)
.await
.map(|_| ())
}
pub async fn reserve_anonymous_units(
state: &AppState,
session_id: &str,
ip: IpAddr,
units: u32,
) -> Result<NaiveDate, AppError> {
let date = utc8_date();
if units == 0 {
return Ok(());
return Ok(date);
}
let date = utc8_date();
let session_key = format!("anon_quota:{session_id}:{date}");
let ip_key = format!("anon_quota_ip:{ip}:{date}");
let session_key = anonymous_session_key(session_id, date);
let ip_key = anonymous_ip_key(ip, date);
let mut conn = state.redis.clone();
@@ -306,7 +317,7 @@ pub async fn consume_anonymous_units(
));
}
Ok(())
Ok(date)
}
pub async fn refund_anonymous_units(
@@ -314,14 +325,23 @@ pub async fn refund_anonymous_units(
session_id: &str,
ip: IpAddr,
units: u32,
) -> Result<(), AppError> {
refund_anonymous_units_for_date(state, session_id, ip, utc8_date(), units).await
}
async fn refund_anonymous_units_for_date(
state: &AppState,
session_id: &str,
ip: IpAddr,
date: NaiveDate,
units: u32,
) -> Result<(), AppError> {
if units == 0 {
return Ok(());
}
let date = utc8_date();
let session_key = format!("anon_quota:{session_id}:{date}");
let ip_key = format!("anon_quota_ip:{ip}:{date}");
let session_key = anonymous_session_key(session_id, date);
let ip_key = anonymous_ip_key(ip, date);
let mut conn = state.redis.clone();
let script = redis::Script::new(
r#"
@@ -349,9 +369,171 @@ pub async fn refund_anonymous_units(
Ok(())
}
fn utc8_date() -> String {
pub async fn refund_anonymous_reservation_once(
state: &AppState,
task_id: Uuid,
session_id: &str,
ip: IpAddr,
date: NaiveDate,
units: u32,
) -> Result<(), AppError> {
if units == 0 {
return Ok(());
}
let session_key = anonymous_session_key(session_id, date);
let ip_key = anonymous_ip_key(ip, date);
let marker_key = format!("anon_quota_refund:{task_id}");
let mut conn = state.redis.clone();
let script = redis::Script::new(
r#"
local dec = tonumber(ARGV[1])
local ttl = tonumber(ARGV[2])
if redis.call('EXISTS', KEYS[3]) == 1 then
return 0
end
local function refund(key)
local current = tonumber(redis.call('GET', key) or '0')
if current <= 0 then return 0 end
return redis.call('DECRBY', key, math.min(current, dec))
end
refund(KEYS[1])
refund(KEYS[2])
redis.call('SET', KEYS[3], '1', 'EX', ttl)
return 1
"#,
);
let _: i64 = script
.key(session_key)
.key(ip_key)
.key(marker_key)
.arg(units as i64)
.arg(48 * 60 * 60)
.invoke_async(&mut conn)
.await
.map_err(|err| {
AppError::new(ErrorCode::Internal, "退还匿名任务配额失败").with_source(err)
})?;
Ok(())
}
#[derive(Debug, FromRow)]
struct AnonymousTaskReservationRow {
anonymous_units_reserved: i32,
anonymous_quota_date: Option<NaiveDate>,
total_files: i32,
consumed_units: i32,
session_id: Option<String>,
client_ip: Option<String>,
created_at: DateTime<Utc>,
}
pub async fn settle_anonymous_task_reservation(
state: &AppState,
task_id: Uuid,
) -> Result<Option<u32>, AppError> {
let mut tx = state.db.begin().await.map_err(|err| {
AppError::new(ErrorCode::Internal, "开启匿名配额结算事务失败").with_source(err)
})?;
let row = sqlx::query_as::<_, AnonymousTaskReservationRow>(
r#"
SELECT anonymous_units_reserved,
anonymous_quota_date,
total_files,
COALESCE((
SELECT COUNT(*)::integer
FROM task_files f
WHERE f.task_id = tasks.id
AND f.status = 'completed'
AND f.compressed_size < f.original_size
AND NOT (
tasks.compression_rate = 100
AND f.original_format = f.output_format
AND tasks.max_width IS NULL
AND tasks.max_height IS NULL
)
), 0) AS consumed_units,
session_id,
host(client_ip) AS client_ip,
created_at
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(row) = row else {
tx.rollback().await.ok();
return Ok(None);
};
if row.anonymous_units_reserved <= 0 {
tx.rollback().await.ok();
return Ok(Some(0));
}
let refundable = refundable_reserved_units(
row.anonymous_units_reserved,
row.total_files,
row.consumed_units,
);
if refundable > 0 {
let session_id = row
.session_id
.as_deref()
.filter(|value| !value.is_empty())
.ok_or_else(|| AppError::new(ErrorCode::Internal, "匿名任务缺少 session_id"))?;
let ip = row
.client_ip
.as_deref()
.and_then(|value| value.parse::<IpAddr>().ok())
.ok_or_else(|| AppError::new(ErrorCode::Internal, "匿名任务缺少 client_ip"))?;
let date = row
.anonymous_quota_date
.unwrap_or_else(|| (row.created_at + Duration::hours(8)).date_naive());
refund_anonymous_reservation_once(state, task_id, session_id, ip, date, refundable).await?;
}
sqlx::query("UPDATE tasks SET anonymous_units_reserved = 0 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(Some(refundable))
}
fn refundable_reserved_units(reserved: i32, total_files: i32, consumed_units: i32) -> u32 {
let reserved = reserved.max(0);
let total_files = total_files.max(0);
let consumed_units = consumed_units.clamp(0, total_files);
let unused_reservation = reserved.saturating_sub(consumed_units);
let unfinished_files = total_files.saturating_sub(consumed_units);
u32::try_from(unused_reservation.min(unfinished_files)).unwrap_or(0)
}
fn anonymous_session_key(session_id: &str, date: NaiveDate) -> String {
format!("anon_quota:{session_id}:{}", date.format("%Y-%m-%d"))
}
fn anonymous_ip_key(ip: IpAddr, date: NaiveDate) -> String {
format!("anon_quota_ip:{ip}:{}", date.format("%Y-%m-%d"))
}
fn utc8_date() -> NaiveDate {
let now = Utc::now() + Duration::hours(8);
now.format("%Y-%m-%d").to_string()
now.date_naive()
}
#[cfg(test)]
@@ -391,4 +573,32 @@ mod tests {
now + Duration::days(20),
));
}
#[test]
fn anonymous_reservation_refunds_only_unconsumed_units() {
assert_eq!(refundable_reserved_units(10, 10, 6), 4);
assert_eq!(refundable_reserved_units(10, 10, 10), 0);
assert_eq!(refundable_reserved_units(10, 10, 0), 10);
}
#[test]
fn anonymous_reservation_refund_is_clamped_to_valid_bounds() {
assert_eq!(refundable_reserved_units(3, 10, 2), 1);
assert_eq!(refundable_reserved_units(10, 3, 1), 2);
assert_eq!(refundable_reserved_units(-1, 10, 0), 0);
assert_eq!(refundable_reserved_units(10, -1, -2), 0);
}
#[test]
fn anonymous_quota_keys_use_the_reserved_date() {
let date = NaiveDate::from_ymd_opt(2026, 7, 25).unwrap();
assert_eq!(
anonymous_session_key("session", date),
"anon_quota:session:2026-07-25"
);
assert_eq!(
anonymous_ip_key("127.0.0.1".parse().unwrap(), date),
"anon_quota_ip:127.0.0.1:2026-07-25"
);
}
}

View File

@@ -28,7 +28,7 @@ pub async fn run(state: AppState) -> Result<(), AppError> {
crate::services::bootstrap::ensure_schema(&state).await?;
let consumer = format!("worker_{}", Uuid::new_v4());
ensure_group(&state, &consumer).await?;
ensure_group(&state).await?;
let mut last_maintenance = Instant::now();
@@ -47,7 +47,7 @@ pub async fn run(state: AppState) -> Result<(), AppError> {
}
}
async fn ensure_group(state: &AppState, _consumer: &str) -> Result<(), AppError> {
async fn ensure_group(state: &AppState) -> Result<(), AppError> {
let mut conn = state.redis.clone();
let res: Result<redis::Value, redis::RedisError> = redis::cmd("XGROUP")
@@ -284,6 +284,7 @@ async fn mark_task_dead_letter(
tx.commit()
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "提交死信事务失败").with_source(err))?;
quota::settle_anonymous_task_reservation(state, task_id).await?;
Ok(())
}
@@ -384,7 +385,12 @@ async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
return Ok(());
};
if matches!(task.status.as_str(), "completed" | "failed" | "cancelled") {
if task.status == "cancelled" {
finalize_task_status(state, task_id).await?;
return Ok(());
}
if matches!(task.status.as_str(), "completed" | "failed") {
quota::settle_anonymous_task_reservation(state, task_id).await?;
return Ok(());
}
@@ -752,6 +758,18 @@ async fn finalize_file(
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "开启事务失败").with_source(err))?;
let task_status: 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, "锁定任务状态失败").with_source(err)
})?;
if !matches!(task_status.as_deref(), Some("pending" | "processing")) {
return Err(AppError::new(ErrorCode::InvalidRequest, "任务已结束"));
}
// Paid users: charge before marking file completed (atomic w/ status update).
if charge_units {
if let Some(billing) = billing_ctx {
@@ -923,11 +941,18 @@ async fn finalize_task_status(state: &AppState, task_id: Uuid) -> Result<(), App
.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",
r#"
UPDATE tasks
SET completed_files = (SELECT COUNT(*) FROM task_files WHERE task_id = $1 AND status = 'completed'),
failed_files = (SELECT COUNT(*) FROM task_files WHERE task_id = $1 AND status = 'failed'),
completed_at = COALESCE(completed_at, NOW())
WHERE id = $1
"#,
)
.bind(task_id)
.execute(&state.db)
.await;
quota::settle_anonymous_task_reservation(state, task_id).await?;
return Ok(());
}
@@ -945,6 +970,7 @@ async fn finalize_task_status(state: &AppState, task_id: Uuid) -> Result<(), App
.execute(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "更新任务状态失败").with_source(err))?;
quota::settle_anonymous_task_reservation(state, task_id).await?;
}
Ok(())
@@ -995,12 +1021,34 @@ async fn charge_one_unit(
}
async fn maintenance(state: &AppState) -> Result<(), AppError> {
settle_finished_anonymous_reservations(state).await?;
cleanup_expired_tasks(state).await?;
cleanup_stale_zip_temp(state).await?;
cleanup_expired_records(state).await?;
Ok(())
}
async fn settle_finished_anonymous_reservations(state: &AppState) -> Result<(), AppError> {
let task_ids: Vec<Uuid> = sqlx::query_scalar(
r#"
SELECT id
FROM tasks
WHERE anonymous_units_reserved > 0
AND status IN ('completed', 'failed', 'cancelled')
ORDER BY completed_at ASC NULLS FIRST
LIMIT 200
"#,
)
.fetch_all(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "查询待结算匿名任务失败").with_source(err))?;
for task_id in task_ids {
quota::settle_anonymous_task_reservation(state, task_id).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 {
@@ -1092,6 +1140,15 @@ 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,