feat(saas): add quota middleware and usage aggregation worker

B1.3 Quota middleware:
- quota_check_middleware for relay route chain
- Checks monthly relay_requests quota before processing
- Gracefully degrades on billing service failure

B1.5 AggregateUsageWorker:
- Aggregates usage_records into billing_usage_quotas monthly
- Supports single-account and all-accounts modes
- Scheduled hourly via Worker dispatcher (6 workers total)
This commit is contained in:
iven
2026-04-02 00:06:39 +08:00
parent d06ecded34
commit b66087de0e
3 changed files with 152 additions and 3 deletions

View File

@@ -0,0 +1,123 @@
//! 计费用量聚合 Worker
//!
//! 从 usage_records 聚合当月用量到 billing_usage_quotas 表。
//! 由 Scheduler 每小时触发,或在 relay 请求完成时直接派发。
use async_trait::async_trait;
use chrono::{Datelike, Timelike};
use serde::{Deserialize, Serialize};
use sqlx::PgPool;
use crate::error::SaasResult;
use super::Worker;
/// 用量聚合参数
#[derive(Debug, Serialize, Deserialize)]
pub struct AggregateUsageArgs {
/// 聚合的目标账户 IDNone = 聚合所有活跃账户)
pub account_id: Option<String>,
}
pub struct AggregateUsageWorker;
#[async_trait]
impl Worker for AggregateUsageWorker {
type Args = AggregateUsageArgs;
fn name(&self) -> &str {
"aggregate_usage"
}
async fn perform(&self, db: &PgPool, args: Self::Args) -> SaasResult<()> {
match args.account_id {
Some(account_id) => {
aggregate_single_account(db, &account_id).await?;
}
None => {
aggregate_all_accounts(db).await?;
}
}
Ok(())
}
}
/// 聚合单个账户的当月用量
async fn aggregate_single_account(db: &PgPool, account_id: &str) -> SaasResult<()> {
// 获取或创建用量记录(确保存在)
let usage = crate::billing::service::get_or_create_usage(db, account_id).await?;
// 从 usage_records 聚合当月实际 token 用量
let now = chrono::Utc::now();
let period_start = now
.with_day(1).unwrap_or(now)
.with_hour(0).unwrap_or(now)
.with_minute(0).unwrap_or(now)
.with_second(0).unwrap_or(now)
.with_nanosecond(0).unwrap_or(now);
let aggregated: Option<(i64, i64, i64)> = sqlx::query_as(
"SELECT COALESCE(SUM(input_tokens), 0), \
COALESCE(SUM(output_tokens), 0), \
COUNT(*) \
FROM usage_records \
WHERE account_id = $1 AND created_at >= $2 AND status = 'success'"
)
.bind(account_id)
.bind(period_start)
.fetch_optional(db)
.await?;
if let Some((input_tokens, output_tokens, request_count)) = aggregated {
sqlx::query(
"UPDATE billing_usage_quotas \
SET input_tokens = $1, \
output_tokens = $2, \
relay_requests = GREATEST(relay_requests, $3::int), \
updated_at = NOW() \
WHERE id = $4"
)
.bind(input_tokens)
.bind(output_tokens)
.bind(request_count as i32)
.bind(&usage.id)
.execute(db)
.await?;
tracing::debug!(
"Aggregated usage for account {}: in={}, out={}, reqs={}",
account_id, input_tokens, output_tokens, request_count
);
}
Ok(())
}
/// 聚合所有活跃账户
async fn aggregate_all_accounts(db: &PgPool) -> SaasResult<()> {
let account_ids: Vec<String> = sqlx::query_scalar(
"SELECT DISTINCT account_id FROM billing_subscriptions \
WHERE status IN ('trial', 'active', 'past_due') \
UNION \
SELECT DISTINCT account_id FROM billing_usage_quotas \
WHERE period_start >= date_trunc('month', NOW())"
)
.fetch_all(db)
.await?;
let total = account_ids.len();
let mut errors = 0;
for account_id in &account_ids {
if let Err(e) = aggregate_single_account(db, account_id).await {
tracing::warn!("Failed to aggregate usage for {}: {}", account_id, e);
errors += 1;
}
}
tracing::info!(
"Usage aggregation complete: {} accounts, {} errors",
total, errors
);
Ok(())
}