b66087de0e
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)
124 lines
3.6 KiB
Rust
124 lines
3.6 KiB
Rust
//! 计费用量聚合 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(())
|
||
}
|