fix: settle anonymous single-file reservations
This commit is contained in:
20
migrations/022_anonymous_single_reservations.sql
Normal file
20
migrations/022_anonymous_single_reservations.sql
Normal file
@@ -0,0 +1,20 @@
|
|||||||
|
CREATE TABLE anonymous_single_reservations (
|
||||||
|
task_id UUID PRIMARY KEY,
|
||||||
|
session_id VARCHAR(100) NOT NULL,
|
||||||
|
client_ip INET NOT NULL,
|
||||||
|
quota_date DATE NOT NULL,
|
||||||
|
units INTEGER NOT NULL,
|
||||||
|
status VARCHAR(20) NOT NULL DEFAULT 'pending',
|
||||||
|
refund_after TIMESTAMPTZ NOT NULL DEFAULT (NOW() + INTERVAL '15 minutes'),
|
||||||
|
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||||
|
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||||
|
settled_at TIMESTAMPTZ,
|
||||||
|
CONSTRAINT anonymous_single_reservations_units_check
|
||||||
|
CHECK (units > 0),
|
||||||
|
CONSTRAINT anonymous_single_reservations_status_check
|
||||||
|
CHECK (status IN ('pending', 'charged', 'refund_pending', 'refunded'))
|
||||||
|
);
|
||||||
|
|
||||||
|
CREATE INDEX anonymous_single_reservations_unsettled
|
||||||
|
ON anonymous_single_reservations(refund_after, created_at)
|
||||||
|
WHERE status IN ('pending', 'refund_pending');
|
||||||
@@ -79,6 +79,10 @@ fn default_units_charged() -> i32 {
|
|||||||
1
|
1
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn metered_units(charged: bool) -> i32 {
|
||||||
|
i32::from(charged)
|
||||||
|
}
|
||||||
|
|
||||||
fn direct_response<B: IntoResponse>(
|
fn direct_response<B: IntoResponse>(
|
||||||
body: B,
|
body: B,
|
||||||
format: ImageFmt,
|
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 {
|
let op: Result<CompressResponse, AppError> = (async {
|
||||||
match "a_ctx {
|
match "a_ctx {
|
||||||
QuotaContext::User(billing) => ensure_quota_available(&state, billing, 1).await?,
|
QuotaContext::User(billing) => ensure_quota_available(&state, billing, 1).await?,
|
||||||
QuotaContext::ApiKey(billing, _) => ensure_quota_available(&state, billing, 1).await?,
|
QuotaContext::ApiKey(billing, _) => ensure_quota_available(&state, billing, 1).await?,
|
||||||
QuotaContext::Anonymous { session_id, ip } => {
|
QuotaContext::Anonymous { session_id, ip } => {
|
||||||
quota::consume_anonymous_units(&state, session_id, *ip, 1).await?;
|
anonymous_reservation_date = Some(
|
||||||
anonymous_reserved = true;
|
quota::reserve_anonymous_single_unit(&state, task_id, session_id, *ip).await?,
|
||||||
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -299,17 +305,15 @@ async fn compress_json(
|
|||||||
} else {
|
} else {
|
||||||
(saved_bytes as f64) * 100.0 / (original_size as f64)
|
(saved_bytes as f64) * 100.0 / (original_size as f64)
|
||||||
};
|
};
|
||||||
let charge_units = anonymous_reserved
|
let charge_units = quota::output_consumes_unit(
|
||||||
|| quota::output_consumes_unit(
|
req.compression_rate,
|
||||||
req.compression_rate,
|
format_in == format_out,
|
||||||
format_in == format_out,
|
req.max_width.is_some() || req.max_height.is_some(),
|
||||||
req.max_width.is_some() || req.max_height.is_some(),
|
req.target_size_bytes.is_some(),
|
||||||
req.target_size_bytes.is_some(),
|
original_size,
|
||||||
original_size,
|
compressed_size,
|
||||||
compressed_size,
|
);
|
||||||
);
|
|
||||||
|
|
||||||
let task_id = Uuid::new_v4();
|
|
||||||
let file_id = Uuid::new_v4();
|
let file_id = Uuid::new_v4();
|
||||||
let retention_hours = retention.num_hours();
|
let retention_hours = retention.num_hours();
|
||||||
let object_key =
|
let object_key =
|
||||||
@@ -356,6 +360,17 @@ async fn compress_json(
|
|||||||
return Err(err);
|
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 {
|
Ok(CompressResponse {
|
||||||
task_id,
|
task_id,
|
||||||
file_id,
|
file_id,
|
||||||
@@ -368,7 +383,7 @@ async fn compress_json(
|
|||||||
download_url: format!("/downloads/{file_id}"),
|
download_url: format!("/downloads/{file_id}"),
|
||||||
expires_at,
|
expires_at,
|
||||||
billing: BillingView {
|
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) => {
|
Err(err) => {
|
||||||
if anonymous_reserved {
|
if anonymous_reservation_date.is_some() {
|
||||||
if let QuotaContext::Anonymous { session_id, ip } = "a_ctx {
|
if let Err(refund_err) =
|
||||||
let _ = quota::refund_anonymous_units(&state, session_id, *ip, 1).await;
|
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)) = (
|
if let (Some(scope), Some(idem_key), Some(request_hash)) = (
|
||||||
@@ -653,7 +670,7 @@ async fn compress_direct(
|
|||||||
compressed_size,
|
compressed_size,
|
||||||
saved_bytes,
|
saved_bytes,
|
||||||
saved_percent,
|
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);
|
let response = direct_response(compressed, format_out, &idem_data);
|
||||||
Ok((response, 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))?;
|
.map_err(|err| AppError::new(ErrorCode::Internal, "创建文件记录失败").with_source(err))?;
|
||||||
|
|
||||||
match quota_ctx {
|
match quota_ctx {
|
||||||
QuotaContext::Anonymous { .. } => {}
|
QuotaContext::Anonymous { .. } => {
|
||||||
|
quota::mark_anonymous_single_result(&mut tx, task_id, charge_units).await?;
|
||||||
|
}
|
||||||
QuotaContext::User(billing) => {
|
QuotaContext::User(billing) => {
|
||||||
if charge_units {
|
if charge_units {
|
||||||
charge_one_unit(
|
charge_one_unit(
|
||||||
@@ -1247,4 +1266,34 @@ mod tests {
|
|||||||
assert_eq!(response.headers()["imageforge-saved-percent"], "37.50");
|
assert_eq!(response.headers()["imageforge-saved-percent"], "37.50");
|
||||||
assert_eq!(response.headers()["imageforge-units-charged"], "1");
|
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);
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -320,55 +320,257 @@ pub async fn reserve_anonymous_units(
|
|||||||
Ok(date)
|
Ok(date)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn refund_anonymous_units(
|
pub async fn reserve_anonymous_single_unit(
|
||||||
state: &AppState,
|
state: &AppState,
|
||||||
|
task_id: Uuid,
|
||||||
session_id: &str,
|
session_id: &str,
|
||||||
ip: IpAddr,
|
ip: IpAddr,
|
||||||
units: u32,
|
) -> Result<NaiveDate, AppError> {
|
||||||
) -> Result<(), AppError> {
|
reserve_anonymous_single_unit_for_date(state, task_id, session_id, ip, utc8_date()).await
|
||||||
refund_anonymous_units_for_date(state, session_id, ip, utc8_date(), units).await
|
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn refund_anonymous_units_for_date(
|
async fn reserve_anonymous_single_unit_for_date(
|
||||||
state: &AppState,
|
state: &AppState,
|
||||||
|
task_id: Uuid,
|
||||||
session_id: &str,
|
session_id: &str,
|
||||||
ip: IpAddr,
|
ip: IpAddr,
|
||||||
date: NaiveDate,
|
date: NaiveDate,
|
||||||
units: u32,
|
) -> Result<NaiveDate, AppError> {
|
||||||
) -> Result<(), AppError> {
|
sqlx::query(
|
||||||
if units == 0 {
|
r#"
|
||||||
return Ok(());
|
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 session_key = anonymous_session_key(session_id, date);
|
||||||
let ip_key = anonymous_ip_key(ip, 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 mut conn = state.redis.clone();
|
||||||
let script = redis::Script::new(
|
let script = redis::Script::new(
|
||||||
r#"
|
r#"
|
||||||
local dec = tonumber(ARGV[1])
|
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)
|
if redis.call('EXISTS', KEYS[3]) == 1 then
|
||||||
local current = tonumber(redis.call('GET', key) or '0')
|
local function refund(key)
|
||||||
if current <= 0 then return 0 end
|
local current = tonumber(redis.call('GET', key) or '0')
|
||||||
return redis.call('DECRBY', key, math.min(current, dec))
|
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
|
end
|
||||||
|
redis.call('SET', KEYS[4], '1', 'EX', ttl)
|
||||||
refund(KEYS[1])
|
|
||||||
refund(KEYS[2])
|
|
||||||
return 1
|
return 1
|
||||||
"#,
|
"#,
|
||||||
);
|
);
|
||||||
|
|
||||||
let _: i64 = script
|
let _: i64 = script
|
||||||
.key(session_key)
|
.key(session_key)
|
||||||
.key(ip_key)
|
.key(ip_key)
|
||||||
|
.key(reservation_key)
|
||||||
|
.key(refund_key)
|
||||||
.arg(units as i64)
|
.arg(units as i64)
|
||||||
|
.arg(48 * 60 * 60)
|
||||||
.invoke_async(&mut conn)
|
.invoke_async(&mut conn)
|
||||||
.await
|
.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(())
|
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(
|
pub async fn refund_anonymous_reservation_once(
|
||||||
state: &AppState,
|
state: &AppState,
|
||||||
task_id: Uuid,
|
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"))
|
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 {
|
pub(crate) fn anonymous_ip_scope(ip: IpAddr) -> String {
|
||||||
match ip {
|
match ip {
|
||||||
IpAddr::V4(ip) => ip.to_string(),
|
IpAddr::V4(ip) => ip.to_string(),
|
||||||
@@ -980,4 +1186,230 @@ mod tests {
|
|||||||
.expect("delete quota settlement Redis keys");
|
.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");
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1715,6 +1715,7 @@ async fn charge_one_unit(
|
|||||||
}
|
}
|
||||||
|
|
||||||
async fn maintenance(state: &AppState) -> Result<(), AppError> {
|
async fn maintenance(state: &AppState) -> Result<(), AppError> {
|
||||||
|
settle_stale_anonymous_single_reservations(state).await?;
|
||||||
settle_finished_anonymous_reservations(state).await?;
|
settle_finished_anonymous_reservations(state).await?;
|
||||||
cleanup_expired_tasks(state).await?;
|
cleanup_expired_tasks(state).await?;
|
||||||
cleanup_stale_zip_temp(state).await?;
|
cleanup_stale_zip_temp(state).await?;
|
||||||
@@ -1722,6 +1723,19 @@ async fn maintenance(state: &AppState) -> Result<(), AppError> {
|
|||||||
Ok(())
|
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> {
|
async fn settle_finished_anonymous_reservations(state: &AppState) -> Result<(), AppError> {
|
||||||
for _ in 0..MAX_MAINTENANCE_BATCHES {
|
for _ in 0..MAX_MAINTENANCE_BATCHES {
|
||||||
let task_ids: Vec<Uuid> = sqlx::query_scalar(
|
let task_ids: Vec<Uuid> = sqlx::query_scalar(
|
||||||
@@ -1816,6 +1830,16 @@ async fn cleanup_expired_records(state: &AppState) -> Result<(), AppError> {
|
|||||||
.execute(&state.db)
|
.execute(&state.db)
|
||||||
.await;
|
.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 _ =
|
let _ =
|
||||||
sqlx::query("DELETE FROM webhook_events WHERE received_at < NOW() - INTERVAL '90 days'")
|
sqlx::query("DELETE FROM webhook_events WHERE received_at < NOW() - INTERVAL '90 days'")
|
||||||
.execute(&state.db)
|
.execute(&state.db)
|
||||||
|
|||||||
Reference in New Issue
Block a user