mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 08:06:19 +00:00
perf(store): move client budget settlement off request path
Track pending charges in memory and settle them with a supervised worker. Include pending costs in budget checks without retrying historical writes at admission, and reject new admissions when the backlog reaches its limit. Use idempotent per-key batches with write deadlines, adaptive batch sizes, and bounded shutdown draining. Keep load tests outside the regular suite. Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
Codex
@@ -891,7 +891,7 @@ impl StreamingExecutionBody {
|
||||
}
|
||||
continue;
|
||||
}
|
||||
// 已提交的流在确认 Core 终结后才交付成功终态，不能先输出 completed 再报结算失败
|
||||
// 已提交的流先确认 Core 终结，避免后续 EOF 验证推翻已交付的成功终态
|
||||
if self.awaiting_terminal_eof && self.committed {
|
||||
self.verify_terminal_eof().await?;
|
||||
encoded.push_back(done_frame());
|
||||
|
||||
@@ -44,9 +44,9 @@ pub struct ClientBudgetCharge {
|
||||
pub struct ClientBudgetError(#[source] pub Option<crate::error::ErrorSource>);
|
||||
|
||||
pub trait ClientBudgetPort: Send + Sync {
|
||||
/// 原子检查当前限额与已用金额，不创建预扣费或待结算记录
|
||||
/// 检查当前限额、已用金额与已接管的费用，不预占未来费用
|
||||
fn admit(&self, key_id: ClientApiKeyId) -> BoxFuture<'_, Result<(), GatewayError>>;
|
||||
|
||||
/// 按网关请求 ID 幂等累计已取得费用；写入失败由 Store 重试
|
||||
/// 按网关请求 ID 幂等接管已取得费用；运行时入队后返回，由 Store 后台落库和重试
|
||||
fn settle(&self, charge: ClientBudgetCharge) -> BoxFuture<'_, Result<(), ClientBudgetError>>;
|
||||
}
|
||||
|
||||
@@ -32,7 +32,7 @@ pub(super) async fn settle_budget(
|
||||
diagnostics,
|
||||
&request_id,
|
||||
"settle_client_budget",
|
||||
"Client budget settlement failed; storage will retry on the next request",
|
||||
"Client budget settlement could not be accepted",
|
||||
error,
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -241,6 +241,10 @@ async fn connect(
|
||||
)?);
|
||||
let (admissions, admission_release_writer) =
|
||||
redis::BufferedClientAdmissionPort::new(admissions, diagnostics.clone());
|
||||
let (budgets, budget_writer) = postgres::BufferedClientBudgetStore::new(
|
||||
postgres::PgClientBudgetStore::new(pool.clone()).with_timezone(config.timezone),
|
||||
diagnostics.clone(),
|
||||
);
|
||||
let core_ports = CoreStorePorts::new(
|
||||
execution,
|
||||
(
|
||||
@@ -260,9 +264,7 @@ async fn connect(
|
||||
Arc::new(client_key_usage),
|
||||
diagnostics.clone(),
|
||||
)
|
||||
.with_budget(Arc::new(
|
||||
postgres::PgClientBudgetStore::new(pool.clone()).with_timezone(config.timezone),
|
||||
))
|
||||
.with_budget(Arc::new(budgets))
|
||||
.with_session_affinity(Arc::clone(&provider_session_affinity))
|
||||
.with_client_retry_store(Arc::new(redis::RedisClientRetryRepository::new(
|
||||
redis_connection.clone(),
|
||||
@@ -304,6 +306,7 @@ async fn connect(
|
||||
execution_writer,
|
||||
client_key_usage_writer,
|
||||
admission_release_writer,
|
||||
budget_writer,
|
||||
)?,
|
||||
None,
|
||||
),
|
||||
@@ -313,6 +316,7 @@ async fn connect(
|
||||
execution: execution_writer,
|
||||
client_key_usage: client_key_usage_writer,
|
||||
admission_release: admission_release_writer,
|
||||
budget: budget_writer,
|
||||
}),
|
||||
),
|
||||
};
|
||||
|
||||
@@ -0,0 +1,344 @@
|
||||
//! 费用登记与后台落库分离，准入同时检查尚未落库的精确费用
|
||||
|
||||
use std::{
|
||||
collections::{BTreeMap, VecDeque},
|
||||
num::NonZeroUsize,
|
||||
sync::{Arc, Mutex},
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
|
||||
use futures::future::BoxFuture;
|
||||
use gateway_core::{
|
||||
diagnostics::{OperationalDiagnostics, OperationalFailure},
|
||||
engine::budget::{ClientBudgetCharge, ClientBudgetError, ClientBudgetPort},
|
||||
error::{ErrorDetails, GatewayError},
|
||||
policy::ClientApiKeyId,
|
||||
task::{DaemonTask, WorkerDaemonContext, WorkerTaskError},
|
||||
};
|
||||
use tokio::sync::Notify;
|
||||
|
||||
use super::{PgClientBudgetStore, client_budgets::unavailable};
|
||||
|
||||
const DEFAULT_PENDING_CAPACITY: NonZeroUsize =
|
||||
NonZeroUsize::new(4_096).expect("non-zero queue capacity");
|
||||
const ADMISSION_TIMEOUT: Duration = Duration::from_secs(1);
|
||||
const WRITE_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
const DRAIN_TIMEOUT: Duration = Duration::from_secs(3);
|
||||
const WRITE_BATCH_SIZE: usize = 32;
|
||||
|
||||
#[derive(Default)]
|
||||
struct PendingCharges {
|
||||
by_key: BTreeMap<String, BTreeMap<String, ClientBudgetCharge>>,
|
||||
ready: VecDeque<(String, String)>,
|
||||
count: usize,
|
||||
closing: bool,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct BudgetQueue {
|
||||
pending: Mutex<PendingCharges>,
|
||||
changed: Notify,
|
||||
}
|
||||
|
||||
/// 单副本运行时先接管费用，再由受监督的 writer 落库和重试
|
||||
pub struct BufferedClientBudgetStore {
|
||||
inner: Arc<PgClientBudgetStore>,
|
||||
queue: Arc<BudgetQueue>,
|
||||
capacity: NonZeroUsize,
|
||||
}
|
||||
|
||||
impl BufferedClientBudgetStore {
|
||||
#[must_use]
|
||||
pub fn new(
|
||||
inner: PgClientBudgetStore,
|
||||
diagnostics: Arc<dyn OperationalDiagnostics>,
|
||||
) -> (Self, ClientBudgetWriter) {
|
||||
Self::with_capacity(inner, DEFAULT_PENDING_CAPACITY, diagnostics)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_capacity(
|
||||
inner: PgClientBudgetStore,
|
||||
capacity: NonZeroUsize,
|
||||
diagnostics: Arc<dyn OperationalDiagnostics>,
|
||||
) -> (Self, ClientBudgetWriter) {
|
||||
let inner = Arc::new(inner);
|
||||
let queue = Arc::new(BudgetQueue::default());
|
||||
(
|
||||
Self {
|
||||
inner: inner.clone(),
|
||||
queue: queue.clone(),
|
||||
capacity,
|
||||
},
|
||||
ClientBudgetWriter {
|
||||
inner,
|
||||
queue,
|
||||
diagnostics,
|
||||
batch_limits: tokio::sync::Mutex::new(BTreeMap::new()),
|
||||
},
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl ClientBudgetPort for BufferedClientBudgetStore {
|
||||
fn admit(&self, key_id: ClientApiKeyId) -> BoxFuture<'_, Result<(), GatewayError>> {
|
||||
Box::pin(async move {
|
||||
let pending = {
|
||||
let state = self.queue.pending.lock().map_err(|_| unavailable())?;
|
||||
if state.closing || state.count >= self.capacity.get() {
|
||||
return Err(unavailable());
|
||||
}
|
||||
state
|
||||
.by_key
|
||||
.get(key_id.as_str())
|
||||
.map(|charges| charges.values().cloned().collect::<Vec<_>>())
|
||||
.unwrap_or_default()
|
||||
};
|
||||
tokio::time::timeout(
|
||||
ADMISSION_TIMEOUT,
|
||||
self.inner.admit_with_pending(key_id, &pending),
|
||||
)
|
||||
.await
|
||||
.map_err(|source| unavailable().with_source(source))?
|
||||
})
|
||||
}
|
||||
|
||||
fn settle(&self, charge: ClientBudgetCharge) -> BoxFuture<'_, Result<(), ClientBudgetError>> {
|
||||
Box::pin(async move {
|
||||
let mut state = self
|
||||
.queue
|
||||
.pending
|
||||
.lock()
|
||||
.map_err(|_| ClientBudgetError(None))?;
|
||||
let key = charge.key_id.as_str().to_owned();
|
||||
let request = charge.request_id.as_str().to_owned();
|
||||
let charges = state.by_key.entry(key.clone()).or_default();
|
||||
if charges.contains_key(&request) {
|
||||
return Ok(());
|
||||
}
|
||||
charges.insert(request.clone(), charge);
|
||||
state.ready.push_back((key, request));
|
||||
state.count += 1;
|
||||
// 达到容量后拒绝新准入，但已准入请求的费用必须接收，不能丢弃或反向等待写入
|
||||
self.queue.changed.notify_one();
|
||||
if state.closing {
|
||||
return Err(ClientBudgetError(None));
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// 每次仅占一个写入任务；任务重启、取消和超时后仍保留同一笔费用
|
||||
pub struct ClientBudgetWriter {
|
||||
inner: Arc<PgClientBudgetStore>,
|
||||
queue: Arc<BudgetQueue>,
|
||||
diagnostics: Arc<dyn OperationalDiagnostics>,
|
||||
// 独占 writer，同时保留各 Key 跨任务重启的缩批进度
|
||||
batch_limits: tokio::sync::Mutex<BTreeMap<String, usize>>,
|
||||
}
|
||||
|
||||
struct PendingWrite {
|
||||
queue: Arc<BudgetQueue>,
|
||||
key: String,
|
||||
charges: Vec<ClientBudgetCharge>,
|
||||
}
|
||||
|
||||
impl PendingWrite {
|
||||
fn complete(&self) -> Result<bool, WorkerTaskError> {
|
||||
let mut state = self.queue.pending.lock().map_err(|_| queue_error())?;
|
||||
if let Some(charges) = state.by_key.get_mut(&self.key) {
|
||||
let before = charges.len();
|
||||
for charge in &self.charges {
|
||||
charges.remove(charge.request_id.as_str());
|
||||
}
|
||||
let removed = before - charges.len();
|
||||
let empty = charges.is_empty();
|
||||
state.count -= removed;
|
||||
if empty {
|
||||
state.by_key.remove(&self.key);
|
||||
}
|
||||
}
|
||||
Ok(!state.by_key.contains_key(&self.key))
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for PendingWrite {
|
||||
fn drop(&mut self) {
|
||||
let Ok(mut state) = self.queue.pending.lock() else {
|
||||
return;
|
||||
};
|
||||
// 只有提交确认后才移除；future 被取消时重排队，由请求 ID 处理提交结果未知
|
||||
for charge in &self.charges {
|
||||
let request = charge.request_id.as_str();
|
||||
if state
|
||||
.by_key
|
||||
.get(&self.key)
|
||||
.is_some_and(|charges| charges.contains_key(request))
|
||||
{
|
||||
state
|
||||
.ready
|
||||
.push_back((self.key.clone(), request.to_owned()));
|
||||
self.queue.changed.notify_one();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ClientBudgetWriter {
|
||||
fn next(
|
||||
&self,
|
||||
batch_limits: &BTreeMap<String, usize>,
|
||||
) -> Result<Option<PendingWrite>, WorkerTaskError> {
|
||||
let mut state = self.queue.pending.lock().map_err(|_| queue_error())?;
|
||||
let Some((key, _)) = state.ready.front() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let key = key.clone();
|
||||
let limit = batch_limits.get(&key).copied().unwrap_or(WRITE_BATCH_SIZE);
|
||||
let mut requests = Vec::with_capacity(limit);
|
||||
// 混合 Key 的请求也能合批；限制单批数量，并保持其他 Key 的相对顺序
|
||||
state.ready.retain(|(candidate, request)| {
|
||||
if candidate == &key && requests.len() < limit {
|
||||
requests.push(request.clone());
|
||||
false
|
||||
} else {
|
||||
true
|
||||
}
|
||||
});
|
||||
let charges = requests
|
||||
.iter()
|
||||
.map(|request| {
|
||||
state
|
||||
.by_key
|
||||
.get(&key)
|
||||
.and_then(|charges| charges.get(request))
|
||||
.cloned()
|
||||
.ok_or_else(queue_error)
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
Ok(Some(PendingWrite {
|
||||
queue: self.queue.clone(),
|
||||
key,
|
||||
charges,
|
||||
}))
|
||||
}
|
||||
|
||||
async fn write(
|
||||
&self,
|
||||
batch_limits: &mut BTreeMap<String, usize>,
|
||||
pending: PendingWrite,
|
||||
) -> Result<(), WorkerTaskError> {
|
||||
let started = Instant::now();
|
||||
let result =
|
||||
match tokio::time::timeout(WRITE_TIMEOUT, self.inner.settle_charges(&pending.charges))
|
||||
.await
|
||||
{
|
||||
Ok(result) => result,
|
||||
Err(source) => Err(ClientBudgetError(Some(source.into()))),
|
||||
};
|
||||
if let Err(source) = result {
|
||||
// 失败后只缩小当前 Key，避免固定大批次持续超过写入时限
|
||||
batch_limits.insert(pending.key.clone(), (pending.charges.len() / 2).max(1));
|
||||
let mut failure = OperationalFailure::new(
|
||||
"store",
|
||||
"settle_client_budget",
|
||||
"budget_unavailable",
|
||||
"Client budget settlement failed; background writer will retry",
|
||||
);
|
||||
failure.correlation_id = pending
|
||||
.charges
|
||||
.first()
|
||||
.map(|charge| charge.request_id.as_str().to_owned());
|
||||
failure.details = ErrorDetails::capture(Some(&source.clone().into()), None, false);
|
||||
let _ = self.diagnostics.record_failure(failure).await;
|
||||
return Err(
|
||||
WorkerTaskError::safe("client budget settlement failed").with_source(source)
|
||||
);
|
||||
}
|
||||
if pending.complete()? {
|
||||
batch_limits.remove(&pending.key);
|
||||
} else if let Some(limit) = batch_limits.get_mut(&pending.key)
|
||||
&& pending.charges.len() == *limit
|
||||
&& started.elapsed() < WRITE_TIMEOUT / 4
|
||||
{
|
||||
// 满批写入足够快才恢复吞吐，保留余量避免慢库反复放大到超时
|
||||
*limit = (*limit * 2).min(WRITE_BATCH_SIZE);
|
||||
if *limit == WRITE_BATCH_SIZE {
|
||||
batch_limits.remove(&pending.key);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn drain(
|
||||
&self,
|
||||
batch_limits: &mut BTreeMap<String, usize>,
|
||||
) -> Result<(), WorkerTaskError> {
|
||||
self.queue
|
||||
.pending
|
||||
.lock()
|
||||
.map_err(|_| queue_error())?
|
||||
.closing = true;
|
||||
let result = tokio::time::timeout(DRAIN_TIMEOUT, async {
|
||||
while let Some(pending) = self.next(batch_limits)? {
|
||||
self.write(batch_limits, pending).await?;
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
.await;
|
||||
let result = result
|
||||
.map_err(|source| {
|
||||
WorkerTaskError::safe("client budget shutdown drain timed out").with_source(source)
|
||||
})
|
||||
.and_then(std::convert::identity);
|
||||
if let Err(error) = &result {
|
||||
tracing::warn!(operation = "settle_client_budget", %error, "费用队列关闭排空未完成");
|
||||
}
|
||||
result
|
||||
}
|
||||
}
|
||||
|
||||
impl DaemonTask for ClientBudgetWriter {
|
||||
fn run(&self, context: WorkerDaemonContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
Box::pin(async move {
|
||||
let mut batch_limits = self.batch_limits.lock().await;
|
||||
loop {
|
||||
let write = async {
|
||||
match self.next(&batch_limits)? {
|
||||
Some(pending) => {
|
||||
self.write(&mut batch_limits, pending).await?;
|
||||
context.succeeded();
|
||||
}
|
||||
None => self.queue.changed.notified().await,
|
||||
}
|
||||
Ok::<(), WorkerTaskError>(())
|
||||
};
|
||||
tokio::select! {
|
||||
biased;
|
||||
() = context.cancellation().cancelled() => return self.drain(&mut batch_limits).await,
|
||||
result = write => result?,
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for ClientBudgetWriter {
|
||||
fn drop(&mut self) {
|
||||
if let Ok(mut state) = self.queue.pending.lock() {
|
||||
state.closing = true;
|
||||
if state.count > 0 {
|
||||
tracing::warn!(
|
||||
operation = "settle_client_budget",
|
||||
pending_charges = state.count,
|
||||
"费用写入任务已销毁，仍有未落库费用"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn queue_error() -> WorkerTaskError {
|
||||
WorkerTaskError::safe("client budget settlement queue is unavailable")
|
||||
}
|
||||
@@ -1,7 +1,7 @@
|
||||
//! 按 Key 串行检查限额，并幂等累计已取得的 USD 费用
|
||||
//! 读取 Key 预算快照，并按 Key 幂等累计已取得的 USD 费用
|
||||
|
||||
use gateway_admin::model::audit::MutationAuditOperation;
|
||||
use std::{collections::BTreeMap, sync::Mutex, time::Duration};
|
||||
use std::{collections::BTreeMap, time::Duration};
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use futures::future::BoxFuture;
|
||||
@@ -129,7 +129,6 @@ async fn reset_client_key_budget_in_transaction(
|
||||
pub struct PgClientBudgetStore {
|
||||
timezone: gateway_core::time::DeploymentTimeZone,
|
||||
pool: PgPool,
|
||||
retry: Mutex<BTreeMap<String, ClientBudgetCharge>>,
|
||||
}
|
||||
|
||||
impl PgClientBudgetStore {
|
||||
@@ -138,7 +137,6 @@ impl PgClientBudgetStore {
|
||||
Self {
|
||||
timezone: Default::default(),
|
||||
pool,
|
||||
retry: Mutex::new(BTreeMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -148,88 +146,103 @@ impl PgClientBudgetStore {
|
||||
self
|
||||
}
|
||||
|
||||
async fn admit_inner(&self, key_id: ClientApiKeyId) -> Result<(), GatewayError> {
|
||||
// 短暂存储故障后按原金额重试；进程退出丢失的费用不转成人工核账或阻断 Key
|
||||
let retries = self
|
||||
.retry
|
||||
.lock()
|
||||
.map_err(|_| unavailable())?
|
||||
.values()
|
||||
.filter(|charge| charge.key_id == key_id)
|
||||
.cloned()
|
||||
pub(super) async fn admit_with_pending(
|
||||
&self,
|
||||
key_id: ClientApiKeyId,
|
||||
pending: &[ClientBudgetCharge],
|
||||
) -> Result<(), GatewayError> {
|
||||
// 同一语句快照同时读取窗口和已落库请求，避免提交与队列移除之间重复计算费用
|
||||
// 普通准入只读，不等待后台结算持有的 Key 行锁
|
||||
let ids = pending
|
||||
.iter()
|
||||
.map(|charge| charge.request_id.as_str())
|
||||
.collect::<Vec<_>>();
|
||||
for charge in retries {
|
||||
self.settle(charge)
|
||||
.await
|
||||
.map_err(|source| unavailable().with_source(source))?;
|
||||
}
|
||||
let mut tx = self
|
||||
.pool
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|source| unavailable().with_source(source))?;
|
||||
let row = sqlx::query(
|
||||
"select daily_limit_usd::text, weekly_limit_usd::text, enabled
|
||||
from client_api_keys where id = $1 for update",
|
||||
"select k.enabled, k.daily_limit_usd::text, k.weekly_limit_usd::text,
|
||||
w.daily_start, w.daily_end, w.weekly_start, w.weekly_end,
|
||||
coalesce(w.daily_used_usd, 0)::text as daily_used,
|
||||
coalesce(w.weekly_used_usd, 0)::text as weekly_used,
|
||||
case when k.daily_limit_usd = 0 and k.weekly_limit_usd = 0 then array[]::text[]
|
||||
else array(select request_id from client_api_key_charge_events where request_id = any($2))
|
||||
end as persisted
|
||||
from client_api_keys k left join client_api_key_budget_windows w on w.client_api_key_id = k.id
|
||||
where k.id = $1",
|
||||
)
|
||||
.bind(key_id.as_str())
|
||||
.fetch_optional(&mut *tx)
|
||||
.await
|
||||
.bind(key_id.as_str()).bind(ids).fetch_optional(&self.pool).await
|
||||
.map_err(|source| unavailable().with_source(source))?
|
||||
.ok_or_else(|| {
|
||||
GatewayError::new(
|
||||
GatewayErrorKind::Unauthorized,
|
||||
"client API key no longer exists",
|
||||
)
|
||||
})?;
|
||||
.ok_or_else(|| GatewayError::new(GatewayErrorKind::Unauthorized, "client API key no longer exists"))?;
|
||||
if !row.get::<bool, _>("enabled") {
|
||||
return Err(GatewayError::new(
|
||||
GatewayErrorKind::PolicyDenied,
|
||||
"client API key is disabled",
|
||||
));
|
||||
}
|
||||
let parse = |field: &str| -> Result<Decimal, GatewayError> {
|
||||
row.get::<String, _>(field)
|
||||
.parse()
|
||||
.map_err(|source| unavailable().with_source(source))
|
||||
};
|
||||
let limits = ClientBudgetLimits {
|
||||
daily_usd: row
|
||||
.get::<String, _>("daily_limit_usd")
|
||||
.parse()
|
||||
.map_err(|source| unavailable().with_source(source))?,
|
||||
weekly_usd: row
|
||||
.get::<String, _>("weekly_limit_usd")
|
||||
.parse()
|
||||
.map_err(|source| unavailable().with_source(source))?,
|
||||
daily_usd: parse("daily_limit_usd")?,
|
||||
weekly_usd: parse("weekly_limit_usd")?,
|
||||
};
|
||||
let now = Utc::now();
|
||||
advance_windows(&mut tx, key_id.as_str(), now, now, self.timezone)
|
||||
.await
|
||||
.map_err(|source| unavailable().with_source(source))?;
|
||||
let day = self.timezone.day_start(now).ok_or_else(unavailable)?;
|
||||
let window = |start: &str, end: &str, used: &str, days| -> Result<_, GatewayError> {
|
||||
let start = row.get::<Option<DateTime<Utc>>, _>(start);
|
||||
let end = row.get::<Option<DateTime<Utc>>, _>(end);
|
||||
let reopen = end.is_none_or(|end| end <= now);
|
||||
if reopen {
|
||||
Ok((
|
||||
day.max(start.unwrap_or(day)).max(end.unwrap_or(day)),
|
||||
self.timezone
|
||||
.days_after(now, days)
|
||||
.ok_or_else(unavailable)?,
|
||||
Decimal::ZERO,
|
||||
true,
|
||||
))
|
||||
} else {
|
||||
Ok((
|
||||
start.ok_or_else(unavailable)?,
|
||||
end.ok_or_else(unavailable)?,
|
||||
parse(used)?,
|
||||
false,
|
||||
))
|
||||
}
|
||||
};
|
||||
let (daily_start, daily_end, mut daily, open_daily) =
|
||||
window("daily_start", "daily_end", "daily_used", 1)?;
|
||||
let (weekly_start, weekly_end, mut weekly, open_weekly) =
|
||||
window("weekly_start", "weekly_end", "weekly_used", 7)?;
|
||||
if limits.is_limited() {
|
||||
let window = sqlx::query(
|
||||
"select daily_used_usd::text, weekly_used_usd::text, daily_end, weekly_end
|
||||
from client_api_key_budget_windows where client_api_key_id = $1",
|
||||
)
|
||||
.bind(key_id.as_str())
|
||||
.fetch_one(&mut *tx)
|
||||
.await
|
||||
.map_err(|source| unavailable().with_source(source))?;
|
||||
let daily: Decimal = window
|
||||
.get::<String, _>("daily_used_usd")
|
||||
.parse()
|
||||
.map_err(|source| unavailable().with_source(source))?;
|
||||
let weekly: Decimal = window
|
||||
.get::<String, _>("weekly_used_usd")
|
||||
.parse()
|
||||
.map_err(|source| unavailable().with_source(source))?;
|
||||
let persisted = row
|
||||
.get::<Vec<String>, _>("persisted")
|
||||
.into_iter()
|
||||
.collect::<std::collections::BTreeSet<_>>();
|
||||
for charge in pending {
|
||||
if persisted.contains(charge.request_id.as_str()) {
|
||||
continue;
|
||||
}
|
||||
let completed_at = DateTime::<Utc>::from(charge.completed_at);
|
||||
if completed_at >= daily_start && completed_at < daily_end {
|
||||
daily = daily
|
||||
.checked_add(charge.amount_usd)
|
||||
.ok_or_else(unavailable)?;
|
||||
}
|
||||
if completed_at >= weekly_start && completed_at < weekly_end {
|
||||
weekly = weekly
|
||||
.checked_add(charge.amount_usd)
|
||||
.ok_or_else(unavailable)?;
|
||||
}
|
||||
}
|
||||
let daily_exceeded = limits.daily_usd != Decimal::ZERO && daily >= limits.daily_usd;
|
||||
let weekly_exceeded = limits.weekly_usd != Decimal::ZERO && weekly >= limits.weekly_usd;
|
||||
if daily_exceeded || weekly_exceeded {
|
||||
let daily_end: DateTime<Utc> = window.get("daily_end");
|
||||
let weekly_end: DateTime<Utc> = window.get("weekly_end");
|
||||
let reset = if weekly_exceeded {
|
||||
weekly_end
|
||||
} else {
|
||||
daily_end
|
||||
};
|
||||
let retry = (reset - now).to_std().unwrap_or(Duration::from_secs(1));
|
||||
return Err(GatewayError::new(
|
||||
GatewayErrorKind::RateLimited,
|
||||
"client API key budget is exhausted",
|
||||
@@ -239,30 +252,59 @@ impl PgClientBudgetStore {
|
||||
} else {
|
||||
"key_daily_budget_exceeded"
|
||||
})
|
||||
.with_retry_after(retry));
|
||||
.with_retry_after((reset - now).to_std().unwrap_or(Duration::from_secs(1))));
|
||||
}
|
||||
}
|
||||
tx.commit()
|
||||
.await
|
||||
.map_err(|source| unavailable().with_source(source))
|
||||
// 首次使用或窗口到期才写入；与重置、结算保持相同锁序
|
||||
if open_daily || open_weekly {
|
||||
let mut tx = self
|
||||
.pool
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|source| unavailable().with_source(source))?;
|
||||
sqlx::query("select id from client_api_keys where id = $1 for update")
|
||||
.bind(key_id.as_str())
|
||||
.fetch_optional(&mut *tx)
|
||||
.await
|
||||
.map_err(|source| unavailable().with_source(source))?
|
||||
.ok_or_else(|| {
|
||||
GatewayError::new(
|
||||
GatewayErrorKind::Unauthorized,
|
||||
"client API key no longer exists",
|
||||
)
|
||||
})?;
|
||||
advance_windows(&mut tx, key_id.as_str(), now, now, self.timezone)
|
||||
.await
|
||||
.map_err(|source| unavailable().with_source(source))?;
|
||||
tx.commit()
|
||||
.await
|
||||
.map_err(|source| unavailable().with_source(source))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn settle_inner(&self, charge: &ClientBudgetCharge) -> Result<(), ClientBudgetError> {
|
||||
pub(super) async fn settle_charges(
|
||||
&self,
|
||||
charges: &[ClientBudgetCharge],
|
||||
) -> Result<(), ClientBudgetError> {
|
||||
let Some(first) = charges.first() else {
|
||||
return Ok(());
|
||||
};
|
||||
let mut tx = self
|
||||
.pool
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|source| ClientBudgetError(Some(source.into())))?;
|
||||
// 与准入统一先锁 Key，再写窗口和费用，串行化同一 Key 的并发结算
|
||||
// 与窗口开启和重置统一先锁 Key，再写窗口和费用
|
||||
let key = sqlx::query_scalar::<_, String>(
|
||||
"select id from client_api_keys where id = $1 for update",
|
||||
)
|
||||
.bind(charge.key_id.as_str())
|
||||
.bind(first.key_id.as_str())
|
||||
.fetch_optional(&mut *tx)
|
||||
.await
|
||||
.map_err(|source| ClientBudgetError(Some(source.into())))?;
|
||||
let Some(key) = key else { return Ok(()) }; // 删除 Key 时也会删除其费用记录
|
||||
settle_in_transaction(&mut tx, &key, charge, self.timezone)
|
||||
settle_in_transaction(&mut tx, &key, charges, self.timezone)
|
||||
.await
|
||||
.map_err(|source| ClientBudgetError(Some(source.into())))?;
|
||||
tx.commit()
|
||||
@@ -274,58 +316,65 @@ impl PgClientBudgetStore {
|
||||
async fn settle_in_transaction(
|
||||
tx: &mut Transaction<'_, Postgres>,
|
||||
key: &str,
|
||||
charge: &ClientBudgetCharge,
|
||||
charges: &[ClientBudgetCharge],
|
||||
timezone: gateway_core::time::DeploymentTimeZone,
|
||||
) -> Result<(), sqlx::Error> {
|
||||
let completed_at = DateTime::<Utc>::from(charge.completed_at);
|
||||
// 仅在请求结束时写入费用；请求 ID 冲突时不重复累计
|
||||
let changed = sqlx::query(
|
||||
let ids = charges
|
||||
.iter()
|
||||
.map(|charge| charge.request_id.as_str())
|
||||
.collect::<Vec<_>>();
|
||||
let amounts = charges
|
||||
.iter()
|
||||
.map(|charge| charge.amount_usd.canonical())
|
||||
.collect::<Vec<_>>();
|
||||
let times = charges
|
||||
.iter()
|
||||
.map(|charge| DateTime::<Utc>::from(charge.completed_at))
|
||||
.collect::<Vec<_>>();
|
||||
// 同一 Key 的小批次共用事务，只累计本次真正插入的事件
|
||||
let inserted = sqlx::query(
|
||||
"insert into client_api_key_charge_events (request_id, client_api_key_id, amount_usd, completed_at)
|
||||
values ($1, $2, $3::text::numeric, $4)
|
||||
on conflict (request_id) do nothing",
|
||||
select request_id, $1, amount::numeric, completed_at
|
||||
from unnest($2::text[], $3::text[], $4::timestamptz[]) as charges(request_id, amount, completed_at)
|
||||
on conflict (request_id) do nothing returning amount_usd::text as amount, completed_at",
|
||||
).bind(key).bind(ids).bind(amounts).bind(times).fetch_all(&mut **tx).await?;
|
||||
let amounts = inserted
|
||||
.iter()
|
||||
.map(|row| row.get::<String, _>("amount"))
|
||||
.collect::<Vec<_>>();
|
||||
let times = inserted
|
||||
.iter()
|
||||
.map(|row| row.get::<DateTime<Utc>, _>("completed_at"))
|
||||
.collect::<Vec<_>>();
|
||||
let Some(latest) = times.iter().max().copied() else {
|
||||
return Ok(());
|
||||
};
|
||||
advance_windows(tx, key, Utc::now(), latest, timezone).await?;
|
||||
sqlx::query(
|
||||
"update client_api_key_budget_windows w set
|
||||
daily_used_usd = daily_used_usd + coalesce((select sum(amount::numeric)
|
||||
from unnest($2::text[], $3::timestamptz[]) as charges(amount, completed_at)
|
||||
where completed_at >= w.daily_start and completed_at < w.daily_end), 0),
|
||||
weekly_used_usd = weekly_used_usd + coalesce((select sum(amount::numeric)
|
||||
from unnest($2::text[], $3::timestamptz[]) as charges(amount, completed_at)
|
||||
where completed_at >= w.weekly_start and completed_at < w.weekly_end), 0)
|
||||
where client_api_key_id = $1",
|
||||
)
|
||||
.bind(charge.request_id.as_str())
|
||||
.bind(key)
|
||||
.bind(charge.amount_usd.canonical())
|
||||
.bind(completed_at)
|
||||
.bind(amounts)
|
||||
.bind(times)
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
if changed == 1 {
|
||||
advance_windows(tx, key, Utc::now(), completed_at, timezone).await?;
|
||||
sqlx::query("update client_api_key_budget_windows set
|
||||
daily_used_usd = daily_used_usd + case when $3 >= daily_start and $3 < daily_end then $2::text::numeric else 0 end,
|
||||
weekly_used_usd = weekly_used_usd + case when $3 >= weekly_start and $3 < weekly_end then $2::text::numeric else 0 end
|
||||
where client_api_key_id = $1")
|
||||
.bind(key).bind(charge.amount_usd.canonical()).bind(completed_at)
|
||||
.execute(&mut **tx).await?;
|
||||
}
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
impl ClientBudgetPort for PgClientBudgetStore {
|
||||
fn admit(&self, key_id: ClientApiKeyId) -> BoxFuture<'_, Result<(), GatewayError>> {
|
||||
Box::pin(async move { self.admit_inner(key_id).await })
|
||||
Box::pin(async move { self.admit_with_pending(key_id, &[]).await })
|
||||
}
|
||||
|
||||
fn settle(&self, charge: ClientBudgetCharge) -> BoxFuture<'_, Result<(), ClientBudgetError>> {
|
||||
Box::pin(async move {
|
||||
let result = self.settle_inner(&charge).await;
|
||||
// 重试缓存的锁异常不能覆盖已捕获的账本失败
|
||||
let mut retry = self.retry.lock().map_err(|_| {
|
||||
result
|
||||
.as_ref()
|
||||
.err()
|
||||
.cloned()
|
||||
.unwrap_or(ClientBudgetError(None))
|
||||
})?;
|
||||
if result.is_err() {
|
||||
retry.insert(charge.request_id.as_str().to_owned(), charge);
|
||||
} else {
|
||||
retry.remove(charge.request_id.as_str());
|
||||
}
|
||||
result
|
||||
})
|
||||
Box::pin(async move { self.settle_charges(std::slice::from_ref(&charge)).await })
|
||||
}
|
||||
}
|
||||
|
||||
@@ -346,7 +395,7 @@ async fn advance_windows(
|
||||
let weekly_end = timezone
|
||||
.days_after(now, 7)
|
||||
.ok_or_else(|| sqlx::Error::Protocol("invalid budget week end".to_owned()))?;
|
||||
// Key 行锁仍串行化准入和结算；未滚动的窗口无需产生新元组与 WAL
|
||||
// 调用方持有 Key 行锁；未滚动的窗口无需产生新元组与 WAL
|
||||
sqlx::query("insert into client_api_key_budget_windows
|
||||
(client_api_key_id, daily_start, daily_end, weekly_start, weekly_end)
|
||||
select $1, $4, $5, $4, $6
|
||||
@@ -426,7 +475,7 @@ pub(super) async fn load_client_key_budgets(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn unavailable() -> GatewayError {
|
||||
pub(super) fn unavailable() -> GatewayError {
|
||||
GatewayError::new(
|
||||
GatewayErrorKind::ProviderInfrastructureUnavailable,
|
||||
"client budget service is temporarily unavailable",
|
||||
|
||||
@@ -4,6 +4,7 @@ mod account_groups;
|
||||
mod admin_security_audit;
|
||||
mod admission_recovery;
|
||||
mod backup;
|
||||
mod client_budget_buffer;
|
||||
mod client_budgets;
|
||||
mod client_keys;
|
||||
mod connection;
|
||||
@@ -27,6 +28,7 @@ pub use account_groups::*;
|
||||
pub use admin_security_audit::*;
|
||||
pub use admission_recovery::*;
|
||||
pub use backup::*;
|
||||
pub use client_budget_buffer::{BufferedClientBudgetStore, ClientBudgetWriter};
|
||||
pub use client_budgets::PgClientBudgetStore;
|
||||
pub use client_keys::*;
|
||||
pub use connection::connect_and_migrate;
|
||||
|
||||
@@ -11,9 +11,10 @@ pub(crate) struct CommandStoreWriters {
|
||||
pub(crate) execution: postgres::ExecutionObservationWriter<postgres::PgExecutionStore>,
|
||||
pub(crate) client_key_usage: postgres::PgClientApiKeyUsageWriter,
|
||||
pub(crate) admission_release: redis::ClientAdmissionReleaseWriter,
|
||||
pub(crate) budget: postgres::ClientBudgetWriter,
|
||||
}
|
||||
|
||||
/// 短生命周期 CLI 只运行数据面必需的三个写泵，不注册恢复、保留或维护 Worker
|
||||
/// 短生命周期 CLI 只运行数据面必需的四个写泵，不注册恢复、保留或维护 Worker
|
||||
pub struct CommandStoreDrain {
|
||||
cancellation: gateway_core::lifecycle::CancellationToken,
|
||||
tasks: Vec<tokio::task::JoinHandle<Result<(), WorkerTaskError>>>,
|
||||
@@ -31,12 +32,14 @@ impl CommandStoreWriters {
|
||||
execution,
|
||||
client_key_usage,
|
||||
admission_release,
|
||||
budget,
|
||||
} = self;
|
||||
let execution_idle = execution.idle();
|
||||
let tasks = vec![
|
||||
spawn_command_writer(execution, cancellation.child_token()),
|
||||
spawn_command_writer(client_key_usage, cancellation.child_token()),
|
||||
spawn_command_writer(admission_release, cancellation.child_token()),
|
||||
spawn_command_writer(budget, cancellation.child_token()),
|
||||
];
|
||||
CommandStoreDrain {
|
||||
cancellation,
|
||||
@@ -101,6 +104,7 @@ pub(crate) fn store_worker_contributions(
|
||||
execution_writer: postgres::ExecutionObservationWriter<postgres::PgExecutionStore>,
|
||||
client_key_usage_writer: postgres::PgClientApiKeyUsageWriter,
|
||||
admission_release_writer: redis::ClientAdmissionReleaseWriter,
|
||||
budget_writer: postgres::ClientBudgetWriter,
|
||||
) -> StoreResult<Vec<WorkerContribution>> {
|
||||
let stale_id = WorkerId::try_new(WorkerKind::StaleModelRequestRecovery, "postgres")
|
||||
.map_err(worker_definition_error)?;
|
||||
@@ -111,10 +115,22 @@ pub(crate) fn store_worker_contributions(
|
||||
.map_err(worker_definition_error)?;
|
||||
let admission_flush_id = WorkerId::try_new(WorkerKind::OpsFlush, "redis_admission")
|
||||
.map_err(worker_definition_error)?;
|
||||
let budget_flush_id = WorkerId::try_new(WorkerKind::OpsFlush, "postgres_client_budget")
|
||||
.map_err(worker_definition_error)?;
|
||||
let ops_flush_restart =
|
||||
DaemonRestartPolicy::try_new(Duration::from_secs(1), Duration::from_secs(60))
|
||||
.map_err(worker_definition_error)?;
|
||||
Ok(vec![
|
||||
WorkerContribution::Registration(
|
||||
WorkerRegistration::try_new(
|
||||
budget_flush_id,
|
||||
WorkerRunnable::Daemon {
|
||||
restart: ops_flush_restart,
|
||||
task: Box::new(budget_writer),
|
||||
},
|
||||
)
|
||||
.map_err(worker_definition_error)?,
|
||||
),
|
||||
WorkerContribution::Registration(scheduled_worker(
|
||||
stale_id,
|
||||
Duration::from_secs(30),
|
||||
|
||||
@@ -0,0 +1,368 @@
|
||||
//! 验证慢结算隔离、待结算额度、幂等重试和后台写入生命周期
|
||||
|
||||
use std::{num::NonZeroUsize, sync::Arc, time::Duration};
|
||||
|
||||
use gateway_admin::{
|
||||
model::client_keys::{
|
||||
ClientKeyBudgetMutationOrigin, ClientKeyBudgetPeriod, ResetClientKeyBudget,
|
||||
},
|
||||
ports::store::ClientKeyStore as _,
|
||||
};
|
||||
use gateway_core::{
|
||||
engine::budget::ClientBudgetPort,
|
||||
lifecycle::CancellationToken,
|
||||
task::{DaemonTask, WorkerDaemonContext},
|
||||
};
|
||||
use gateway_store::postgres::{
|
||||
BufferedClientBudgetStore, ClientBudgetWriter, PgAdminClientKeyStore, PgClientBudgetStore,
|
||||
};
|
||||
|
||||
use super::{
|
||||
TestDatabase,
|
||||
client_budgets::{charge, context, key_id, seed, status},
|
||||
};
|
||||
|
||||
fn buffered(pool: sqlx::PgPool) -> (BufferedClientBudgetStore, Arc<ClientBudgetWriter>) {
|
||||
let (store, writer) = BufferedClientBudgetStore::new(
|
||||
PgClientBudgetStore::new(pool),
|
||||
Arc::new(crate::RecordingDiagnostics::default()),
|
||||
);
|
||||
(store, Arc::new(writer))
|
||||
}
|
||||
|
||||
async fn wait_for_writes(context: &WorkerDaemonContext, expected: u64) {
|
||||
tokio::time::timeout(Duration::from_secs(5), async {
|
||||
loop {
|
||||
let completed = context.completed_cycles();
|
||||
if completed >= expected {
|
||||
break;
|
||||
}
|
||||
context.progress_since(completed).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("background budget writes must complete");
|
||||
}
|
||||
|
||||
async fn drain(writer: &ClientBudgetWriter) {
|
||||
let cancellation = CancellationToken::new();
|
||||
cancellation.cancel();
|
||||
writer
|
||||
.run(WorkerDaemonContext::new(cancellation))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn slow_settlement_does_not_block_receipt_or_concurrent_admission() {
|
||||
let Some(database) = TestDatabase::create("budget_buffer_slow").await else {
|
||||
return;
|
||||
};
|
||||
let direct = PgClientBudgetStore::new(database.pool.clone());
|
||||
for (key, daily) in [("allowed", "10"), ("exhausted", "1")] {
|
||||
seed(&database, key, daily, "0").await;
|
||||
direct.admit(key_id(key)).await.unwrap();
|
||||
}
|
||||
// 控制连接独立于数据面连接池，保证测试只阻塞结算所需行锁
|
||||
let pool = sqlx::postgres::PgPoolOptions::new()
|
||||
.max_connections(4)
|
||||
.connect_with((*database.pool.connect_options()).clone())
|
||||
.await
|
||||
.unwrap();
|
||||
let (store, writer) = buffered(pool.clone());
|
||||
let mut blocker = database.pool.begin().await.unwrap();
|
||||
sqlx::query("select id from client_api_keys for update")
|
||||
.fetch_all(&mut *blocker)
|
||||
.await
|
||||
.unwrap();
|
||||
for key in ["allowed", "exhausted"] {
|
||||
let fee = charge(key, key, "1.25");
|
||||
tokio::time::timeout(Duration::from_millis(250), async {
|
||||
store.settle(fee.clone()).await.unwrap();
|
||||
store.settle(fee).await.unwrap();
|
||||
})
|
||||
.await
|
||||
.expect("completion must only enqueue, including duplicate receipts");
|
||||
}
|
||||
let worker_context = WorkerDaemonContext::new(CancellationToken::new());
|
||||
let running_context = worker_context.clone();
|
||||
let task = tokio::spawn(async move { writer.run(running_context).await });
|
||||
tokio::time::timeout(Duration::from_millis(500), async {
|
||||
for result in
|
||||
futures::future::join_all((0..4).map(|_| store.admit(key_id("allowed")))).await
|
||||
{
|
||||
result.unwrap();
|
||||
}
|
||||
let error = store.admit(key_id("exhausted")).await.unwrap_err();
|
||||
assert_eq!(error.client_error_code(), Some("key_daily_budget_exceeded"));
|
||||
})
|
||||
.await
|
||||
.expect("admission must not wait for settlement row locks");
|
||||
let before: i64 = sqlx::query_scalar("select count(*) from client_api_key_charge_events")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(before, 0);
|
||||
blocker.commit().await.unwrap();
|
||||
wait_for_writes(&worker_context, 2).await;
|
||||
worker_context.cancellation().cancel();
|
||||
task.await.unwrap().unwrap();
|
||||
assert_eq!(
|
||||
status(&database, "allowed")
|
||||
.await
|
||||
.daily_used_usd
|
||||
.canonical(),
|
||||
"1.25"
|
||||
);
|
||||
assert_eq!(
|
||||
status(&database, "exhausted")
|
||||
.await
|
||||
.daily_used_usd
|
||||
.canonical(),
|
||||
"1.25"
|
||||
);
|
||||
pool.close().await;
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pending_cost_respects_committed_ids_and_manual_reset_boundaries() {
|
||||
let Some(database) = TestDatabase::create("budget_buffer_reset").await else {
|
||||
return;
|
||||
};
|
||||
seed(&database, "key", "2", "0").await;
|
||||
let direct = PgClientBudgetStore::new(database.pool.clone());
|
||||
let (store, writer) = buffered(database.pool.clone());
|
||||
store.admit(key_id("key")).await.unwrap();
|
||||
let committed = charge("key", "committed", "1.25");
|
||||
store.settle(committed.clone()).await.unwrap();
|
||||
direct.settle(committed).await.unwrap();
|
||||
// 模拟提交成功、队列尚未移除，金额不能按 2.5 重复计算
|
||||
store.admit(key_id("key")).await.unwrap();
|
||||
store
|
||||
.settle(charge("key", "before-reset", "3"))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
store
|
||||
.admit(key_id("key"))
|
||||
.await
|
||||
.unwrap_err()
|
||||
.client_error_code(),
|
||||
Some("key_daily_budget_exceeded")
|
||||
);
|
||||
PgAdminClientKeyStore::new(database.pool.clone())
|
||||
.reset_client_key_budget(
|
||||
ResetClientKeyBudget {
|
||||
id: key_id("key"),
|
||||
period: ClientKeyBudgetPeriod::All,
|
||||
},
|
||||
ClientKeyBudgetMutationOrigin::Admin,
|
||||
&context(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
store.admit(key_id("key")).await.unwrap();
|
||||
store
|
||||
.settle(charge("key", "after-reset", "2.5"))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
store
|
||||
.admit(key_id("key"))
|
||||
.await
|
||||
.unwrap_err()
|
||||
.client_error_code(),
|
||||
Some("key_daily_budget_exceeded")
|
||||
);
|
||||
drain(&writer).await;
|
||||
assert_eq!(
|
||||
status(&database, "key").await.daily_used_usd.canonical(),
|
||||
"2.5"
|
||||
);
|
||||
assert_eq!(
|
||||
status(&database, "key").await.weekly_used_usd.canonical(),
|
||||
"2.5"
|
||||
);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn failed_write_retains_cost_and_retries_without_another_request() {
|
||||
let Some(database) = TestDatabase::create("budget_buffer_retry").await else {
|
||||
return;
|
||||
};
|
||||
seed(&database, "key", "1", "0").await;
|
||||
let (store, writer) = buffered(database.pool.clone());
|
||||
store.admit(key_id("key")).await.unwrap();
|
||||
sqlx::raw_sql(
|
||||
"create sequence write_attempts;
|
||||
create function reject_budget_update() returns trigger language plpgsql as $$ begin
|
||||
perform nextval('write_attempts');
|
||||
raise exception 'injected settlement failure';
|
||||
end $$;
|
||||
create trigger reject_budget_update before update on client_api_key_budget_windows
|
||||
for each row execute function reject_budget_update();",
|
||||
)
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
for (request, amount) in [("retry-a", "0.25"), ("retry-b", "0.5"), ("retry-c", "0.5")] {
|
||||
store.settle(charge("key", request, amount)).await.unwrap();
|
||||
}
|
||||
writer
|
||||
.run(WorkerDaemonContext::new(CancellationToken::new()))
|
||||
.await
|
||||
.unwrap_err();
|
||||
for result in futures::future::join_all((0..4).map(|_| store.admit(key_id("key")))).await {
|
||||
assert_eq!(
|
||||
result.unwrap_err().client_error_code(),
|
||||
Some("key_daily_budget_exceeded")
|
||||
);
|
||||
}
|
||||
let attempts: i64 = sqlx::query_scalar("select last_value from write_attempts")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(attempts, 1, "admission never retries writes");
|
||||
let rows: i64 = sqlx::query_scalar("select count(*) from client_api_key_charge_events")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(rows, 0, "the entire batch rolls back");
|
||||
assert_eq!(
|
||||
status(&database, "key").await.daily_used_usd.canonical(),
|
||||
"0"
|
||||
);
|
||||
sqlx::query("drop trigger reject_budget_update on client_api_key_budget_windows")
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let worker_context = WorkerDaemonContext::new(CancellationToken::new());
|
||||
let running_context = worker_context.clone();
|
||||
let task = tokio::spawn(async move { writer.run(running_context).await });
|
||||
wait_for_writes(&worker_context, 1).await;
|
||||
worker_context.cancellation().cancel();
|
||||
task.await.unwrap().unwrap();
|
||||
assert_eq!(
|
||||
status(&database, "key").await.daily_used_usd.canonical(),
|
||||
"1.25"
|
||||
);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cancelled_writer_requeues_inflight_cost_for_shutdown_drain() {
|
||||
let Some(database) = TestDatabase::create("budget_buffer_cancel").await else {
|
||||
return;
|
||||
};
|
||||
seed(&database, "key", "0", "0").await;
|
||||
let name = format!("budget_cancel_{}", uuid::Uuid::new_v4().simple());
|
||||
let pool = sqlx::postgres::PgPoolOptions::new()
|
||||
.max_connections(2)
|
||||
.connect_with(
|
||||
(*database.pool.connect_options())
|
||||
.clone()
|
||||
.application_name(&name),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let (store, writer) = buffered(pool.clone());
|
||||
store.admit(key_id("key")).await.unwrap();
|
||||
let mut blocker = database.pool.begin().await.unwrap();
|
||||
sqlx::query("select id from client_api_keys for update")
|
||||
.fetch_all(&mut *blocker)
|
||||
.await
|
||||
.unwrap();
|
||||
store
|
||||
.settle(charge("key", "cancelled", "0.75"))
|
||||
.await
|
||||
.unwrap();
|
||||
let running = writer.clone();
|
||||
let task = tokio::spawn(async move {
|
||||
running
|
||||
.run(WorkerDaemonContext::new(CancellationToken::new()))
|
||||
.await
|
||||
});
|
||||
tokio::time::timeout(Duration::from_secs(2), async {
|
||||
loop {
|
||||
let blocked: bool = sqlx::query_scalar("select exists(select 1 from pg_stat_activity where application_name=$1 and wait_event_type='Lock')")
|
||||
.bind(&name).fetch_one(&database.pool).await.unwrap();
|
||||
if blocked { break; }
|
||||
tokio::time::sleep(Duration::from_millis(5)).await;
|
||||
}
|
||||
}).await.unwrap();
|
||||
task.abort();
|
||||
assert!(task.await.unwrap_err().is_cancelled());
|
||||
blocker.commit().await.unwrap();
|
||||
drain(&writer).await;
|
||||
assert_eq!(
|
||||
status(&database, "key").await.daily_used_usd.canonical(),
|
||||
"0.75"
|
||||
);
|
||||
pool.close().await;
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn full_queue_rejects_new_admission_but_keeps_already_completed_costs() {
|
||||
let Some(database) = TestDatabase::create("budget_buffer_full").await else {
|
||||
return;
|
||||
};
|
||||
for key in ["key", "other"] {
|
||||
seed(&database, key, "0", "0").await;
|
||||
}
|
||||
let (store, writer) = BufferedClientBudgetStore::with_capacity(
|
||||
PgClientBudgetStore::new(database.pool.clone()),
|
||||
NonZeroUsize::new(2).unwrap(),
|
||||
Arc::new(crate::RecordingDiagnostics::default()),
|
||||
);
|
||||
for (i, (key, amount)) in [("key", "1"), ("other", "0.25"), ("key", "2")]
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
{
|
||||
store
|
||||
.settle(charge(key, &format!("fee-{i}"), amount))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
assert_eq!(
|
||||
store
|
||||
.admit(key_id("key"))
|
||||
.await
|
||||
.unwrap_err()
|
||||
.client_error_code(),
|
||||
Some("key_budget_unavailable")
|
||||
);
|
||||
drain(&writer).await;
|
||||
assert_eq!(
|
||||
status(&database, "key").await.daily_used_usd.canonical(),
|
||||
"3"
|
||||
);
|
||||
assert_eq!(
|
||||
status(&database, "other").await.daily_used_usd.canonical(),
|
||||
"0.25"
|
||||
);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn stalled_budget_read_has_a_total_admission_deadline() {
|
||||
let Some(database) = TestDatabase::create("budget_buffer_deadline").await else {
|
||||
return;
|
||||
};
|
||||
seed(&database, "key", "0", "0").await;
|
||||
let (store, _writer) = buffered(database.pool.clone());
|
||||
let mut blocker = database.pool.begin().await.unwrap();
|
||||
sqlx::query("lock table client_api_keys in access exclusive mode")
|
||||
.execute(&mut *blocker)
|
||||
.await
|
||||
.unwrap();
|
||||
let error = tokio::time::timeout(Duration::from_millis(1_500), store.admit(key_id("key")))
|
||||
.await
|
||||
.expect("budget query must have a total deadline")
|
||||
.unwrap_err();
|
||||
assert_eq!(error.client_error_code(), Some("key_budget_unavailable"));
|
||||
blocker.rollback().await.unwrap();
|
||||
database.close().await;
|
||||
}
|
||||
@@ -33,11 +33,11 @@ use gateway_store::postgres::{
|
||||
|
||||
use super::TestDatabase;
|
||||
|
||||
fn key_id(key: &str) -> ClientApiKeyId {
|
||||
pub(super) fn key_id(key: &str) -> ClientApiKeyId {
|
||||
ClientApiKeyId::new(key).unwrap()
|
||||
}
|
||||
|
||||
fn charge(key: &str, request: &str, amount: &str) -> ClientBudgetCharge {
|
||||
pub(super) fn charge(key: &str, request: &str, amount: &str) -> ClientBudgetCharge {
|
||||
ClientBudgetCharge {
|
||||
key_id: key_id(key),
|
||||
request_id: ModelRequestId::new(format!("req_{request}")).unwrap(),
|
||||
@@ -46,7 +46,7 @@ fn charge(key: &str, request: &str, amount: &str) -> ClientBudgetCharge {
|
||||
}
|
||||
}
|
||||
|
||||
async fn seed(database: &TestDatabase, key: &str, daily: &str, weekly: &str) {
|
||||
pub(super) async fn seed(database: &TestDatabase, key: &str, daily: &str, weekly: &str) {
|
||||
sqlx::query(
|
||||
"insert into client_api_keys (id, name, key, daily_limit_usd, weekly_limit_usd, created_at, updated_at)
|
||||
values ($1, $1, $2, $3::text::numeric, $4::text::numeric, now(), now())",
|
||||
@@ -60,7 +60,7 @@ async fn seed(database: &TestDatabase, key: &str, daily: &str, weekly: &str) {
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
async fn status(database: &TestDatabase, key: &str) -> ClientBudgetStatus {
|
||||
pub(super) async fn status(database: &TestDatabase, key: &str) -> ClientBudgetStatus {
|
||||
PgClientApiKeyRepository::new(database.pool.clone())
|
||||
.get_client_api_key(key)
|
||||
.await
|
||||
@@ -69,7 +69,7 @@ async fn status(database: &TestDatabase, key: &str) -> ClientBudgetStatus {
|
||||
.budget
|
||||
}
|
||||
|
||||
fn context() -> MutationContext {
|
||||
pub(super) fn context() -> MutationContext {
|
||||
MutationContext {
|
||||
actor: MutationActor::System,
|
||||
request_id: "budget-test".to_owned(),
|
||||
@@ -687,7 +687,7 @@ async fn zero_cost_and_interrupted_requests_never_block_limited_keys() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn transient_settlement_failure_rolls_back_and_retries_exact_cost_before_admission() {
|
||||
async fn settlement_failure_rolls_back_and_explicit_retry_is_idempotent() {
|
||||
let Some(database) = TestDatabase::create("budgets_retry").await else {
|
||||
return;
|
||||
};
|
||||
@@ -721,6 +721,7 @@ async fn transient_settlement_failure_rolls_back_and_retries_exact_cost_before_a
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
store.settle(charge("key", "retry", "1.25")).await.unwrap();
|
||||
let error = store.admit(key_id("key")).await.unwrap_err();
|
||||
assert_eq!(error.client_error_code(), Some("key_daily_budget_exceeded"));
|
||||
store.settle(charge("key", "retry", "1.25")).await.unwrap();
|
||||
|
||||
@@ -20,6 +20,7 @@ mod analytics_buckets;
|
||||
mod analytics_reset;
|
||||
mod backup;
|
||||
mod backup_export;
|
||||
mod client_budget_buffer;
|
||||
mod client_budgets;
|
||||
mod client_keys;
|
||||
mod connection;
|
||||
|
||||
+4
-3
@@ -1487,14 +1487,15 @@ WebSocket 使用对应错误事件。等待期间不发送上游请求，取消
|
||||
SSE 在 Key 准入完成并取得执行会话前不发送保活帧，此阶段保留 HTTP 错误状态；
|
||||
取得会话后的账号等待受 SSE 保活覆盖，已开始交付的失败沿用流内错误合同
|
||||
|
||||
任一已结算金额达到限额后拒绝新请求，已准入请求可完成并使金额超过阈值。
|
||||
任一周期的已落库金额与当前进程已接管的待结算费用合计达到限额后拒绝新请求，已准入请求可完成并使金额超过阈值。
|
||||
HTTP 返回 `429`，`error.code` 为 `key_daily_budget_exceeded` 或 `key_weekly_budget_exceeded`，
|
||||
并附 `Retry-After`；WebSocket 每次 `response.create` 执行相同检查并返回协议错误事件。
|
||||
只累计上游上报或按用量与模型价格计算出的 USD 费用；无法取得费用的尝试按零累计，
|
||||
保留错误和用量诊断，不产生待核账记录或阻断。内部重试中已经取得的费用仍会累计。
|
||||
预算存储不可用时返回 `503`、`key_budget_unavailable`
|
||||
预算检查超时、存储不可用或待结算队列无法接受新请求时返回 `503`、`key_budget_unavailable`
|
||||
|
||||
自动结算按网关请求 ID 幂等执行。账本独立于使用统计日志，记录保留至删除 Key，
|
||||
自动结算按网关请求 ID 幂等执行，响应交付不等待费用落库，额度查询中的已用金额在后台写入后可见。
|
||||
账本独立于使用统计日志，记录保留至删除 Key，
|
||||
不受 `usageRetentionDays` 影响
|
||||
|
||||
## 8. 运行设置
|
||||
|
||||
+12
-6
@@ -782,20 +782,26 @@ Admin 的手动重置复用同一账本与 Key 行锁，在一个事务中清零
|
||||
已打开窗口不随时区切换清零或改写 UTC 边界，续接起点不得早于旧终点与人工重置边界
|
||||
- 金额优先使用 Provider 上报的 USD，否则按现有模型价格估算；订阅账号的估算费用不代表上游订阅账单。
|
||||
费用归属请求完成时间，跨日请求计入完成日，延迟写入仍保留原完成时间
|
||||
- 准入检查已结算金额，不预占未来费用。达到任一金额阈值后拒绝新请求，已准入请求仍可完成并使总额超过阈值
|
||||
- 准入检查已落库金额与当前进程已接管的待结算费用，不预占未来费用。达到任一金额阈值后拒绝新请求，已准入请求仍可完成并使总额超过阈值
|
||||
- SSE 持有并发名额直到终态；WebSocket 每个 `response.create` 独立准入，空闲连接不占名额。
|
||||
Core 在每次请求开始时重新鉴权并冻结当前策略，既有连接也应用已发布的限额与授权范围变更。
|
||||
换号和内部重试复用同一名额，Key 与账号并发上限同时生效；预算拒绝或启动失败释放名额
|
||||
|
||||
完成、失败、取消和断连使用同一结算路径，在结束时以网关请求 ID 幂等累计已取得费用。
|
||||
执行会话持有结算与准入释放的同一个清理 future；协议等待被取消后，后续驱动或宿主 detached finalizer
|
||||
继续原清理，不重复结算。只有执行终态、结算与释放均已返回，协议层才可视为 finalized 并放弃清理责任。
|
||||
完成、失败、取消和断连使用同一结算路径，在结束时以网关请求 ID 幂等登记已取得费用。
|
||||
执行会话持有费用登记与准入释放的同一个清理 future；协议等待被取消后，后续驱动或宿主 detached finalizer
|
||||
继续原清理。Store 接管精确费用后立即返回，响应交付及并发释放不等待 PostgreSQL 落库；
|
||||
只有执行终态、费用登记与释放均已返回，协议层才可视为 finalized 并放弃清理责任。
|
||||
缺少用量或价格、无法取得 USD 费用的尝试按零累计，不创建待核账记录，也不会阻断 Key；
|
||||
内部重试中已经取得的费用仍须累计。发送状态只决定重放是否安全，不能据此推断费用。
|
||||
请求日志继续保留真实错误、用量与费用来源，账本按零累计不表示上游实际免费
|
||||
|
||||
结算写入失败时，进程内保留精确费用，在同一 Key 下次请求前重试；进程退出后无法恢复的费用不会形成欠账。
|
||||
PostgreSQL 不可用时拒绝所有新的计费请求，Redis 继续管理并发/RPM 租约。
|
||||
受监督的费用 writer 按 Key 合并小批次事务，负责后台落库、退避重试及关闭时有界排空，
|
||||
写入取消或结果未知时保留费用并按请求 ID 重试。
|
||||
写入失败后缩小对应 Key 的批次，快速写入恢复后逐步还原批量，避免慢库反复重试无法完成的大批次。
|
||||
普通准入通过同一数据库快照读取预算窗口与已落库请求，合并待结算费用并避免重复计算，不执行补账或等待结算行锁；
|
||||
仅首次使用或窗口到期时开启窗口，仍与重置和结算使用相同锁序。预算准入总等待上限为一秒。
|
||||
待结算费用达到 4096 笔后拒绝新准入，已准入请求的费用继续接收；缓存上限因此还包含在途请求的收尾费用。
|
||||
进程退出后无法恢复的费用不会形成欠账。PostgreSQL 不可用或费用队列不可用时拒绝新的计费请求，Redis 继续管理并发/RPM 租约。
|
||||
账本独立于可丢弃的请求观测日志，日志清理不重置金额；费用事件保留至删除 Key，已有日志不会回填为账本费用。
|
||||
字段与错误合同见 [Client Key API](api.md#7-client-key)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user