use crate::error::{AppError, ErrorCode}; use crate::services::billing::BillingContext; use crate::state::AppState; use chrono::{DateTime, Duration, NaiveDate, Utc}; use sqlx::{FromRow, Postgres, Transaction}; use std::net::IpAddr; use uuid::Uuid; #[derive(Debug, Clone, Copy)] pub struct UserUsageBalance { pub used_units: i64, pub included_units: i64, pub bonus_units: i64, pub redeemed_units: i64, pub total_units: i64, pub remaining_units: i64, } #[derive(Debug, FromRow)] struct UsagePeriodRow { used_units: i32, bonus_units: i32, grant_used_units: i32, } #[derive(Debug, FromRow)] struct AvailableGrantRow { id: Uuid, expires_at: DateTime, } pub async fn user_usage_balance( state: &AppState, billing: &BillingContext, ) -> Result { let usage = sqlx::query_as::<_, UsagePeriodRow>( r#" SELECT used_units, bonus_units, grant_used_units FROM usage_periods WHERE user_id = $1 AND period_start = $2 AND period_end = $3 "#, ) .bind(billing.user_id) .bind(billing.period_start) .bind(billing.period_end) .fetch_optional(&state.db) .await .map_err(|err| AppError::new(ErrorCode::Internal, "查询用量失败").with_source(err))? .unwrap_or(UsagePeriodRow { used_units: 0, bonus_units: 0, grant_used_units: 0, }); let redeemed_units: i64 = sqlx::query_scalar( r#" SELECT COALESCE(SUM(remaining_units), 0)::bigint FROM unit_grants WHERE user_id = $1 AND starts_at <= NOW() AND expires_at > NOW() AND remaining_units > 0 "#, ) .bind(billing.user_id) .fetch_one(&state.db) .await .map_err(|err| AppError::new(ErrorCode::Internal, "查询兑换额度失败").with_source(err))?; Ok(calculate_user_balance( billing.plan.included_units_per_period, usage.used_units, usage.bonus_units, usage.grant_used_units, redeemed_units, )) } pub async fn ensure_user_units( state: &AppState, billing: &BillingContext, needed_units: i32, ) -> Result<(), AppError> { if needed_units <= 0 { return Ok(()); } let balance = user_usage_balance(state, billing).await?; if balance.remaining_units < i64::from(needed_units) { return Err(AppError::new(ErrorCode::QuotaExceeded, "可用配额已用完")); } Ok(()) } pub async fn consume_user_unit( tx: &mut Transaction<'_, Postgres>, billing: &BillingContext, bytes_in: u64, bytes_out: u64, ) -> Result<(), AppError> { sqlx::query( r#" INSERT INTO usage_periods (user_id, subscription_id, period_start, period_end) VALUES ($1, $2, $3, $4) ON CONFLICT (user_id, period_start, period_end) DO NOTHING "#, ) .bind(billing.user_id) .bind(billing.subscription_id) .bind(billing.period_start) .bind(billing.period_end) .execute(&mut **tx) .await .map_err(|err| AppError::new(ErrorCode::Internal, "初始化用量周期失败").with_source(err))?; let usage = sqlx::query_as::<_, UsagePeriodRow>( r#" SELECT used_units, bonus_units, grant_used_units FROM usage_periods WHERE user_id = $1 AND period_start = $2 AND period_end = $3 FOR UPDATE "#, ) .bind(billing.user_id) .bind(billing.period_start) .bind(billing.period_end) .fetch_one(&mut **tx) .await .map_err(|err| AppError::new(ErrorCode::Internal, "锁定用量周期失败").with_source(err))?; let plan_used = usage.used_units.saturating_sub(usage.grant_used_units); let plan_capacity = billing .plan .included_units_per_period .saturating_add(usage.bonus_units); let plan_available = plan_used < plan_capacity; let grant = sqlx::query_as::<_, AvailableGrantRow>( r#" SELECT id, expires_at FROM unit_grants WHERE user_id = $1 AND starts_at <= NOW() AND expires_at > NOW() AND remaining_units > 0 ORDER BY expires_at ASC, created_at ASC FOR UPDATE LIMIT 1 "#, ) .bind(billing.user_id) .fetch_optional(&mut **tx) .await .map_err(|err| AppError::new(ErrorCode::Internal, "锁定兑换额度失败").with_source(err))?; let use_grant = should_consume_grant( plan_available, grant.as_ref().map(|grant| grant.expires_at), billing.period_end, ); if !plan_available && !use_grant { return Err(AppError::new(ErrorCode::QuotaExceeded, "可用配额已用完")); } if use_grant { let grant_id = grant .ok_or_else(|| AppError::new(ErrorCode::QuotaExceeded, "可用配额已用完"))? .id; let updated = sqlx::query( r#" UPDATE unit_grants SET remaining_units = remaining_units - 1, updated_at = NOW() WHERE id = $1 AND remaining_units > 0 AND expires_at > NOW() "#, ) .bind(grant_id) .execute(&mut **tx) .await .map_err(|err| AppError::new(ErrorCode::Internal, "扣减兑换额度失败").with_source(err))?; if updated.rows_affected() != 1 { return Err(AppError::new(ErrorCode::QuotaExceeded, "可用配额已用完")); } } sqlx::query( r#" UPDATE usage_periods SET used_units = used_units + 1, grant_used_units = grant_used_units + $1, bytes_in = bytes_in + $2, bytes_out = bytes_out + $3, updated_at = NOW() WHERE user_id = $4 AND period_start = $5 AND period_end = $6 "#, ) .bind(use_grant as i32) .bind(bytes_in as i64) .bind(bytes_out as i64) .bind(billing.user_id) .bind(billing.period_start) .bind(billing.period_end) .execute(&mut **tx) .await .map_err(|err| AppError::new(ErrorCode::Internal, "记录用量失败").with_source(err))?; Ok(()) } fn calculate_user_balance( included_units: i32, used_units: i32, bonus_units: i32, grant_used_units: i32, redeemed_units: i64, ) -> UserUsageBalance { let base_capacity = i64::from(included_units.saturating_add(bonus_units)); let base_used = i64::from(used_units.saturating_sub(grant_used_units)); let base_remaining = base_capacity.saturating_sub(base_used).max(0); let remaining_units = base_remaining.saturating_add(redeemed_units.max(0)); let used_units = i64::from(used_units.max(0)); UserUsageBalance { used_units, included_units: i64::from(included_units), bonus_units: i64::from(bonus_units), redeemed_units: redeemed_units.max(0), total_units: used_units.saturating_add(remaining_units), remaining_units, } } fn should_consume_grant( plan_available: bool, grant_expires_at: Option>, plan_expires_at: DateTime, ) -> bool { grant_expires_at.is_some_and(|expires_at| !plan_available || expires_at <= plan_expires_at) } pub async fn consume_anonymous_units( state: &AppState, session_id: &str, 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 { let date = utc8_date(); if units == 0 { return Ok(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 limit = crate::services::settings::runtime_policy(state) .await? .rate_limits .anonymous_units_per_day as i64; let ttl_seconds = 48 * 60 * 60; let inc = units as i64; let script = redis::Script::new( r#" local limit = tonumber(ARGV[1]) local ttl = tonumber(ARGV[2]) local inc = tonumber(ARGV[3]) local v1 = tonumber(redis.call('GET', KEYS[1]) or '0') local v2 = tonumber(redis.call('GET', KEYS[2]) or '0') if v1 + inc > limit or v2 + inc > limit then return -1 end v1 = redis.call('INCRBY', KEYS[1], inc) v2 = redis.call('INCRBY', KEYS[2], inc) if v1 == inc then redis.call('EXPIRE', KEYS[1], ttl) end if v2 == inc then redis.call('EXPIRE', KEYS[2], ttl) end return v1 "#, ); let new_value: i64 = script .key(session_key) .key(ip_key) .arg(limit) .arg(ttl_seconds) .arg(inc) .invoke_async(&mut conn) .await .map_err(|err| AppError::new(ErrorCode::Internal, "匿名配额检查失败").with_source(err))?; if new_value < 0 { return Err(AppError::new( ErrorCode::QuotaExceeded, format!("匿名试用次数已用完(每日 {limit} 次)"), )); } Ok(date) } pub async fn refund_anonymous_units( state: &AppState, 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 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#" local dec = tonumber(ARGV[1]) 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]) return 1 "#, ); let _: i64 = script .key(session_key) .key(ip_key) .arg(units as i64) .invoke_async(&mut conn) .await .map_err(|err| AppError::new(ErrorCode::Internal, "退还匿名配额失败").with_source(err))?; Ok(()) } 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, total_files: i32, consumed_units: i32, session_id: Option, client_ip: Option, created_at: DateTime, } pub async fn settle_anonymous_task_reservation( state: &AppState, task_id: Uuid, ) -> Result, 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 ( -- NULL means the caller did not request the explicit 100% passthrough. COALESCE(tasks.compression_rate = 100, false) 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::().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)) } pub(crate) fn output_consumes_unit( compression_rate: Option, same_format: bool, has_resize: bool, has_target_size: bool, original_size: u64, output_size: u64, ) -> bool { let is_unmetered_passthrough = compression_rate == Some(100) && same_format && !has_resize && !has_target_size; !is_unmetered_passthrough && output_size < original_size } 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")) } pub(crate) fn anonymous_ip_scope(ip: IpAddr) -> String { match ip { IpAddr::V4(ip) => ip.to_string(), IpAddr::V6(ip) => { if let Some(mapped) = ip.to_ipv4_mapped() { return mapped.to_string(); } let network = std::net::Ipv6Addr::from(u128::from(ip) & (u128::MAX << 64)); format!("{network}/64") } } } fn anonymous_ip_key(ip: IpAddr, date: NaiveDate) -> String { format!( "anon_quota_ip:{}:{}", anonymous_ip_scope(ip), date.format("%Y-%m-%d") ) } fn utc8_date() -> NaiveDate { let now = Utc::now() + Duration::hours(8); now.date_naive() } #[cfg(test)] mod tests { use super::*; use crate::config::Config; use crate::services::mail::Mailer; use sqlx::postgres::PgPoolOptions; use std::sync::Arc; use tokio::sync::Semaphore; struct AnonymousSettlementFixture<'a> { session_id: &'a str, ip: IpAddr, date: NaiveDate, reserved_units: u32, compression_rate: Option, total_files: usize, completed_files: usize, } async fn build_test_state( pool: sqlx::PgPool, database_url: String, redis_url: String, ) -> AppState { let mut config = Config::from_env().expect("load quota test config"); config.database_url = database_url; config.redis_url = redis_url; config.mail_enabled = false; config.mail_log_links_when_disabled = false; config.anon_daily_units = 10; let redis = redis::Client::open(config.redis_url.clone()) .expect("create quota test Redis client") .get_connection_manager() .await .expect("connect quota test Redis"); AppState { mailer: Arc::new(Mailer::new(&config).expect("create disabled quota test mailer")), image_processing_semaphore: Arc::new(Semaphore::new(2)), runtime_policy_cache: crate::services::settings::RuntimePolicyCache::new(), storage_cache: crate::services::storage::StorageCache::new(), config, db: pool, redis, } } async fn insert_anonymous_settlement_fixture( pool: &sqlx::PgPool, fixture: AnonymousSettlementFixture<'_>, ) -> Uuid { assert!(fixture.completed_files <= fixture.total_files); let task_id = Uuid::new_v4(); sqlx::query( r#" INSERT INTO tasks ( id, session_id, client_ip, status, compression_rate, total_files, completed_files, failed_files, anonymous_units_reserved, anonymous_quota_date ) VALUES ( $1, $2, $3::inet, 'completed', $4, $5, $6, $7, $8, $9 ) "#, ) .bind(task_id) .bind(fixture.session_id) .bind(fixture.ip.to_string()) .bind(fixture.compression_rate) .bind(fixture.total_files as i32) .bind(fixture.completed_files as i32) .bind((fixture.total_files - fixture.completed_files) as i32) .bind(fixture.reserved_units as i32) .bind(fixture.date) .execute(pool) .await .expect("insert anonymous settlement task"); for index in 0..fixture.total_files { let completed = index < fixture.completed_files; sqlx::query( r#" INSERT INTO task_files ( id, task_id, original_name, original_format, output_format, original_size, compressed_size, status ) VALUES ( $1, $2, $3, 'jpeg', 'jpeg', 100, $4, $5::file_status ) "#, ) .bind(Uuid::new_v4()) .bind(task_id) .bind(format!("fixture-{index}.jpg")) .bind(completed.then_some(50_i64)) .bind(if completed { "completed" } else { "failed" }) .execute(pool) .await .expect("insert anonymous settlement file"); } task_id } async fn anonymous_quota_counts( state: &AppState, session_id: &str, ip: IpAddr, date: NaiveDate, ) -> (i64, i64) { let mut redis = state.redis.clone(); let session_count: Option = redis::cmd("GET") .arg(anonymous_session_key(session_id, date)) .query_async(&mut redis) .await .expect("read anonymous session quota"); let ip_count: Option = redis::cmd("GET") .arg(anonymous_ip_key(ip, date)) .query_async(&mut redis) .await .expect("read anonymous IP quota"); (session_count.unwrap_or(0), ip_count.unwrap_or(0)) } #[test] fn balance_keeps_redeemed_units_separate_from_plan_usage() { let balance = calculate_user_balance(10, 13, 0, 3, 7); assert_eq!(balance.remaining_units, 7); assert_eq!(balance.total_units, 20); } #[test] fn balance_uses_plan_capacity_before_redeemed_units() { let balance = calculate_user_balance(10, 4, 2, 0, 5); assert_eq!(balance.remaining_units, 13); assert_eq!(balance.total_units, 17); } #[test] fn earlier_expiring_entitlement_is_consumed_first() { let now = Utc::now(); assert!(should_consume_grant( true, Some(now + Duration::days(2)), now + Duration::days(20), )); assert!(!should_consume_grant( true, Some(now + Duration::days(30)), now + Duration::days(20), )); assert!(should_consume_grant( false, Some(now + Duration::days(30)), 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 output_metering_matches_passthrough_contract() { assert!(output_consumes_unit(None, true, false, false, 100, 50)); assert!(!output_consumes_unit( Some(100), true, false, false, 100, 50 )); assert!(output_consumes_unit( Some(100), false, false, false, 100, 50 )); assert!(output_consumes_unit(Some(100), true, true, false, 100, 50)); assert!(!output_consumes_unit(None, true, false, false, 100, 100)); } #[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" ); } #[test] fn anonymous_ipv6_addresses_share_a_slash_64_scope() { let first: IpAddr = "2001:db8:1234:5678::1".parse().unwrap(); let second: IpAddr = "2001:db8:1234:5678:ffff::99".parse().unwrap(); let other: IpAddr = "2001:db8:1234:5679::1".parse().unwrap(); assert_eq!(anonymous_ip_scope(first), "2001:db8:1234:5678::/64"); assert_eq!(anonymous_ip_scope(first), anonymous_ip_scope(second)); assert_ne!(anonymous_ip_scope(first), anonymous_ip_scope(other)); assert_eq!( anonymous_ip_scope("::ffff:203.0.113.7".parse().unwrap()), "203.0.113.7" ); } #[tokio::test(flavor = "multi_thread", worker_threads = 4)] #[ignore = "requires isolated IMAGEFORGE_TEST_DATABASE_URL and IMAGEFORGE_TEST_REDIS_URL"] async fn anonymous_batch_settlement_charges_successes_and_refunds_only_unused_units() { 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 quota 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 mut cleanup = Vec::new(); let null_session = format!("quota-null-{marker}"); let null_ip: IpAddr = "198.51.100.11".parse().expect("parse fixture IP"); let null_date = reserve_anonymous_units(&state, &null_session, null_ip, 3) .await .expect("reserve NULL-rate batch quota"); let null_task = insert_anonymous_settlement_fixture( &pool, AnonymousSettlementFixture { session_id: &null_session, ip: null_ip, date: null_date, reserved_units: 3, compression_rate: None, total_files: 3, completed_files: 3, }, ) .await; cleanup.push((null_task, null_session.clone(), null_ip, null_date)); assert_eq!( settle_anonymous_task_reservation(&state, null_task) .await .expect("settle NULL-rate batch"), Some(0) ); assert_eq!( anonymous_quota_counts(&state, &null_session, null_ip, null_date).await, (3, 3) ); let passthrough_session = format!("quota-passthrough-{marker}"); let passthrough_ip: IpAddr = "198.51.100.12".parse().expect("parse fixture IP"); let passthrough_date = reserve_anonymous_units(&state, &passthrough_session, passthrough_ip, 3) .await .expect("reserve passthrough batch quota"); let passthrough_task = insert_anonymous_settlement_fixture( &pool, AnonymousSettlementFixture { session_id: &passthrough_session, ip: passthrough_ip, date: passthrough_date, reserved_units: 3, compression_rate: Some(100), total_files: 3, completed_files: 3, }, ) .await; cleanup.push(( passthrough_task, passthrough_session.clone(), passthrough_ip, passthrough_date, )); assert_eq!( settle_anonymous_task_reservation(&state, passthrough_task) .await .expect("settle passthrough batch"), Some(3) ); assert_eq!( anonymous_quota_counts( &state, &passthrough_session, passthrough_ip, passthrough_date, ) .await, (0, 0) ); let partial_session = format!("quota-partial-{marker}"); let partial_ip: IpAddr = "198.51.100.13".parse().expect("parse fixture IP"); let partial_date = reserve_anonymous_units(&state, &partial_session, partial_ip, 3) .await .expect("reserve partial batch quota"); let partial_task = insert_anonymous_settlement_fixture( &pool, AnonymousSettlementFixture { session_id: &partial_session, ip: partial_ip, date: partial_date, reserved_units: 3, compression_rate: None, total_files: 3, completed_files: 2, }, ) .await; cleanup.push(( partial_task, partial_session.clone(), partial_ip, partial_date, )); assert_eq!( settle_anonymous_task_reservation(&state, partial_task) .await .expect("settle partial batch"), Some(1) ); assert_eq!( anonymous_quota_counts(&state, &partial_session, partial_ip, partial_date).await, (2, 2) ); let limit_session = format!("quota-limit-{marker}"); let limit_ip: IpAddr = "198.51.100.14".parse().expect("parse fixture IP"); let mut limit_date = None; for batch in 0..2 { let date = reserve_anonymous_units(&state, &limit_session, limit_ip, 5) .await .expect("reserve consecutive anonymous batch"); limit_date = Some(date); let task_id = insert_anonymous_settlement_fixture( &pool, AnonymousSettlementFixture { session_id: &limit_session, ip: limit_ip, date, reserved_units: 5, compression_rate: None, total_files: 5, completed_files: 5, }, ) .await; cleanup.push((task_id, limit_session.clone(), limit_ip, date)); assert_eq!( settle_anonymous_task_reservation(&state, task_id) .await .expect("settle consecutive anonymous batch"), Some(0), "batch {batch} unexpectedly refunded consumed units" ); } let limit_error = reserve_anonymous_units(&state, &limit_session, limit_ip, 1) .await .expect_err("daily anonymous quota was bypassed"); assert_eq!(limit_error.code, ErrorCode::QuotaExceeded); assert_eq!( anonymous_quota_counts( &state, &limit_session, limit_ip, limit_date.expect("limit quota date"), ) .await, (10, 10) ); let mut redis = state.redis.clone(); for (task_id, session_id, ip, date) in cleanup { sqlx::query("DELETE FROM tasks WHERE id = $1") .bind(task_id) .execute(&pool) .await .expect("delete quota settlement fixture"); let _: i64 = redis::cmd("DEL") .arg(anonymous_session_key(&session_id, date)) .arg(anonymous_ip_key(ip, date)) .arg(format!("anon_quota_refund:{task_id}")) .query_async(&mut redis) .await .expect("delete quota settlement Redis keys"); } } }