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:
123
crates/zclaw-saas/src/workers/aggregate_usage.rs
Normal file
123
crates/zclaw-saas/src/workers/aggregate_usage.rs
Normal 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 {
|
||||
/// 聚合的目标账户 ID(None = 聚合所有活跃账户)
|
||||
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(())
|
||||
}
|
||||
Reference in New Issue
Block a user