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

fix: settle anonymous single-file reservations

Browse Source
This commit is contained in:
237899745
2026-07-26 05:47:54 +08:00
parent 65694cee15
commit 910e60ab59
4 changed files with 563 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

87
src/api/compress.rs
View File

@@ -79,6 +79,10 @@ fn default_units_charged() -> i32 {
1
}
fn metered_units(charged: bool) -> i32 {
i32::from(charged)
}
fn direct_response<B: IntoResponse>(
body: B,
format: ImageFmt,
@@ -264,14 +268,16 @@ async fn compress_json(
}
}
let mut anonymous_reserved = false;
let task_id = Uuid::new_v4();
let mut anonymous_reservation_date = None;
let op: Result<CompressResponse, AppError> = (async {
match &quota_ctx {
QuotaContext::User(billing) => ensure_quota_available(&state, billing, 1).await?,
QuotaContext::ApiKey(billing, _) => ensure_quota_available(&state, billing, 1).await?,
QuotaContext::Anonymous { session_id, ip } => {
quota::consume_anonymous_units(&state, session_id, *ip, 1).await?;
anonymous_reserved = true;
anonymous_reservation_date = Some(
quota::reserve_anonymous_single_unit(&state, task_id, session_id, *ip).await?,
);
}
}
@@ -299,17 +305,15 @@ async fn compress_json(
} else {
(saved_bytes as f64) * 100.0 / (original_size as f64)
};
let charge_units = anonymous_reserved
|| quota::output_consumes_unit(
req.compression_rate,
format_in == format_out,
req.max_width.is_some() || req.max_height.is_some(),
req.target_size_bytes.is_some(),
original_size,
compressed_size,
);
let charge_units = quota::output_consumes_unit(
req.compression_rate,
format_in == format_out,
req.max_width.is_some() || req.max_height.is_some(),
req.target_size_bytes.is_some(),
original_size,
compressed_size,
);
let task_id = Uuid::new_v4();
let file_id = Uuid::new_v4();
let retention_hours = retention.num_hours();
let object_key =
@@ -356,6 +360,17 @@ async fn compress_json(
return Err(err);
}
if anonymous_reservation_date.is_some() {
if let Err(err) =
quota::finalize_anonymous_single_reservation(&state, task_id, charge_units).await
{
// The durable reservation remains visible to maintenance, so a
// transient Redis failure must not turn a successful image into
// a failed, non-idempotent request.
tracing::warn!(task_id = %task_id, charged = charge_units, error = %err, "anonymous single reservation finalization deferred");
}
}
Ok(CompressResponse {
task_id,
file_id,
@@ -368,7 +383,7 @@ async fn compress_json(
download_url: format!("/downloads/{file_id}"),
expires_at,
billing: BillingView {
units_charged: if charge_units { 1 } else { 0 },
units_charged: metered_units(charge_units),
},
})
})
@@ -402,9 +417,11 @@ async fn compress_json(
))
}
Err(err) => {
if anonymous_reserved {
if let QuotaContext::Anonymous { session_id, ip } = &quota_ctx {
let _ = quota::refund_anonymous_units(&state, session_id, *ip, 1).await;
if anonymous_reservation_date.is_some() {
if let Err(refund_err) =
quota::refund_anonymous_single_reservation(&state, task_id).await
{
tracing::warn!(task_id = %task_id, error = %refund_err, "anonymous single reservation refund deferred");
}
}
if let (Some(scope), Some(idem_key), Some(request_hash)) = (
@@ -653,7 +670,7 @@ async fn compress_direct(
compressed_size,
saved_bytes,
saved_percent,
units_charged: if charge_units { 1 } else { 0 },
units_charged: metered_units(charge_units),
};
let response = direct_response(compressed, format_out, &idem_data);
Ok((response, idem_data))
@@ -1137,7 +1154,9 @@ async fn record_task_and_metering(
.map_err(|err| AppError::new(ErrorCode::Internal, "创建文件记录失败").with_source(err))?;
match quota_ctx {
QuotaContext::Anonymous { .. } => {}
QuotaContext::Anonymous { .. } => {
quota::mark_anonymous_single_result(&mut tx, task_id, charge_units).await?;
}
QuotaContext::User(billing) => {
if charge_units {
charge_one_unit(
@@ -1247,4 +1266,34 @@ mod tests {
assert_eq!(response.headers()["imageforge-saved-percent"], "37.50");
assert_eq!(response.headers()["imageforge-units-charged"], "1");
}
#[test]
fn anonymous_response_units_follow_actual_output_metering() {
let cases = [
(Some(100), true, false, false, 100, 50, 0),
(None, true, false, false, 100, 100, 0),
(None, true, false, false, 100, 50, 1),
(Some(100), true, true, false, 100, 50, 1),
];
for (
compression_rate,
same_format,
has_resize,
has_target_size,
original_size,
output_size,
expected_units,
) in cases
{
let charged = quota::output_consumes_unit(
compression_rate,
same_format,
has_resize,
has_target_size,
original_size,
output_size,
);
assert_eq!(metered_units(charged), expected_units);
}
}
}

470
src/services/quota.rs
View File

@@ -320,55 +320,257 @@ pub async fn reserve_anonymous_units(
Ok(date)
}
pub async fn refund_anonymous_units(
pub async fn reserve_anonymous_single_unit(
state: &AppState,
task_id: Uuid,
session_id: &str,
ip: IpAddr,
units: u32,
) -> Result<(), AppError> {
refund_anonymous_units_for_date(state, session_id, ip, utc8_date(), units).await
) -> Result<NaiveDate, AppError> {
reserve_anonymous_single_unit_for_date(state, task_id, session_id, ip, utc8_date()).await
}
async fn refund_anonymous_units_for_date(
async fn reserve_anonymous_single_unit_for_date(
state: &AppState,
task_id: Uuid,
session_id: &str,
ip: IpAddr,
date: NaiveDate,
units: u32,
) -> Result<(), AppError> {
if units == 0 {
return Ok(());
}
) -> Result<NaiveDate, AppError> {
sqlx::query(
r#"
INSERT INTO anonymous_single_reservations (
task_id, session_id, client_ip, quota_date, units
) VALUES ($1, $2, $3::inet, $4, 1)
"#,
)
.bind(task_id)
.bind(session_id)
.bind(ip.to_string())
.bind(date)
.execute(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "创建匿名单文件预留失败").with_source(err))?;
let limit = crate::services::settings::runtime_policy(state)
.await?
.rate_limits
.anonymous_units_per_day as i64;
let session_key = anonymous_session_key(session_id, date);
let ip_key = anonymous_ip_key(ip, date);
let reservation_key = anonymous_single_reservation_key(task_id);
let mut conn = state.redis.clone();
let script = redis::Script::new(
r#"
local limit = tonumber(ARGV[1])
local ttl = tonumber(ARGV[2])
if redis.call('EXISTS', KEYS[3]) == 1 then
return tonumber(redis.call('GET', KEYS[1]) or '0')
end
local session_value = tonumber(redis.call('GET', KEYS[1]) or '0')
local ip_value = tonumber(redis.call('GET', KEYS[2]) or '0')
if session_value + 1 > limit or ip_value + 1 > limit then
return -1
end
session_value = redis.call('INCRBY', KEYS[1], 1)
ip_value = redis.call('INCRBY', KEYS[2], 1)
if session_value == 1 then redis.call('EXPIRE', KEYS[1], ttl) end
if ip_value == 1 then redis.call('EXPIRE', KEYS[2], ttl) end
redis.call('SET', KEYS[3], '1', 'EX', ttl)
return session_value
"#,
);
let reserved: i64 = script
.key(session_key)
.key(ip_key)
.key(reservation_key)
.arg(limit)
.arg(48 * 60 * 60)
.invoke_async(&mut conn)
.await
.map_err(|err| {
// Keep the durable pending row: the script may have committed even
// if the response was lost, and maintenance can safely reconcile it.
AppError::new(ErrorCode::Internal, "匿名配额检查失败").with_source(err)
})?;
if reserved < 0 {
sqlx::query(
r#"
UPDATE anonymous_single_reservations
SET status = 'refunded', settled_at = NOW(), updated_at = NOW()
WHERE task_id = $1 AND status = 'pending'
"#,
)
.bind(task_id)
.execute(&state.db)
.await
.map_err(|err| {
AppError::new(ErrorCode::Internal, "关闭匿名单文件预留失败").with_source(err)
})?;
return Err(AppError::new(
ErrorCode::QuotaExceeded,
format!("匿名试用次数已用完（每日 {limit} 次）"),
));
}
Ok(date)
}
pub async fn mark_anonymous_single_result(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
task_id: Uuid,
charged: bool,
) -> Result<(), AppError> {
let status = if charged { "charged" } else { "refund_pending" };
let updated = sqlx::query(
r#"
UPDATE anonymous_single_reservations
SET status = $2,
settled_at = CASE WHEN $2 = 'charged' THEN NOW() ELSE settled_at END,
updated_at = NOW()
WHERE task_id = $1 AND status = 'pending'
"#,
)
.bind(task_id)
.bind(status)
.execute(&mut **tx)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "更新匿名单文件预留失败").with_source(err))?;
if updated.rows_affected() != 1 {
return Err(AppError::new(
ErrorCode::StorageUnavailable,
"匿名单文件预留已失效，请重试",
));
}
Ok(())
}
pub async fn finalize_anonymous_single_reservation(
state: &AppState,
task_id: Uuid,
charged: bool,
) -> Result<(), AppError> {
if charged {
let mut conn = state.redis.clone();
let _: i64 = redis::cmd("DEL")
.arg(anonymous_single_reservation_key(task_id))
.query_async(&mut conn)
.await
.map_err(|err| {
AppError::new(ErrorCode::Internal, "完成匿名单文件计费失败").with_source(err)
})?;
return Ok(());
}
refund_anonymous_single_reservation(state, task_id).await
}
pub async fn refund_anonymous_single_reservation(
state: &AppState,
task_id: Uuid,
) -> Result<(), AppError> {
let row: Option<(String, String, NaiveDate, i32)> = sqlx::query_as(
r#"
SELECT session_id, host(client_ip), quota_date, units
FROM anonymous_single_reservations
WHERE task_id = $1
AND status IN ('pending', 'refund_pending', 'refunded')
"#,
)
.bind(task_id)
.fetch_optional(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "查询匿名单文件预留失败").with_source(err))?;
let Some((session_id, ip, date, units)) = row else {
return Ok(());
};
let ip: IpAddr = ip.parse().map_err(|err| {
AppError::new(ErrorCode::Internal, "匿名单文件预留 IP 无效").with_source(err)
})?;
let units = u32::try_from(units).unwrap_or(0);
let session_key = anonymous_session_key(&session_id, date);
let ip_key = anonymous_ip_key(ip, date);
let reservation_key = anonymous_single_reservation_key(task_id);
let refund_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[4]) == 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))
if redis.call('EXISTS', KEYS[3]) == 1 then
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('DEL', KEYS[3])
end
refund(KEYS[1])
refund(KEYS[2])
redis.call('SET', KEYS[4], '1', 'EX', ttl)
return 1
"#,
);
let _: i64 = script
.key(session_key)
.key(ip_key)
.key(reservation_key)
.key(refund_key)
.arg(units as i64)
.arg(48 * 60 * 60)
.invoke_async(&mut conn)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "退还匿名配额失败").with_source(err))?;
.map_err(|err| {
AppError::new(ErrorCode::Internal, "退还匿名单文件配额失败").with_source(err)
})?;
sqlx::query(
r#"
UPDATE anonymous_single_reservations
SET status = 'refunded', settled_at = NOW(), updated_at = NOW()
WHERE task_id = $1 AND status IN ('pending', 'refund_pending', 'refunded')
"#,
)
.bind(task_id)
.execute(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "记录匿名单文件退款失败").with_source(err))?;
Ok(())
}
pub async fn settle_stale_anonymous_single_reservations(
state: &AppState,
limit: i64,
) -> Result<usize, AppError> {
let task_ids: Vec<Uuid> = sqlx::query_scalar(
r#"
SELECT task_id
FROM anonymous_single_reservations
WHERE status = 'refund_pending'
OR (status = 'pending' AND refund_after <= NOW())
ORDER BY refund_after ASC
LIMIT $1
"#,
)
.bind(limit)
.fetch_all(&state.db)
.await
.map_err(|err| {
AppError::new(ErrorCode::Internal, "查询待补偿匿名单文件预留失败").with_source(err)
})?;
let mut settled = 0;
for task_id in task_ids {
match refund_anonymous_single_reservation(state, task_id).await {
Ok(()) => settled += 1,
Err(err) => {
tracing::warn!(task_id = %task_id, error = %err, "anonymous single reservation refund deferred")
}
}
}
Ok(settled)
}
pub async fn refund_anonymous_reservation_once(
state: &AppState,
task_id: Uuid,
@@ -541,6 +743,10 @@ fn anonymous_session_key(session_id: &str, date: NaiveDate) -> String {
format!("anon_quota:{session_id}:{}", date.format("%Y-%m-%d"))
}
fn anonymous_single_reservation_key(task_id: Uuid) -> String {
format!("anon_quota_reservation:{task_id}")
}
pub(crate) fn anonymous_ip_scope(ip: IpAddr) -> String {
match ip {
IpAddr::V4(ip) => ip.to_string(),
@@ -980,4 +1186,230 @@ mod tests {
.expect("delete quota settlement Redis keys");
}
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
#[ignore = "requires isolated IMAGEFORGE_TEST_DATABASE_URL and IMAGEFORGE_TEST_REDIS_URL"]
async fn anonymous_single_reservations_charge_actual_work_and_refund_original_date() {
let database_url = std::env::var("IMAGEFORGE_TEST_DATABASE_URL")
.expect("IMAGEFORGE_TEST_DATABASE_URL must be set");
assert!(
database_url.to_ascii_lowercase().contains("test"),
"refusing to run destructive integration test outside a test database"
);
let redis_url = std::env::var("IMAGEFORGE_TEST_REDIS_URL")
.expect("IMAGEFORGE_TEST_REDIS_URL must be set");
let pool = PgPoolOptions::new()
.max_connections(16)
.connect(&database_url)
.await
.expect("connect anonymous single test database");
sqlx::migrate!().run(&pool).await.expect("run migrations");
let state = build_test_state(pool.clone(), database_url, redis_url).await;
let marker = Uuid::new_v4().simple().to_string();
let current_date = utc8_date();
let previous_date = current_date.pred_opt().expect("previous quota date");
let mut cleanup = Vec::new();
async fn mark_and_finalize(
state: &AppState,
task_id: Uuid,
charged: bool,
) -> Result<(), AppError> {
let mut tx = state.db.begin().await.map_err(|err| {
AppError::new(ErrorCode::Internal, "begin anonymous single test").with_source(err)
})?;
mark_anonymous_single_result(&mut tx, task_id, charged).await?;
tx.commit().await.map_err(|err| {
AppError::new(ErrorCode::Internal, "commit anonymous single test").with_source(err)
})?;
finalize_anonymous_single_reservation(state, task_id, charged).await
}
let passthrough_task = Uuid::new_v4();
let passthrough_session = format!("single-passthrough-{marker}");
let passthrough_ip: IpAddr = "198.51.100.31".parse().expect("parse passthrough IP");
reserve_anonymous_single_unit_for_date(
&state,
passthrough_task,
&passthrough_session,
passthrough_ip,
current_date,
)
.await
.expect("reserve passthrough unit");
let passthrough_charged = output_consumes_unit(Some(100), true, false, false, 100, 50);
assert!(!passthrough_charged);
mark_and_finalize(&state, passthrough_task, passthrough_charged)
.await
.expect("refund passthrough unit");
assert_eq!(
anonymous_quota_counts(&state, &passthrough_session, passthrough_ip, current_date)
.await,
(0, 0)
);
cleanup.push((
passthrough_task,
passthrough_session,
passthrough_ip,
current_date,
));
let unchanged_task = Uuid::new_v4();
let unchanged_session = format!("single-unchanged-{marker}");
let unchanged_ip: IpAddr = "198.51.100.32".parse().expect("parse unchanged IP");
reserve_anonymous_single_unit_for_date(
&state,
unchanged_task,
&unchanged_session,
unchanged_ip,
current_date,
)
.await
.expect("reserve unchanged-output unit");
let unchanged_charged = output_consumes_unit(None, true, false, false, 100, 100);
assert!(!unchanged_charged);
mark_and_finalize(&state, unchanged_task, unchanged_charged)
.await
.expect("refund unchanged-output unit");
assert_eq!(
anonymous_quota_counts(&state, &unchanged_session, unchanged_ip, current_date).await,
(0, 0)
);
cleanup.push((
unchanged_task,
unchanged_session,
unchanged_ip,
current_date,
));
let compressed_task = Uuid::new_v4();
let compressed_session = format!("single-compressed-{marker}");
let compressed_ip: IpAddr = "198.51.100.33".parse().expect("parse compressed IP");
reserve_anonymous_single_unit_for_date(
&state,
compressed_task,
&compressed_session,
compressed_ip,
current_date,
)
.await
.expect("reserve compressed unit");
let compressed_charged = output_consumes_unit(None, true, false, false, 100, 50);
assert!(compressed_charged);
mark_and_finalize(&state, compressed_task, compressed_charged)
.await
.expect("finalize compressed unit");
assert_eq!(
anonymous_quota_counts(&state, &compressed_session, compressed_ip, current_date).await,
(1, 1)
);
let charged_status: String = sqlx::query_scalar(
"SELECT status FROM anonymous_single_reservations WHERE task_id = $1",
)
.bind(compressed_task)
.fetch_one(&pool)
.await
.expect("query charged reservation");
assert_eq!(charged_status, "charged");
cleanup.push((
compressed_task,
compressed_session,
compressed_ip,
current_date,
));
let cross_day_task = Uuid::new_v4();
let cross_day_session = format!("single-cross-day-{marker}");
let cross_day_ip: IpAddr = "198.51.100.34".parse().expect("parse cross-day IP");
reserve_anonymous_units(&state, &cross_day_session, cross_day_ip, 2)
.await
.expect("seed current-day quota");
reserve_anonymous_single_unit_for_date(
&state,
cross_day_task,
&cross_day_session,
cross_day_ip,
previous_date,
)
.await
.expect("reserve previous-day unit");
refund_anonymous_single_reservation(&state, cross_day_task)
.await
.expect("refund previous-day failure");
assert_eq!(
anonymous_quota_counts(&state, &cross_day_session, cross_day_ip, previous_date).await,
(0, 0)
);
assert_eq!(
anonymous_quota_counts(&state, &cross_day_session, cross_day_ip, current_date).await,
(2, 2),
"cross-day refund changed the current quota bucket"
);
cleanup.push((
cross_day_task,
cross_day_session.clone(),
cross_day_ip,
previous_date,
));
let interrupted_task = Uuid::new_v4();
let interrupted_session = format!("single-interrupted-{marker}");
let interrupted_ip: IpAddr = "198.51.100.35".parse().expect("parse interrupted IP");
reserve_anonymous_single_unit_for_date(
&state,
interrupted_task,
&interrupted_session,
interrupted_ip,
current_date,
)
.await
.expect("reserve interrupted unit");
sqlx::query(
"UPDATE anonymous_single_reservations SET refund_after = NOW() - INTERVAL '1 second' WHERE task_id = $1",
)
.bind(interrupted_task)
.execute(&pool)
.await
.expect("expire interrupted reservation");
assert_eq!(
settle_stale_anonymous_single_reservations(&state, 10)
.await
.expect("settle interrupted reservation"),
1
);
assert_eq!(
anonymous_quota_counts(&state, &interrupted_session, interrupted_ip, current_date)
.await,
(0, 0)
);
cleanup.push((
interrupted_task,
interrupted_session,
interrupted_ip,
current_date,
));
let mut redis = state.redis.clone();
for (task_id, session_id, ip, date) in cleanup {
sqlx::query("DELETE FROM anonymous_single_reservations WHERE task_id = $1")
.bind(task_id)
.execute(&pool)
.await
.expect("delete anonymous single reservation");
let _: i64 = redis::cmd("DEL")
.arg(anonymous_session_key(&session_id, date))
.arg(anonymous_ip_key(ip, date))
.arg(anonymous_single_reservation_key(task_id))
.arg(format!("anon_quota_refund:{task_id}"))
.query_async(&mut redis)
.await
.expect("delete anonymous single Redis keys");
}
let _: i64 = redis::cmd("DEL")
.arg(anonymous_session_key(&cross_day_session, current_date))
.arg(anonymous_ip_key(cross_day_ip, current_date))
.query_async(&mut redis)
.await
.expect("delete cross-day current Redis keys");
}
}

24
src/worker/mod.rs
View File

@@ -1715,6 +1715,7 @@ async fn charge_one_unit(
}
async fn maintenance(state: &AppState) -> Result<(), AppError> {
settle_stale_anonymous_single_reservations(state).await?;
settle_finished_anonymous_reservations(state).await?;
cleanup_expired_tasks(state).await?;
cleanup_stale_zip_temp(state).await?;
@@ -1722,6 +1723,19 @@ async fn maintenance(state: &AppState) -> Result<(), AppError> {
Ok(())
}
async fn settle_stale_anonymous_single_reservations(state: &AppState) -> Result<(), AppError> {
for _ in 0..MAX_MAINTENANCE_BATCHES {
let settled =
quota::settle_stale_anonymous_single_reservations(state, MAINTENANCE_BATCH_SIZE)
.await?;
if settled < MAINTENANCE_BATCH_SIZE as usize {
break;
}
tokio::task::yield_now().await;
}
Ok(())
}
async fn settle_finished_anonymous_reservations(state: &AppState) -> Result<(), AppError> {
for _ in 0..MAX_MAINTENANCE_BATCHES {
let task_ids: Vec<Uuid> = sqlx::query_scalar(
@@ -1816,6 +1830,16 @@ async fn cleanup_expired_records(state: &AppState) -> Result<(), AppError> {
.execute(&state.db)
.await;
let _ = sqlx::query(
r#"
DELETE FROM anonymous_single_reservations
WHERE status IN ('charged', 'refunded')
AND settled_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)