fix: settle quota and bound compression work

src/services/quota.rs

@@ -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"
+        );
+    }
 }