From 535f7593e644ba4efe3428c71bf09a6b371592db Mon Sep 17 00:00:00 2001 From: CMGS Date: Sat, 17 Oct 2026 18:18:40 +0800 Subject: [PATCH 01/22] billing: spool ledger rows the store keeps refusing and replay them A Postgres outage longer than the retry budget (about 13 s) dropped every queued billing row. With storage.ledger_spool_path set, a batch goes to a local JSON-lines file after three attempts and replays every 30 s; request_id uniqueness keeps a replay idempotent. --- crates/config/src/lib.rs | 3 + crates/state/src/admission.rs | 236 ++++++++++++++++++++++++++++++++-- crates/state/src/lib.rs | 3 +- crates/state/src/store.rs | 2 +- docs/configuration.md | 1 + docs/deployment.md | 1 + docs/observability.md | 5 + 7 files changed, 236 insertions(+), 15 deletions(-) diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index 47642e6..d7d2808 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -580,6 +580,9 @@ pub struct StorageConf { /// Prune rolled billing records beyond this many; 0 = unlimited. #[serde(default)] pub ledger_max_rows: u64, + /// Local file holding ledger rows the store refuses until they replay; empty drops them after the retries. + #[serde(default)] + pub ledger_spool_path: String, /// Postgres pool size for the store and key-table pools; 0 = their built-in defaults. #[serde(default)] pub postgres_max_connections: u32, diff --git a/crates/state/src/admission.rs b/crates/state/src/admission.rs index 1a5d7b4..b588fde 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -4,6 +4,9 @@ //! two admission paths cannot drift. Denials carry the user-facing message; //! callers wrap it in their own wire shape. +use std::io::Write; +use std::os::unix::fs::FileExt; +use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; @@ -20,6 +23,9 @@ const LEDGER_RETRY_INITIAL: Duration = Duration::from_millis(100); const LEDGER_RETRY_MAX: Duration = Duration::from_secs(30); // a store that refuses writes must drain the queue rather than wedge it const LEDGER_RETRY_ATTEMPTS: usize = 8; +// with a spool, rows leave the queue sooner, so an outage does not back up into requests +const LEDGER_SPOOL_AFTER_ATTEMPTS: usize = 3; +const LEDGER_REPLAY_INTERVAL: Duration = Duration::from_secs(30); // a month's counter must outlive the next month's rollover read const MONTH_COUNTER_TTL: Duration = Duration::from_secs(62 * 86_400); @@ -77,10 +83,23 @@ impl BillingLedger { } } - pub(crate) fn repairing(store: Arc) -> Self { + /// A batching writer; with `spool_path`, rows the store keeps refusing wait there instead of being dropped. + pub(crate) fn repairing(store: Arc, spool_path: &str) -> Self { let (queue, mut pending) = mpsc::channel::(LEDGER_QUEUE_CAPACITY); let deferred = store.defers_ledger_writes(); let worker_store = store.clone(); + let spool = (!spool_path.is_empty()).then(|| Arc::new(LedgerSpool::new(spool_path))); + if let Some(spool) = spool.clone() { + let replay_store = store.clone(); + // detached: replays the spool for the life of the process + tokio::spawn(async move { + let mut tick = tokio::time::interval(LEDGER_REPLAY_INTERVAL); + loop { + tick.tick().await; + spool.replay(&replay_store).await; + } + }); + } // the bounded worker owns accepted rows through caller cancellation tokio::spawn(async move { let mut batch = Vec::with_capacity(LEDGER_BATCH_MAX); @@ -96,7 +115,7 @@ impl BillingLedger { next.take(&mut batch, &mut row_acks, &mut ack); } let unacked = (batch.len() - row_acks.len()) as u64; - let committed = Self::commit(&worker_store, &mut batch).await; + let committed = Self::commit(&worker_store, spool.as_deref(), &mut batch).await; if !committed { dropped += unacked; } @@ -175,20 +194,37 @@ impl BillingLedger { rx.await.unwrap_or(false) } - async fn commit(store: &Arc, batch: &mut Vec) -> bool { + async fn commit( + store: &Arc, + spool: Option<&LedgerSpool>, + batch: &mut Vec, + ) -> bool { if batch.is_empty() { return true; } + let attempts = if spool.is_some() { + LEDGER_SPOOL_AFTER_ATTEMPTS + } else { + LEDGER_RETRY_ATTEMPTS + }; let mut delay = LEDGER_RETRY_INITIAL; let mut committed = false; - for attempt in 1..=LEDGER_RETRY_ATTEMPTS { + for attempt in 1..=attempts { let Err(e) = store.ledger_add_batch(batch).await else { committed = true; break; }; metrics::counter!("gateway_ledger_write_failures_total").increment(batch.len() as u64); - if attempt == LEDGER_RETRY_ATTEMPTS { - tracing::error!(error = %e, rows = batch.len(), "billing ledger retries exhausted; rows dropped"); + if attempt == attempts { + committed = match spool { + Some(spool) => spool.append(batch).await, + None => false, + }; + if committed { + tracing::warn!(error = %e, rows = batch.len(), "billing ledger write failed; rows spooled for replay"); + } else { + tracing::error!(error = %e, rows = batch.len(), "billing ledger retries exhausted; rows dropped"); + } break; } tracing::error!(error = %e, rows = batch.len(), "billing ledger write failed; retrying"); @@ -200,6 +236,111 @@ impl BillingLedger { } } +/// Ledger rows the store refused, appended to a local JSON-lines file and replayed in order. +#[derive(Debug)] +struct LedgerSpool { + path: PathBuf, + replaying: PathBuf, + lock: tokio::sync::Mutex<()>, +} + +impl LedgerSpool { + fn new(path: &str) -> Self { + let path = PathBuf::from(path); + let replaying = path.with_extension("replaying"); + Self { + path, + replaying, + lock: tokio::sync::Mutex::new(()), + } + } + + async fn append(&self, rows: &[BillingRecord]) -> bool { + let mut lines = Vec::new(); + for r in rows { + if let Err(e) = serde_json::to_writer(&mut lines, r) { + tracing::error!(error = %e, request_id = r.request_id, "spool billing row"); + return false; + } + lines.push(b'\n'); + } + let path = self.path.clone(); + let _held = self.lock.lock().await; + let written = blocking(move || { + let mut f = std::fs::OpenOptions::new() + .create(true) + .read(true) + .append(true) + .open(&path)?; + let len = f.metadata()?.len(); + let mut last = *b"\n"; + if len > 0 { + f.read_exact_at(&mut last, len - 1)?; + } + // a crash mid-append leaves a torn row; the next row starts on its own line + if last[0] != b'\n' { + f.write_all(b"\n")?; + } + f.write_all(&lines)?; + f.sync_data() + }) + .await; + if let Err(e) = &written { + tracing::error!(error = %e, path = %self.path.display(), "billing spool write failed"); + } + written.is_ok() + } + + /// `request_id` uniqueness makes a re-insert of an already accepted batch a no-op. + async fn replay(&self, store: &Arc) { + let (path, replaying) = (self.path.clone(), self.replaying.clone()); + let read = { + let _held = self.lock.lock().await; + blocking(move || { + if !replaying.exists() { + if !path.exists() { + return Ok(None); + } + std::fs::rename(&path, &replaying)?; + } + std::fs::read(&replaying).map(Some) + }) + .await + }; + let bytes = match read { + Ok(Some(bytes)) => bytes, + Ok(None) => return, + Err(e) => { + tracing::error!(error = %e, path = %self.replaying.display(), "billing spool read failed"); + return; + } + }; + let rows: Vec = bytes + .split(|&b| b == b'\n') + .filter(|line| !line.is_empty()) + .filter_map(|line| { + serde_json::from_str(&String::from_utf8_lossy(line)) + .inspect_err( + |e| tracing::error!(error = %e, "skipping a torn billing spool line"), + ) + .ok() + }) + .collect(); + for chunk in rows.chunks(LEDGER_BATCH_MAX) { + if let Err(e) = store.ledger_add_batch(chunk).await { + tracing::warn!(error = %e, rows = rows.len(), "billing spool replay deferred"); + return; + } + } + let replaying = self.replaying.clone(); + if let Err(e) = blocking(move || std::fs::remove_file(replaying)).await { + tracing::error!(error = %e, path = %self.replaying.display(), "billing spool remove failed"); + return; + } + tracing::info!(rows = rows.len(), "billing spool replayed"); + } +} + /// One budget: its window, governance counter and cap. struct Budget { scope: BudgetScope, @@ -682,6 +823,16 @@ fn month_prefix((y, m): (i64, u32)) -> String { format!("m:{y}{m:02}:") } +async fn blocking(f: F) -> std::io::Result +where + T: Send + 'static, + F: FnOnce() -> std::io::Result + Send + 'static, +{ + tokio::task::spawn_blocking(f) + .await + .unwrap_or_else(|e| Err(std::io::Error::other(e))) +} + #[cfg(test)] mod tests { use gw_config::access_key_id; @@ -744,7 +895,7 @@ mod tests { async fn billing_ledger_repairs_failed_writes() { let store = Arc::new(crate::MemoryStore::default()); store.fail_next_ledger_writes(2); - let ledger = BillingLedger::repairing(store.clone()); + let ledger = BillingLedger::repairing(store.clone(), ""); let record = record("req-repair"); ledger.write(record.clone()).await; let (count, rows) = tokio::time::timeout(Duration::from_secs(2), async { @@ -767,7 +918,7 @@ mod tests { let store = Arc::new(crate::MemoryStore::default()); let ledger = BillingLedger { deferred: true, - ..BillingLedger::repairing(store.clone()) + ..BillingLedger::repairing(store.clone(), "") }; let mut row = record("req-attributed"); row.user_id = "user-42".into(); @@ -785,7 +936,7 @@ mod tests { store.fail_next_ledger_writes(LEDGER_RETRY_ATTEMPTS); let ledger = BillingLedger { deferred: true, - ..BillingLedger::repairing(store.clone()) + ..BillingLedger::repairing(store.clone(), "") }; let mut row = record("req-fallback"); row.user_id = "user-42".into(); @@ -801,7 +952,7 @@ mod tests { async fn canceled_attributed_write_is_reported_as_dropped() { let store = Arc::new(crate::MemoryStore::default()); store.fail_next_ledger_writes(LEDGER_RETRY_ATTEMPTS); - let ledger = BillingLedger::repairing(store); + let ledger = BillingLedger::repairing(store, ""); let (tx, rx) = oneshot::channel(); drop(rx); @@ -820,7 +971,7 @@ mod tests { async fn billing_ledger_backpressures_at_capacity_then_repairs_every_row() { let store = Arc::new(crate::MemoryStore::default()); store.fail_next_ledger_writes(usize::MAX); - let ledger = BillingLedger::repairing(store.clone()); + let ledger = BillingLedger::repairing(store.clone(), ""); let queue = ledger.queue.as_ref().unwrap(); ledger.write(record("req-0")).await; @@ -868,7 +1019,7 @@ mod tests { async fn flush_returns_only_after_queued_rows_commit() { let store = Arc::new(crate::MemoryStore::default()); store.fail_next_ledger_writes(1); - let ledger = BillingLedger::repairing(store.clone()); + let ledger = BillingLedger::repairing(store.clone(), ""); ledger.write(record("req-flush")).await; ledger.flush().await; let (count, rows) = store.ledger_snapshot(usize::MAX).await.unwrap(); @@ -880,7 +1031,7 @@ mod tests { async fn billing_ledger_drains_a_store_that_never_accepts() { let store = Arc::new(crate::MemoryStore::default()); store.fail_next_ledger_writes(usize::MAX); - let ledger = BillingLedger::repairing(store.clone()); + let ledger = BillingLedger::repairing(store.clone(), ""); tokio::time::timeout(Duration::from_secs(3_600), async { for i in 0..=LEDGER_QUEUE_CAPACITY { @@ -905,6 +1056,65 @@ mod tests { ); } + #[tokio::test] + async fn rows_the_store_keeps_refusing_are_spooled_not_dropped() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("ledger-spool.jsonl"); + let store = Arc::new(crate::MemoryStore::default()); + store.fail_next_ledger_writes(usize::MAX); + let ledger = BillingLedger::repairing(store.clone(), path.to_str().unwrap()); + + ledger.write(record("req-spooled")).await; + + assert_eq!( + ledger.flush().await, + 0, + "a spooled row is not a dropped row" + ); + let spooled = std::fs::read_to_string(&path).unwrap(); + assert_eq!(spooled.lines().count(), 1); + assert!(spooled.contains("req-spooled")); + } + + #[tokio::test] + async fn spool_replays_into_the_store_once_it_accepts_and_retries_otherwise() { + let dir = tempfile::tempdir().unwrap(); + let spool = LedgerSpool::new(dir.path().join("ledger-spool.jsonl").to_str().unwrap()); + let store = Arc::new(crate::MemoryStore::default()); + let as_store: Arc = store.clone(); + assert!(spool.append(&[record("req-a"), record("req-b")]).await); + + store.fail_next_ledger_writes(1); + spool.replay(&as_store).await; + assert_eq!(store.ledger_snapshot(usize::MAX).await.unwrap().0, 0); + assert!(spool.replaying.exists(), "a refused replay keeps its rows"); + assert!(spool.append(&[record("req-c")]).await); + + spool.replay(&as_store).await; + spool.replay(&as_store).await; + let (count, rows) = store.ledger_snapshot(usize::MAX).await.unwrap(); + let mut ids: Vec<_> = rows.iter().map(|r| r.request_id.as_str()).collect(); + ids.sort_unstable(); + assert_eq!((count, ids), (3, vec!["req-a", "req-b", "req-c"])); + assert!(!spool.replaying.exists() && !spool.path.exists()); + } + + #[tokio::test] + async fn a_torn_spool_row_never_swallows_the_next_one() { + let dir = tempfile::tempdir().unwrap(); + let spool = LedgerSpool::new(dir.path().join("ledger-spool.jsonl").to_str().unwrap()); + std::fs::write(&spool.path, b"{\"ak\":\"torn").unwrap(); + assert!(spool.append(&[record("req-after-crash")]).await); + let store = Arc::new(crate::MemoryStore::default()); + let as_store: Arc = store.clone(); + + spool.replay(&as_store).await; + + let (count, rows) = store.ledger_snapshot(usize::MAX).await.unwrap(); + assert_eq!(count, 1); + assert_eq!(rows[0].request_id, "req-after-crash"); + } + #[test] fn civil_month_follows_the_gregorian_calendar() { assert_eq!(civil_month(0), (1970, 1)); diff --git a/crates/state/src/lib.rs b/crates/state/src/lib.rs index 53d6cdf..8081b6e 100644 --- a/crates/state/src/lib.rs +++ b/crates/state/src/lib.rs @@ -865,7 +865,8 @@ impl GatewayState { } } } - state.billing = admission::BillingLedger::repairing(state.store.clone()); + state.billing = + admission::BillingLedger::repairing(state.store.clone(), &st.ledger_spool_path); Ok(state) } diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index 908584c..f62c5d9 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -60,7 +60,7 @@ const BATCH_ITEM_CHUNK: usize = 13_000; const PG_MAX_CONNECTIONS: u32 = 10; /// One billing entry (recorded locally only; no reporting upstream). -#[derive(Debug, Clone, serde::Serialize)] +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] pub struct BillingRecord { pub ak: String, pub product: String, diff --git a/docs/configuration.md b/docs/configuration.md index 28201ea..bdd4a17 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -34,6 +34,7 @@ storage: shared_cache: false # also share the response cache in Redis (needs redis_url) ledger_max_rows: 100000 # prune oldest billing rows past the cap; 0 = unlimited postgres_max_connections: 10 # store and key-store connection-pool size; 0 = the built-in defaults (10 store, 5 key store) + ledger_spool_path: /var/lib/gw/ledger-spool.jsonl # rows the store refuses wait here for replay; empty = dropped after the retries ``` The billing ledger, uploaded files, and batch jobs live here. In-memory by diff --git a/docs/deployment.md b/docs/deployment.md index 273c479..b2aae4f 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -70,6 +70,7 @@ storage: redis_url: "redis://redis:6379" # shared rate limits + quotas + account health ledger_max_rows: 1000000 # prune oldest billing rows past the cap postgres_max_connections: 10 # store and key-store connection-pool size + ledger_spool_path: /var/lib/gw/ledger-spool.jsonl # keep billing rows through a store outage # sqlite_path: /var/lib/gw/store.db # single-node alternative to postgres_url ``` diff --git a/docs/observability.md b/docs/observability.md index 1a7f04b..5913082 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -90,6 +90,11 @@ when counts came from an aborted stream rather than a vendor usage payload). Writes are idempotent by `request_id`; a transient store failure is held in a bounded repair queue and retried with backoff — eight attempts, about 13 s — before the batch is dropped and counted in `gateway_ledger_write_failures_total`. +With `storage.ledger_spool_path` set, the batch instead goes to that local +JSON-lines file after three attempts (about 0.3 s), and a replay every 30 s +inserts it once the store accepts writes again (`billing spool replayed` in the +log). A row replayed more than 20 minutes after it was billed is in the ledger +but not in the usage rollup that `GET /admin/usage/users` reads. A full queue applies backpressure to settlement until the repair worker makes room, so overload alone never discards an accepted row. Per-user usage additionally rolls into durable minute buckets every minute, so From 026c3fe37f247565aa980d1eba198ac591340bdd Mon Sep 17 00:00:00 2001 From: CMGS Date: Sat, 17 Oct 2026 18:21:50 +0800 Subject: [PATCH 02/22] budgets: hold the call's estimate against the daily budgets at admission The daily cost and token budgets were check-then-consume, so every in-flight call passed a stale check and a burst could overshoot a user's daily cap many times over. Admission now holds the estimate (the token estimate, and its prompt and max_tokens sides at the tenant price) and admits while spent plus held is under the cap; billing settles the hold and a failed call returns it. --- crates/dag/src/context.rs | 5 +- crates/dag/src/nodes.rs | 60 +++++++--- crates/handler/src/lib.rs | 18 +-- crates/state/src/admission.rs | 210 ++++++++++++++++++++++++++------- crates/state/src/governance.rs | 26 ++-- crates/state/src/lib.rs | 8 +- crates/views/src/lib.rs | 17 ++- docs/governance.md | 24 ++-- 8 files changed, 273 insertions(+), 95 deletions(-) diff --git a/crates/dag/src/context.rs b/crates/dag/src/context.rs index 1f5cb66..9a10e0a 100644 --- a/crates/dag/src/context.rs +++ b/crates/dag/src/context.rs @@ -8,7 +8,7 @@ use std::sync::Arc; use gw_config::GatewayConfig; use gw_engines::{EngineOutcome, SharedTransport}; use gw_models::{GResult, GatewayError, GatewayRequest, ModelParamV2}; -use gw_state::admission::TpmReserve; +use gw_state::admission::{BudgetHold, TpmReserve}; use gw_state::{AkInfo, GatewayState, UserBudget}; pub struct DagContext { @@ -40,6 +40,8 @@ pub struct DagContext { pub tpm_reserved: Option, /// The user's budget override resolved at admission, charged at settlement. pub user_budget: Option, + /// Day-budget amounts held at admission (same lifecycle as `quota_reserved`). + pub budget_hold: BudgetHold, /// Outbound DLP buffered this stream; billing waits for the view's delivery result. pub billing_deferred: bool, } @@ -69,6 +71,7 @@ impl DagContext { quota_at: 0, tpm_reserved: None, user_budget: None, + budget_hold: BudgetHold::default(), billing_deferred: false, } } diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index d980a9f..545a03c 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -241,6 +241,11 @@ impl Write for PromptBytes { /// Admission estimate: ~bytes/4 prompt plus the requested max_tokens, saturating and capped. fn reserve_estimate(req: &gw_models::GatewayRequest) -> i64 { + let (prompt, max_out) = reserve_parts(req); + prompt.saturating_add(max_out).min(MAX_RESERVE) +} + +fn reserve_parts(req: &gw_models::GatewayRequest) -> (i64, i64) { let typed = req.model_param_v2.as_ref().and_then(|p| p.typed.as_ref()); let prompt = if let Some(gw_models::TypedParams::Decisions(p)) = typed { let mut bytes = PromptBytes(0); @@ -256,9 +261,7 @@ fn reserve_estimate(req: &gw_models::GatewayRequest) -> i64 { }) .unwrap_or(DEFAULT_COMPLETION_RESERVE) .clamp(0, MAX_RESERVE); - ((prompt as i64 / 4).max(1)) - .saturating_add(max_out) - .min(MAX_RESERVE) + ((prompt as i64 / 4).clamp(1, MAX_RESERVE), max_out) } /// preprocess/quota_check: reserve the estimate against the AK daily quota; billing settles it. @@ -418,7 +421,7 @@ impl DagNode for AkTpmLimit { } } -/// model_access/budget: the tenant's token and cost budgets (soft caps) for the key and user. +/// model_access/budget: holds the call's estimate against the tenant's daily budgets for the key and user. pub struct BudgetGate; #[async_trait::async_trait] @@ -427,10 +430,31 @@ impl DagNode for BudgetGate { "budget" } async fn execute(&self, ctx: &mut DagContext) -> GResult<()> { - let over = admission::check_budgets(&ctx.state, &ctx.cfg, &ctx.ak, ctx.effective_user_id()) - .await? - .map_err(quota_denied)?; + let (prompt, max_out) = reserve_parts(&ctx.request); + let served = served_model(ctx.request.model_param_v2.as_ref()); + let est = admission::BudgetSpend { + tokens: ctx + .quota_reserved + .filter(|&reserved| reserved > 0) + .unwrap_or_else(|| prompt.saturating_add(max_out).min(MAX_RESERVE)), + cost_micros: gw_models::cost_micros( + prompt, + max_out, + ctx.cfg.prices_for_tenant(&ctx.ak.tenant, served), + ), + }; + let (over, hold) = admission::reserve_budgets( + &ctx.state, + &ctx.cfg, + &ctx.ak, + ctx.effective_user_id(), + est, + ctx.quota_at, + ) + .await? + .map_err(quota_denied)?; ctx.user_budget = Some(over); + ctx.budget_hold = hold; Ok(()) } } @@ -757,15 +781,15 @@ pub async fn settle_deferred_stream(ctx: &mut DagContext, delivery: StreamDelive bill_aborted_stream(ctx, Some(completion)).await } StreamDelivery::None => { - ctx.state - .governance - .refund_reserves( - &ctx.ak.ak_id, - ctx.quota_reserved.take().unwrap_or(0), - ctx.tpm_reserved.take(), - ctx.quota_at, - ) - .await; + let gov = ctx.state.governance.as_ref(); + gov.refund_reserves( + &ctx.ak.ak_id, + ctx.quota_reserved.take().unwrap_or(0), + ctx.tpm_reserved.take(), + ctx.quota_at, + ) + .await; + std::mem::take(&mut ctx.budget_hold).refund(gov).await; if let Some(outcome) = ctx.outcome.as_mut() { outcome.response.prompt_tokens = 0; blank_stream_response(&mut outcome.response); @@ -848,10 +872,11 @@ async fn bill(ctx: &mut DagContext, mut tokens: BillTokens, estimated: bool) -> served_conf.and_then(|m| m.batch_discount).unwrap_or(1.0) }; // locals first: the whole-ctx borrow `effective_user_id` needs can't overlap these takes - let (reserved, tpm_reserved, model_quota_key) = ( + let (reserved, tpm_reserved, model_quota_key, budget_hold) = ( ctx.quota_reserved.take().unwrap_or(0), ctx.tpm_reserved.take(), ctx.model_quota_key.take(), + std::mem::take(&mut ctx.budget_hold), ); let settled = admission::settle_and_bill( ctx.state.as_ref(), @@ -885,6 +910,7 @@ async fn bill(ctx: &mut DagContext, mut tokens: BillTokens, estimated: bool) -> model_quota_key, budget_ak: &ctx.ak, user_budget: ctx.user_budget, + budget_hold, }, ) .await; diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index af0c973..7308d78 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -270,15 +270,15 @@ impl OnlineHandler { break; }; // a failed pipeline refunds its reservations whole, on the reserve's day bucket - ctx.state - .governance - .refund_reserves( - &ctx.ak.ak_id, - ctx.quota_reserved.take().unwrap_or(0), - ctx.tpm_reserved.take(), - ctx.quota_at, - ) - .await; + let gov = ctx.state.governance.as_ref(); + gov.refund_reserves( + &ctx.ak.ak_id, + ctx.quota_reserved.take().unwrap_or(0), + ctx.tpm_reserved.take(), + ctx.quota_at, + ) + .await; + std::mem::take(&mut ctx.budget_hold).refund(gov).await; // signed thinking replays only against the model that produced it if is_upstream_fault(&e) && !ctx.request.replays_reasoning_output() diff --git a/crates/state/src/admission.rs b/crates/state/src/admission.rs index b588fde..1b760f8 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -44,6 +44,46 @@ pub struct TpmReserve { pub window: Option, } +/// Tokens and cost a call charges the budgets: the admission estimate, or the settled actuals. +#[derive(Debug, Clone, Copy, Default)] +pub struct BudgetSpend { + pub tokens: i64, + pub cost_micros: i64, +} + +impl BudgetSpend { + fn of(self, scope: BudgetScope) -> i64 { + if scope.charges_cost() { + self.cost_micros + } else { + self.tokens + } + } +} + +/// Day-budget amounts held at admission, so concurrent calls count before they settle. +#[derive(Debug, Default)] +pub struct BudgetHold { + held: Vec<(String, i64)>, + at: i64, +} + +impl BudgetHold { + /// Return every held amount, for a call that never reached billing. + pub async fn refund(self, gov: &dyn Governance) { + let at = self.at; + futures::future::join_all(self.held.into_iter().map(|(key, amount)| async move { + gov.quota_settle(&key, -amount, at).await; + })) + .await; + } + + fn take(&mut self, key: &str) -> Option { + let i = self.held.iter().position(|(k, _)| k == key)?; + Some(self.held.swap_remove(i).1) + } +} + /// One settled call: identity + reserves to close. pub struct SettleInput<'a> { pub billing: BillingInput<'a>, @@ -59,6 +99,8 @@ pub struct SettleInput<'a> { /// The key whose budgets the call charges, and the user override admission resolved. pub budget_ak: &'a AkInfo, pub user_budget: Option, + /// Day-budget amounts admission held; the settle charges the difference. + pub budget_hold: BudgetHold, } /// What a settled request cost, for the decision trail. @@ -441,34 +483,52 @@ pub async fn flush_billing(state: &GatewayState) { } } -/// Daily and monthly budgets (soft caps: check-then-consume, so concurrent -/// turns can overshoot by one): admit while every configured scope is under its cap. -/// Admitted yields the user's override for settlement; the outer error is a failed lookup. -pub async fn check_budgets( +/// Hold `est` against the day budgets, admitting while spent plus held is under each cap; month budgets and a zero `est` only check. +pub async fn reserve_budgets( state: &GatewayState, cfg: &GatewayConfig, ak: &AkInfo, user: &str, -) -> GResult> { + est: BudgetSpend, + at: i64, +) -> GResult> { let over = user_budget(state, ak, user).await?; let gov = state.governance.as_ref(); let budgets = budgets(gov, cfg, ak, user, over).await; - let under = futures::future::join_all(budgets.iter().map(|b| async move { + let admitted = futures::future::join_all(budgets.iter().map(|b| async move { + let amount = est.of(b.scope); match b.window { - Window::Day => gov.quota_check(&b.key, b.limit).await, - Window::Month => gov.counter_get(&b.key).await < b.limit, + Window::Day if amount > 0 => gov.quota_reserve(&b.key, amount, b.limit, at).await, + Window::Day => gov.quota_check(&b.key, b.limit).await.then_some(0), + Window::Month => (gov.counter_get(&b.key).await < b.limit).then_some(0), } })) .await; - if let Some((b, _)) = budgets.iter().zip(under).find(|(_, under)| !under) { - return Ok(Err(format!( - "{} {} budget exhausted for {}", - b.window.label(), - b.scope.unit(), - b.scope.subject(ak, user) - ))); + let mut hold = BudgetHold { + held: Vec::new(), + at, + }; + let mut denied = None; + for (b, admitted) in budgets.into_iter().zip(admitted) { + match admitted { + Some(held) if held > 0 => hold.held.push((b.key, held)), + Some(_) => {} + None if denied.is_none() => { + denied = Some(format!( + "{} {} budget exhausted for {}", + b.window.label(), + b.scope.unit(), + b.scope.subject(ak, user) + )); + } + None => {} + } } - Ok(Ok(over)) + let Some(denied) = denied else { + return Ok(Ok((over, hold))); + }; + hold.refund(gov).await; + Ok(Err(denied)) } pub fn model_quota_key(ak_id: &str, model: &str) -> String { @@ -647,8 +707,11 @@ pub async fn settle_and_bill( s.budget_ak, s.billing.user_id, s.user_budget, - settled.total_tokens, - settled.cost_micros, + s.budget_hold, + BudgetSpend { + tokens: settled.total_tokens, + cost_micros: settled.cost_micros, + }, ); let write_ledger = state.billing.write(record); tokio::join!( @@ -667,18 +730,20 @@ pub fn month_prefixes() -> [String; 2] { [month_prefix(month), month_prefix(previous_month(month))] } -/// Accrue usage under `over`, the override admission resolved (`None` resolves it now); -/// a scope reaching its cap raises a `budget_exhausted` alert (deduped by the bus). +/// Accrue usage under `over`, the override admission resolved (`None` resolves it now), settling +/// what `hold` held; a scope reaching its cap raises a `budget_exhausted` alert (deduped by the bus). async fn consume_budgets( state: &GatewayState, cfg: &GatewayConfig, ak: &AkInfo, user: &str, over: Option, - tokens: i64, - cost_micros: i64, + mut hold: BudgetHold, + spent: BudgetSpend, ) { - if tokens <= 0 && cost_micros <= 0 { + let gov = state.governance.as_ref(); + if spent.tokens <= 0 && spent.cost_micros <= 0 { + hold.refund(gov).await; return; } let over = match over { @@ -688,23 +753,28 @@ async fn consume_budgets( UserBudget::default() }), }; - let gov = state.governance.as_ref(); let budgets = budgets(gov, cfg, ak, user, over).await; - let used = futures::future::join_all(budgets.iter().map(|b| async move { - let amount = if b.scope.charges_cost() { - cost_micros - } else { - tokens - }; - if amount <= 0 { - return None; + let at = hold.at; + let used = futures::future::join_all(budgets.iter().map(|b| { + let held = hold.take(&b.key); + async move { + let amount = spent.of(b.scope); + match (b.window, held) { + (Window::Day, Some(held)) if held == amount => None, + (Window::Day, Some(held)) => { + Some(gov.quota_settle(&b.key, amount - held, at).await) + } + _ if amount <= 0 => None, + (Window::Day, None) => Some(gov.quota_consume(&b.key, amount).await), + (Window::Month, _) => { + Some(gov.counter_add(&b.key, amount, MONTH_COUNTER_TTL).await) + } + } } - Some(match b.window { - Window::Day => gov.quota_consume(&b.key, amount).await, - Window::Month => gov.counter_add(&b.key, amount, MONTH_COUNTER_TTL).await, - }) })) .await; + // a reload between admission and settle may have dropped a held budget + hold.refund(gov).await; for (b, used) in budgets.iter().zip(used) { if let Some(used) = used && used >= b.limit @@ -1184,13 +1254,70 @@ mod tests { ); } + #[tokio::test] + async fn daily_budgets_hold_in_flight_calls_until_they_settle() { + let yaml = "listen: {host: h, port: 1}\ntenants: [{name: t1, key_daily_cost_quota_micros: 100}]\naccess_keys: [{ak: k1, tenant: t1, product: p, qps: 1, daily_token_quota: 1000}]"; + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let ak = state.auth.get(&access_key_id("k1")).await.unwrap(); + let gov = state.governance.as_ref(); + let key = format!("cb:ak:{}", ak.ak_id); + let est = BudgetSpend { + tokens: 1, + cost_micros: 60, + }; + let now = crate::epoch_secs(); + let reserve = || reserve_budgets(&state, &cfg, &ak, "", est, now); + + let (_, first) = reserve().await.unwrap().unwrap(); + let (_, second) = reserve().await.unwrap().unwrap(); + assert_eq!(gov.quota_used(&key).await, 120, "both calls are in flight"); + let denied = reserve().await.unwrap().unwrap_err(); + assert!( + denied.starts_with("daily cost budget exhausted for key:"), + "{denied}" + ); + assert_eq!(gov.quota_used(&key).await, 120, "a denial holds nothing"); + + let spent = BudgetSpend { + tokens: 1, + cost_micros: 10, + }; + consume_budgets(&state, &cfg, &ak, "", None, first, spent).await; + assert_eq!(gov.quota_used(&key).await, 70, "the settle charges actuals"); + second.refund(gov).await; + assert_eq!(gov.quota_used(&key).await, 10, "a refund returns the hold"); + let (_, third) = reserve().await.unwrap().unwrap(); + consume_budgets(&state, &cfg, &ak, "", None, third, BudgetSpend::default()).await; + assert_eq!( + gov.quota_used(&key).await, + 10, + "a zero-cost settle refunds the hold" + ); + } + #[tokio::test] async fn monthly_budget_denies_at_the_cap_without_touching_daily_counters() { let (cfg, state, ak) = monthly_fixture(false).await; let gov = state.governance.as_ref(); - check_budgets(&state, &cfg, &ak, "").await.unwrap().unwrap(); - consume_budgets(&state, &cfg, &ak, "", None, 10, 3).await; - let err = check_budgets(&state, &cfg, &ak, "") + reserve_budgets(&state, &cfg, &ak, "", BudgetSpend::default(), 0) + .await + .unwrap() + .unwrap(); + consume_budgets( + &state, + &cfg, + &ak, + "", + None, + BudgetHold::default(), + BudgetSpend { + tokens: 10, + cost_micros: 3, + }, + ) + .await; + let err = reserve_budgets(&state, &cfg, &ak, "", BudgetSpend::default(), 0) .await .unwrap() .unwrap_err(); @@ -1201,7 +1328,10 @@ mod tests { assert_eq!(gov.quota_used(&format!("cb:ak:{}", ak.ak_id)).await, 0); gov.quota_reset_all().await; assert!( - check_budgets(&state, &cfg, &ak, "").await.unwrap().is_err(), + reserve_budgets(&state, &cfg, &ak, "", BudgetSpend::default(), 0) + .await + .unwrap() + .is_err(), "the daily reset leaves month counters alone" ); } diff --git a/crates/state/src/governance.rs b/crates/state/src/governance.rs index 169edc4..8b9aaf4 100644 --- a/crates/state/src/governance.rs +++ b/crates/state/src/governance.rs @@ -36,8 +36,8 @@ pub trait Governance: Send + Sync + std::fmt::Debug { at_epoch_secs: i64, ) -> Option; /// Apply the settle delta (actual - reserved; negative refunds) to the day - /// bucket the paired reserve used (`at_epoch_secs`). - async fn quota_settle(&self, key: &str, delta: i64, at_epoch_secs: i64); + /// bucket the paired reserve used (`at_epoch_secs`); returns the new total, or 0 for a zero delta. + async fn quota_settle(&self, key: &str, delta: i64, at_epoch_secs: i64) -> i64; /// Tokens spent today under `key`. async fn quota_used(&self, key: &str) -> i64; /// Add to `key`'s spent tokens; returns the day's new total. @@ -105,8 +105,8 @@ impl Governance for MemoryGovernance { async fn quota_reserve(&self, key: &str, amount: i64, limit: i64, _at: i64) -> Option { self.quota.reserve(key, amount, limit).then_some(amount) } - async fn quota_settle(&self, key: &str, delta: i64, _at: i64) { - self.quota.settle(key, delta); + async fn quota_settle(&self, key: &str, delta: i64, _at: i64) -> i64 { + self.quota.settle(key, delta) } async fn quota_used(&self, key: &str) -> i64 { self.quota.used(key) @@ -275,9 +275,9 @@ impl Governance for RedisGovernance { } }) } - async fn quota_settle(&self, key: &str, delta: i64, at: i64) { + async fn quota_settle(&self, key: &str, delta: i64, at: i64) -> i64 { if delta == 0 { - return; + return 0; } settle_floored( &self.conn, @@ -286,7 +286,7 @@ impl Governance for RedisGovernance { Duration::from_millis(QUOTA_TTL_MS as u64), None, ) - .await; + .await } async fn quota_consume(&self, key: &str, tokens: i64) -> i64 { self.incr_window( @@ -348,13 +348,14 @@ fn quota_key(key: &str) -> String { /// Apply a settle delta and floor at 0 atomically (a rollover between reserve /// and settle must not plant a negative); keeps an existing TTL, arms `window` otherwise. +/// Returns the new total, 0 when the round-trip failed. async fn settle_floored( conn: &redis::aio::ConnectionManager, key: &str, delta: i64, window: Duration, id: Option, -) { +) -> i64 { static SCRIPT: LazyLock = LazyLock::new(|| { redis::Script::new( "local id = tonumber(ARGV[3]) @@ -368,16 +369,17 @@ async fn settle_floored( ) }); let mut conn = conn.clone(); - if let Err(e) = SCRIPT + SCRIPT .key(key) .arg(delta) .arg(window.as_millis() as i64) .arg(id.unwrap_or(-1)) .invoke_async::(&mut conn) .await - { - tracing::warn!(error = %e, key, "redis settle failed"); - } + .unwrap_or_else(|e| { + tracing::warn!(error = %e, key, "redis settle failed"); + 0 + }) } #[cfg(test)] diff --git a/crates/state/src/lib.rs b/crates/state/src/lib.rs index 8081b6e..dcb97d7 100644 --- a/crates/state/src/lib.rs +++ b/crates/state/src/lib.rs @@ -389,9 +389,11 @@ impl QuotaStore { reserve_on(&mut slot_mut(&self.used, key, || 0), amount, limit) } - /// Apply the settle delta (actual - reserved); never below zero. - pub fn settle(&self, key: &str, delta: i64) { - settle_on(&mut slot_mut(&self.used, key, || 0), delta); + /// Apply the settle delta (actual - reserved); never below zero. Returns the new total. + pub fn settle(&self, key: &str, delta: i64) -> i64 { + let mut e = slot_mut(&self.used, key, || 0); + settle_on(&mut e, delta); + *e } pub fn reset_all(&self) { diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index ad0b553..29af601 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -585,10 +585,17 @@ async fn realtime_gate( admission::check_model_qpm(gov, cfg, &m.served) .await .map_err(throttled)?; - let user_budget = admission::check_budgets(state, cfg, &ak, ak.attributed_user(hint)) - .await - .map_err(|e| (ErrClass::ServiceUnavailable, e.message))? - .map_err(quota_exceeded)?; + let (user_budget, _) = admission::reserve_budgets( + state, + cfg, + &ak, + ak.attributed_user(hint), + admission::BudgetSpend::default(), + gw_state::epoch_secs(), + ) + .await + .map_err(|e| (ErrClass::ServiceUnavailable, e.message))? + .map_err(quota_exceeded)?; if let Some(limit) = admission::model_quota_limit(cfg, &ak, &m.requested) && !gov .quota_check(&admission::model_quota_key(&ak.ak_id, &m.requested), limit) @@ -689,6 +696,7 @@ async fn bill_realtime_turn( model_quota_key, budget_ak: ak, user_budget: admit.user_budget, + budget_hold: Default::default(), }, ) .await; @@ -4582,6 +4590,7 @@ async fn admit_video_job( model_quota_key: None, budget_ak: submitter.as_deref().unwrap_or(&ak), user_budget: None, + budget_hold: Default::default(), }, ) .await; diff --git a/docs/governance.md b/docs/governance.md index 57ae1db..e62a417 100644 --- a/docs/governance.md +++ b/docs/governance.md @@ -81,10 +81,10 @@ models: | QPS | pooled per tenant | `tenants[].qps` | | Daily tokens | per access key | `access_keys[].daily_token_quota` (fleet/Redis: rolls at UTC midnight; single-node in-memory: a ~daily background reset) | | Daily tokens | per (key, model) | `tenants[].model_quotas` default, `access_keys[].model_quotas` override (over it: the tenant `fallback_model`, else refused) | -| Daily tokens | per end user | `tenants[].user_daily_token_quota` (soft) | -| Daily cost | pooled per tenant | `tenants[].daily_cost_quota_micros` (micro-dollars of charged price; soft) | -| Daily cost | per access key | `tenants[].key_daily_cost_quota_micros` (soft) | -| Daily cost | per end user | `tenants[].user_daily_cost_quota_micros` (soft), per-user override via `/admin/tenants/{tenant}/users/{user}/budget` | +| Daily tokens | per end user | `tenants[].user_daily_token_quota` (held) | +| Daily cost | pooled per tenant | `tenants[].daily_cost_quota_micros` (micro-dollars of charged price; held) | +| Daily cost | per access key | `tenants[].key_daily_cost_quota_micros` (held) | +| Daily cost | per end user | `tenants[].user_daily_cost_quota_micros` (held), per-user override via `/admin/tenants/{tenant}/users/{user}/budget` | | Monthly cost | pooled per tenant / per key / per end user | `tenants[].monthly_cost_quota_micros`, `key_monthly_cost_quota_micros`, `user_monthly_cost_quota_micros` (UTC calendar month; soft; `monthly_cost_rollover` carries the unspent remainder); the per-user cap takes a per-user override | | TPM | per access key | `access_keys[].tokens_per_minute` | | QPM | per model | `models[].qpm` | @@ -262,15 +262,21 @@ REST, realtime (the `x-gw-user` hint is captured at WS connect), and batch (each item's `user`, or the submitter's `x-gw-user`, is persisted with the item so a fleet drainer still attributes and budgets it). `GET /admin/usage/users?user=&since=&until=` returns per-(user, model) cost over a billing period (add `format=csv` for -export). `TenantConf.user_daily_token_quota` sets a soft per-user daily cap — -enforced with a hard `400 service_quota_exceeded_exception` once exceeded; -"soft" means check-then-consume, so concurrent in-flight turns can overshoot it -by one before the counter accrues. +export). `TenantConf.user_daily_token_quota` sets a per-user daily cap — +enforced with a hard `400 service_quota_exceeded_exception` once exceeded. +On the REST pipeline every daily budget holds the call's estimate at admission +(the token estimate for the token cap; for the cost caps that estimate's prompt +and `max_tokens` sides at the tenant's price), admitting while spent plus held +is under the cap, so concurrent calls overshoot by at most one; billing settles +the hold to the charged amount and a failed call returns it. A call without +`max_tokens` holds 256 output tokens, so a long completion can still carry one +call past the cap. Realtime turns and the monthly caps check before the call +and accrue after it ("soft": concurrent turns can overshoot by one each). The per-(key, model) quota is soft the same way: it is a routing trigger (the fallback-model degrade), while the reserved per-key daily quota hard-caps spend. The daily cost budgets — `daily_cost_quota_micros` pooled over the tenant's keys, `key_daily_cost_quota_micros` per key, `user_daily_cost_quota_micros` -per end user — are soft caps of the same kind over the ledger's charged +per end user — are caps of the same kind over the ledger's charged `cost_micros` (list price, the tenant's `model_prices`, or a decisions reply's `usage.cost`, so unit-priced surfaces count too); they apply on every surface, realtime and batch included, and a response-cache hit, which bills nothing, counts nothing. From 12dafcd7350eba877eade7f31a18d62cb8988579 Mon Sep 17 00:00:00 2001 From: CMGS Date: Sat, 17 Oct 2026 18:24:28 +0800 Subject: [PATCH 03/22] billing: bill an abandoned stream from the upstream's own usage, not a delivered-text estimate A stream the client abandoned never carries its usage frame, so it billed a cl100k estimate of the delivered text: whole prompt at full price (cached agent contexts up to ~10x over) and only delivered completion (generated-but-undelivered tokens under). OpenRouter records every generation, cancelled ones included, within about 5 s; the settle now polls GET /api/v1/generation for the stream's id and bills its native counts and cost, falling back to the estimate when no record lands. A DLP-buffered stream delivered in part already holds the upstream's usage frame and now bills it instead of the estimate. --- crates/dag/src/nodes.rs | 127 +++++++++++++------ crates/engines/src/generation.rs | 181 ++++++++++++++++++++++++++++ crates/engines/src/lib.rs | 1 + crates/engines/src/openai_engine.rs | 5 + crates/handler/src/lib.rs | 114 +++++++++++++++++- crates/models/src/response.rs | 3 + docs/architecture.md | 3 +- docs/governance.md | 10 +- docs/observability.md | 5 + 9 files changed, 405 insertions(+), 44 deletions(-) create mode 100644 crates/engines/src/generation.rs diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index 545a03c..6c94311 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -13,6 +13,13 @@ use crate::executor::{DagNode, Layer}; const DEFAULT_COMPLETION_RESERVE: i64 = 256; /// Reservation cap, so a hostile `max_tokens` cannot overflow the estimate or the counter. const MAX_RESERVE: i64 = 1_000_000; +/// Waits before each lookup of an abandoned stream's recorded usage; OpenRouter records one in about 5 s. +const GENERATION_LOOKUP_DELAYS: [std::time::Duration; 4] = [ + std::time::Duration::from_secs(3), + std::time::Duration::from_secs(5), + std::time::Duration::from_secs(10), + std::time::Duration::from_secs(20), +]; /// preprocess/model_quota: per-(AK, model) daily cap (AK override, else tenant default); /// over it degrades to the tenant fallback model or denies, so it runs before resolve_model. @@ -682,37 +689,7 @@ impl DagNode for CostCalc { if resp.aborted && resp.completion_tokens == 0 { return bill_aborted_stream(ctx, None).await; } - let rate = - gw_state::model_token_rate(&ctx.cfg, served_model(ctx.request.model_param_v2.as_ref())); - // saturating sums so a malformed usage subtree can't overflow the totals - let mut tokens = match &resp.common_usage { - Some(u) => { - let ti = gw_models::TokenInput { - prompt: u.platform_input, - audio_prompt: u.audio_input, - read_cache: u.read_cache, - write_cache: u.write_cache, - write_cache_1h: u.write_cache_1h, - completion: u.completion, - audio_completion: u.audio_output, - reasoning: u.reason, - }; - BillTokens { - prompt: u.prompt_total(), - completion: u.completion_total(), - billable_prompt: gw_models::weighted_prompt(&ti, &rate), - billable_completion: gw_models::weighted_completion(&ti, &rate), - units: 0, - } - } - // total-only vendors (MiniMax v1): meter the total as completion, not zero - None if resp.prompt_tokens == 0 && resp.completion_tokens == 0 => { - BillTokens::weighted(0, resp.total_tokens, &rate) - } - None => BillTokens::weighted(resp.prompt_tokens, resp.completion_tokens, &rate), - }; - tokens.units = resp.billed_units; - bill(ctx, tokens, false).await + bill_usage(ctx).await } } @@ -760,7 +737,7 @@ impl BillTokens { } } -/// Close a buffered stream's billing: refund before delivery, bill estimates after a partial one. +/// Close a buffered stream's billing: refund before delivery; after a partial one bill the upstream's usage, or estimate without it. pub async fn settle_deferred_stream(ctx: &mut DagContext, delivery: StreamDelivery) -> GResult<()> { if !ctx.billing_deferred { return Ok(()); @@ -777,8 +754,13 @@ pub async fn settle_deferred_stream(ctx: &mut DagContext, delivery: StreamDelive .outcome .as_mut() .ok_or_else(|| GatewayError::internal("stream delivery without an outcome"))?; - blank_stream_response(&mut outcome.response); - bill_aborted_stream(ctx, Some(completion)).await + if outcome.response.completion_tokens == 0 { + blank_stream_response(&mut outcome.response); + return bill_aborted_stream(ctx, Some(completion)).await; + } + outcome.response.message.clear(); + outcome.response.aborted = true; + bill_usage(ctx).await } StreamDelivery::None => { let gov = ctx.state.governance.as_ref(); @@ -804,6 +786,22 @@ async fn bill_aborted_stream( ctx: &mut DagContext, delivered_completion: Option, ) -> GResult<()> { + if let Some(usage) = recorded_usage(ctx).await { + let outcome = ctx + .outcome + .as_mut() + .ok_or_else(|| GatewayError::internal("aborted stream without an outcome"))?; + let common = gw_engines::extract_common_usage(&usage, false); + if let Some(u) = &common { + let r = &mut outcome.response; + (r.prompt_tokens, r.completion_tokens) = (u.prompt_total(), u.completion_total()); + r.total_tokens = r.prompt_tokens.saturating_add(r.completion_tokens); + } + outcome.response.common_usage = common; + outcome.response.raw_usage = Some(usage); + ctx.decide("cost_calc", "aborted stream, billed the provider's record"); + return bill_usage(ctx).await; + } let response = &ctx .outcome .as_ref() @@ -833,6 +831,28 @@ async fn bill_aborted_stream( bill(ctx, BillTokens::weighted(prompt, completion, &rate), true).await } +/// OpenRouter's record of an abandoned stream's usage, polled until it lands. +async fn recorded_usage(ctx: &DagContext) -> Option { + let account = ctx.request.account.as_deref()?; + let id = ctx.outcome.as_ref()?.response.upstream_id.as_str(); + if !gw_engines::generation::records_generation(account, id) { + return None; + } + for delay in GENERATION_LOOKUP_DELAYS { + tokio::time::sleep(delay).await; + match gw_engines::generation::openrouter_usage(ctx.transport.as_ref(), account, id).await { + Ok(Some(usage)) => return Some(usage), + Ok(None) => {} + Err(e) => tracing::warn!(error = %e, id, "generation usage lookup failed"), + } + } + tracing::warn!( + id, + "no recorded usage for an aborted stream; billing the estimate" + ); + None +} + fn blank_stream_response(response: &mut gw_models::GatewayResponse) { response.message.clear(); response.completion_tokens = 0; @@ -842,6 +862,45 @@ fn blank_stream_response(response: &mut gw_models::GatewayResponse) { response.aborted = true; } +async fn bill_usage(ctx: &mut DagContext) -> GResult<()> { + let resp = &ctx + .outcome + .as_ref() + .ok_or_else(|| GatewayError::internal("cost_calc without an engine outcome"))? + .response; + let rate = + gw_state::model_token_rate(&ctx.cfg, served_model(ctx.request.model_param_v2.as_ref())); + // saturating sums so a malformed usage subtree can't overflow the totals + let mut tokens = match &resp.common_usage { + Some(u) => { + let ti = gw_models::TokenInput { + prompt: u.platform_input, + audio_prompt: u.audio_input, + read_cache: u.read_cache, + write_cache: u.write_cache, + write_cache_1h: u.write_cache_1h, + completion: u.completion, + audio_completion: u.audio_output, + reasoning: u.reason, + }; + BillTokens { + prompt: u.prompt_total(), + completion: u.completion_total(), + billable_prompt: gw_models::weighted_prompt(&ti, &rate), + billable_completion: gw_models::weighted_completion(&ti, &rate), + units: 0, + } + } + // total-only vendors (MiniMax v1): meter the total as completion, not zero + None if resp.prompt_tokens == 0 && resp.completion_tokens == 0 => { + BillTokens::weighted(0, resp.total_tokens, &rate) + } + None => BillTokens::weighted(resp.prompt_tokens, resp.completion_tokens, &rate), + }; + tokens.units = resp.billed_units; + bill(ctx, tokens, false).await +} + /// Settle reserves and write the ledger; `estimated` marks an aborted stream's estimated counts. async fn bill(ctx: &mut DagContext, mut tokens: BillTokens, estimated: bool) -> GResult<()> { let ptu_spillover = ctx diff --git a/crates/engines/src/generation.rs b/crates/engines/src/generation.rs new file mode 100644 index 0000000..2ea3d0b --- /dev/null +++ b/crates/engines/src/generation.rs @@ -0,0 +1,181 @@ +//! Post-hoc usage of an OpenRouter generation, which OpenRouter records even for a cancelled stream. + +use gw_consts::Protocol; +use gw_models::{Account, GResult, GatewayError}; +use serde_json::{Value, json}; + +use crate::base::versioned_url; +use crate::transport::{Transport, UpstreamBody, UpstreamRequest}; + +const OPENROUTER_KIND: &str = "openrouter"; + +/// Whether `account` records generation `id` for a later [`openrouter_usage`] lookup. +pub fn records_generation(account: &Account, id: &str) -> bool { + account.kind == OPENROUTER_KIND + && !account.endpoint.is_empty() + && !account.api_key_env.is_empty() + && !id.is_empty() + && id + .bytes() + .all(|b| b.is_ascii_alphanumeric() || b == b'-' || b == b'_') +} + +/// The generation's native counts and cost as chat-wire usage; `None` until OpenRouter records it. +pub async fn openrouter_usage( + transport: &dyn Transport, + account: &Account, + id: &str, +) -> GResult> { + let Some(key) = account.api_key() else { + return Ok(None); + }; + let resp = transport + .send(UpstreamRequest { + protocol: Protocol::OpenaiChat, + method: "GET", + url: versioned_url( + account.endpoint.trim_end_matches('/'), + &format!("generation?id={id}"), + ), + headers: vec![("authorization", format!("Bearer {key}"))], + body: Vec::new(), + stream: false, + account: account.name.clone(), + replay_account: None, + output_cap: 0, + }) + .await?; + if resp.status == 404 { + return Ok(None); + } + let UpstreamBody::Json(body) = resp.body else { + return Err(lookup_failed("generation lookup returned no JSON body")); + }; + if resp.status != 200 { + return Err(lookup_failed(format!( + "generation lookup returned {}", + resp.status + ))); + } + let doc: Value = serde_json::from_slice(&body) + .map_err(|e| lookup_failed(format!("generation lookup body: {e}")))?; + Ok(generation_usage(&doc["data"])) +} + +fn lookup_failed(message: impl Into) -> GatewayError { + GatewayError::new(gw_consts::ErrCode::FED_RESP_RPC_FAILED, 502, message) +} + +/// Map `/generation` data to chat-wire usage, reasoning inside completion as in `completion_tokens`. +fn generation_usage(data: &Value) -> Option { + let count = |k: &str| data[k].as_i64(); + let completion = count("native_tokens_completion")?; + Some(json!({ + "prompt_tokens": count("native_tokens_prompt")?, + "completion_tokens": completion, + "prompt_tokens_details": {"cached_tokens": count("native_tokens_cached").unwrap_or(0)}, + "completion_tokens_details": {"reasoning_tokens": count("native_tokens_reasoning").unwrap_or(0).clamp(0, completion)}, + "cost": data["total_cost"], + })) +} + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + + use super::*; + use crate::transport::UpstreamResponse; + + #[derive(Debug)] + struct Canned(u16, &'static str, Mutex>); + + #[async_trait::async_trait] + impl Transport for Canned { + async fn send(&self, req: UpstreamRequest) -> GResult { + self.2 + .lock() + .unwrap() + .push(format!("{} {}", req.method, req.url)); + Ok(UpstreamResponse { + status: self.0, + body: UpstreamBody::Json(bytes::Bytes::from_static(self.1.as_bytes())), + headers: Default::default(), + }) + } + } + + fn openrouter() -> Account { + Account { + name: "or".into(), + kind: OPENROUTER_KIND.into(), + endpoint: "https://openrouter.ai/api".into(), + api_key_env: "PATH".into(), + ..Default::default() + } + } + + #[tokio::test] + async fn a_recorded_generation_maps_to_chat_usage() { + let t = Canned( + 200, + r#"{"data":{"cancelled":true,"native_tokens_prompt":11,"native_tokens_completion":73,"native_tokens_cached":4,"native_tokens_reasoning":73,"total_cost":9.09e-05}}"#, + Mutex::default(), + ); + let usage = openrouter_usage(&t, &openrouter(), "gen-1792232109-kUA52UZsI2") + .await + .unwrap() + .unwrap(); + assert_eq!( + t.2.lock().unwrap().as_slice(), + ["GET https://openrouter.ai/api/v1/generation?id=gen-1792232109-kUA52UZsI2"] + ); + let common = crate::extract_common_usage(&usage, false).unwrap(); + assert_eq!( + ( + common.prompt_total(), + common.read_cache, + common.completion_total(), + common.reason + ), + (11, 4, 73, 73) + ); + assert_eq!(crate::extract_vendor_cost_micros(&usage), Some(91)); + } + + #[tokio::test] + async fn an_unrecorded_generation_yields_nothing() { + let missing = Canned(404, r#"{"error":{"code":404}}"#, Mutex::default()); + assert!( + openrouter_usage(&missing, &openrouter(), "gen-1") + .await + .unwrap() + .is_none() + ); + let pending = Canned( + 200, + r#"{"data":{"native_tokens_prompt":11}}"#, + Mutex::default(), + ); + assert!( + openrouter_usage(&pending, &openrouter(), "gen-1") + .await + .unwrap() + .is_none() + ); + } + + #[test] + fn only_an_openrouter_account_with_a_plain_id_records_generations() { + assert!(records_generation( + &openrouter(), + "gen-1792232109-kUA52UZsI2" + )); + let other = Account { + kind: "openai".into(), + ..openrouter() + }; + assert!(!records_generation(&other, "gen-1")); + assert!(!records_generation(&openrouter(), "gen/../x")); + assert!(!records_generation(&openrouter(), "")); + } +} diff --git a/crates/engines/src/lib.rs b/crates/engines/src/lib.rs index aed0c23..8f446eb 100644 --- a/crates/engines/src/lib.rs +++ b/crates/engines/src/lib.rs @@ -14,6 +14,7 @@ pub mod engine; pub mod eventstream; pub mod factory; pub mod families; +pub mod generation; pub mod http_transport; pub mod mock_transport; pub mod multipart; diff --git a/crates/engines/src/openai_engine.rs b/crates/engines/src/openai_engine.rs index 39bcb6b..525869c 100644 --- a/crates/engines/src/openai_engine.rs +++ b/crates/engines/src/openai_engine.rs @@ -325,6 +325,11 @@ fn apply_sse_event( if resp.model.is_empty() { resp.model = crate::engine::take_string(&mut v, "/model").unwrap_or_default(); } + if resp.upstream_id.is_empty() + && let Some(id) = crate::engine::take_string(&mut v, "/id") + { + resp.upstream_id = id; + } let mut tool_calls = None; if let Some(delta) = v .get_mut("choices") diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index 7308d78..dcfe587 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -2717,6 +2717,106 @@ mod tests { ); } + #[derive(Debug, Default)] + struct RecordedGenerationStream { + lookups: std::sync::atomic::AtomicUsize, + } + + #[async_trait::async_trait] + impl gw_engines::transport::Transport for RecordedGenerationStream { + async fn send( + &self, + req: gw_engines::transport::UpstreamRequest, + ) -> GResult { + use futures::StreamExt; + use gw_engines::transport::{UpstreamBody, UpstreamResponse}; + if req.method == "GET" { + assert!( + req.url.ends_with("/v1/generation?id=gen-test-1"), + "{}", + req.url + ); + let first = self + .lookups + .fetch_add(1, std::sync::atomic::Ordering::SeqCst) + == 0; + let (status, body) = if first { + (404, r#"{"error":{"code":404}}"#) + } else { + ( + 200, + r#"{"data":{"cancelled":true,"native_tokens_prompt":120,"native_tokens_completion":80,"native_tokens_cached":100,"native_tokens_reasoning":0,"total_cost":0.0005}}"#, + ) + }; + return Ok(UpstreamResponse { + status, + body: UpstreamBody::Json(bytes::Bytes::from_static(body.as_bytes())), + headers: Default::default(), + }); + } + let frames: Vec> = vec![ + Ok(bytes::Bytes::from( + "data: {\"id\":\"gen-test-1\",\"choices\":[{\"delta\":{\"content\":\"partial\"}}]}\n\n", + )), + Err(gw_engines::transport::StreamFault { + timeout: false, + message: "connection reset".to_owned(), + }), + ]; + Ok(UpstreamResponse { + status: 200, + body: UpstreamBody::SseStream(futures::stream::iter(frames).boxed()), + headers: Default::default(), + }) + } + } + + #[tokio::test(start_paused = true)] + async fn an_aborted_openrouter_stream_bills_the_recorded_generation() { + let yaml = "listen: {host: h, port: 1}\nproviders: [{name: openrouter, kind: openrouter, api_key_env: PATH}]\nmodels: [{name: m, provider: openrouter, input_price_per_1k_micros: 1000, output_price_per_1k_micros: 2000, token_rate: {read_cache: 0.1}}]\ntenants: [{name: t1}]\naccess_keys: [{ak: or-ak, tenant: t1, product: p, qps: 100, daily_token_quota: 100000}]"; + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let transport = Arc::new(RecordedGenerationStream::default()); + let h = OnlineHandler::new(gw_state::SharedConfig::new(cfg, state), transport.clone()); + let ak = h.state().auth.get(&access_key_id("or-ak")).await.unwrap(); + + let ctx = h + .run(drained_stream_req("m", "count to a thousand"), ak) + .await + .unwrap(); + + let response = ctx.outcome.expect("outcome").response; + assert!(response.aborted); + assert_eq!( + ( + response.prompt_tokens, + response.completion_tokens, + response.total_tokens + ), + (120, 80, 200), + "the access log and metrics read the record too" + ); + assert_eq!( + transport.lookups.load(std::sync::atomic::Ordering::SeqCst), + 2 + ); + let (_, ledger) = h.state().store.ledger_snapshot(usize::MAX).await.unwrap(); + assert_eq!(ledger.len(), 1); + let row = &ledger[0]; + assert!(!row.estimated, "{row:?}"); + assert_eq!( + (row.prompt_tokens, row.completion_tokens), + (120, 80), + "{row:?}" + ); + assert_eq!( + row.cost_micros, + 20 + 10 + 160, + "uncached at full price, cached at 0.1: {row:?}" + ); + assert_eq!(row.vendor_cost_micros, 500, "{row:?}"); + } + #[derive(Debug)] struct PiiStream; @@ -2732,7 +2832,7 @@ mod tests { "data: {\"choices\":[{\"delta\":{\"content\":\"reach me at jane@corp.com now\"}}]}\n\n", )), Ok(bytes::Bytes::from( - "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n", + "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":5,\"completion_tokens\":7,\"total_tokens\":12}}\n\n", )), Ok(bytes::Bytes::from("data: [DONE]\n\n")), ]; @@ -2777,7 +2877,7 @@ mod tests { } #[tokio::test] - async fn dlp_partial_replay_bills_only_delivered_text() { + async fn dlp_partial_replay_bills_the_upstream_usage() { let h = OnlineHandler::new(handler().config.clone(), Arc::new(PiiStream)); let mut request = chat_req("gpt-4o", "hello"); request.stream = true; @@ -2793,9 +2893,13 @@ mod tests { .unwrap(); let (_, ledger) = h.state().store.ledger_snapshot(10).await.unwrap(); assert_eq!(ledger.len(), 1); - assert!(ledger[0].estimated); - assert!(ledger[0].completion_tokens > 0); - assert!(ctx.outcome.as_ref().unwrap().response.aborted); + assert!(!ledger[0].estimated, "{:?}", ledger[0]); + assert_eq!( + (ledger[0].prompt_tokens, ledger[0].completion_tokens), + (5, 7) + ); + let response = &ctx.outcome.as_ref().unwrap().response; + assert!(response.aborted && response.message.is_empty()); } #[tokio::test] diff --git a/crates/models/src/response.rs b/crates/models/src/response.rs index 135a353..7370ed4 100644 --- a/crates/models/src/response.rs +++ b/crates/models/src/response.rs @@ -62,6 +62,9 @@ pub struct GatewayResponse { /// raw usage sub-tree from the vendor body/last SSE frame. #[serde(skip)] pub raw_usage: Option, + /// The vendor's id for this generation (OpenRouter `gen-…`), for a post-hoc usage lookup. + #[serde(skip)] + pub upstream_id: String, } impl GatewayResponse { diff --git a/docs/architecture.md b/docs/architecture.md index 1138f49..991c341 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -62,7 +62,8 @@ default, so the whole pipeline is testable offline: - **`TokenEncoder`** — prompt token estimation capability (tiktoken cl100k_base BPE default, zero-dependency heuristic fallback). `cost_calc` uses it to bill an aborted stream from the delivered text when the vendor's - usage frame never arrives; token-aware PTU sizing is not ported yet. + usage frame never arrives and the account keeps no generation record + (OpenRouter does: `engines::generation`); token-aware PTU sizing is not ported yet. - **`Store`** — durable records (billing ledger, uploaded files, batch jobs). `MemoryStore` by default; `SqliteStore` when `storage.sqlite_path` is configured; `PostgresStore` when `storage.postgres_url` is — shared diff --git a/docs/governance.md b/docs/governance.md index e62a417..f299648 100644 --- a/docs/governance.md +++ b/docs/governance.md @@ -142,8 +142,9 @@ tokens and 1-hour cache writes take their own `token_rate` weights, a scales what `/v1/batches` items cost. A streaming response that breaks after delivery has begun (client disconnect or -upstream failure) is billed for what was delivered: the vendor's usage frame -never arrives, so the token count is estimated from the request and the +upstream failure) never sees the vendor's usage frame. On an OpenRouter account +it is billed from OpenRouter's record of the generation, which counts what the +model generated; elsewhere the token count is estimated from the request and the delivered text. A disconnect *before* any bytes are sent bills nothing. ## MCP tool access @@ -234,8 +235,9 @@ deltas), so **whenever a tenant redacts output — `dlp_redact` or replayed** rather than forwarded token-by-token — DLP trades incremental delivery for a guarantee that no unmasked text reaches the client. Billing is settled after that replay: a disconnect before the first redacted -chunk refunds the reservation, while a partial replay bills only an estimate of -the delivered text. +chunk refunds the reservation, while a partial replay bills the upstream's usage +frame (what the model generated), or an estimate of the delivered text when the +upstream sent none. Turn both `dlp_redact` and `detect_secrets` off to keep incremental delivery; note the embedded demo config ships with `dlp_redact` on. diff --git a/docs/observability.md b/docs/observability.md index 5913082..3c94672 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -87,6 +87,11 @@ the unit-priced surfaces; also summed in `/admin/usage`), charged `cost_micros` and `vendor_cost_micros`, `created_at_epoch_secs`, the PTU-spillover flag, and an `estimated` flag (set when counts came from an aborted stream rather than a vendor usage payload). +An aborted stream served by an OpenRouter account is billed from OpenRouter's +record of the generation instead (`GET /api/v1/generation`, polled for about +40 s; cancelled generations included), with cache reads at their weight and +OpenRouter's cost as `vendor_cost_micros`; only when no record arrives does it +fall back to the estimate. Writes are idempotent by `request_id`; a transient store failure is held in a bounded repair queue and retried with backoff — eight attempts, about 13 s — before the batch is dropped and counted in `gateway_ledger_write_failures_total`. From eace497158493c8db516a3c4b7005b93948c21c7 Mon Sep 17 00:00:00 2001 From: CMGS Date: Sat, 17 Oct 2026 19:10:34 +0800 Subject: [PATCH 04/22] budgets: reuse QuotaCheck's estimate in BudgetGate BudgetGate re-derived the estimate parts, serializing a decisions request's fields a second time on every call. --- crates/dag/src/context.rs | 2 ++ crates/dag/src/nodes.rs | 6 ++++-- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/crates/dag/src/context.rs b/crates/dag/src/context.rs index 9a10e0a..f5f9f8f 100644 --- a/crates/dag/src/context.rs +++ b/crates/dag/src/context.rs @@ -28,6 +28,7 @@ pub struct DagContext { /// A fallback model retries this attempt's upstream fault, so its failure is not the outcome. pub fallback_ahead: bool, pub(crate) request_limits_admitted: bool, + pub(crate) reserve_parts: (i64, i64), /// This request's cache key (computed by cache_lookup, reused by cache_store). pub cache_key: Option, /// Governance key for the (AK, model) daily counter; set only when a cap is configured. @@ -65,6 +66,7 @@ impl DagContext { cache_hit: false, fallback_ahead: false, request_limits_admitted: false, + reserve_parts: (0, 0), cache_key: None, model_quota_key: None, quota_reserved: None, diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index 6c94311..21e7c0f 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -280,7 +280,9 @@ impl DagNode for QuotaCheck { "quota_check" } async fn execute(&self, ctx: &mut DagContext) -> GResult<()> { - let est = reserve_estimate(&ctx.request); + ctx.reserve_parts = reserve_parts(&ctx.request); + let (prompt, max_out) = ctx.reserve_parts; + let est = prompt.saturating_add(max_out).min(MAX_RESERVE); let at = gw_state::epoch_secs(); let reserved = admission::reserve_daily(ctx.state.governance.as_ref(), &ctx.ak, est, at) .await @@ -437,7 +439,7 @@ impl DagNode for BudgetGate { "budget" } async fn execute(&self, ctx: &mut DagContext) -> GResult<()> { - let (prompt, max_out) = reserve_parts(&ctx.request); + let (prompt, max_out) = ctx.reserve_parts; let served = served_model(ctx.request.model_param_v2.as_ref()); let est = admission::BudgetSpend { tokens: ctx From e1aac7b0acb9d62a1eba622f457357422d77a975 Mon Sep 17 00:00:00 2001 From: CMGS Date: Sat, 17 Oct 2026 19:27:16 +0800 Subject: [PATCH 05/22] billing: keep the ledger spool from losing or wedging rows - A spool that cannot be written fell back to dropping rows after three attempts, below the eight retries an unspooled ledger gets; the spool now creates its directory and, on a write failure, the store retries continue. - Postgres TEXT refuses NUL, so a client "user" of "\u0000" failed the whole multi-row insert: before the spool every row in that batch was dropped, with it replay wedged on the row forever. SQL billing inserts now strip NUL. - Replay reads the spool in LEDGER_BATCH_MAX batches instead of loading the whole file, and re-rolls the usage rollup back to the oldest replayed row, so a row billed before the rollup watermark is counted and never pruned unrolled. - BillingRecord deserializes with defaults, so a spool written by an older release replays after an upgrade instead of being skipped as torn. --- crates/state/src/admission.rs | 266 ++++++++++++++++++++++++---------- crates/state/src/store.rs | 128 ++++++++++++---- docs/observability.md | 10 +- 3 files changed, 295 insertions(+), 109 deletions(-) diff --git a/crates/state/src/admission.rs b/crates/state/src/admission.rs index 1b760f8..eba7c15 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -4,15 +4,15 @@ //! two admission paths cannot drift. Denials carry the user-facing message; //! callers wrap it in their own wire shape. -use std::io::Write; -use std::os::unix::fs::FileExt; -use std::path::PathBuf; +use std::fs::{File, OpenOptions}; +use std::io::{BufRead, BufReader, Seek, SeekFrom, Write}; +use std::path::{Path, PathBuf}; use std::sync::Arc; use std::time::Duration; use gw_config::{GatewayConfig, TenantConf}; use gw_models::GResult; -use tokio::sync::{mpsc, oneshot}; +use tokio::sync::{Mutex, mpsc, oneshot}; use crate::store::{BillingInput, BillingRecord, Store, billing_record}; use crate::{AkInfo, GatewayState, Governance, UserBudget, clamp_tokens}; @@ -244,29 +244,30 @@ impl BillingLedger { if batch.is_empty() { return true; } - let attempts = if spool.is_some() { - LEDGER_SPOOL_AFTER_ATTEMPTS - } else { - LEDGER_RETRY_ATTEMPTS - }; let mut delay = LEDGER_RETRY_INITIAL; let mut committed = false; - for attempt in 1..=attempts { + for attempt in 1..=LEDGER_RETRY_ATTEMPTS { let Err(e) = store.ledger_add_batch(batch).await else { committed = true; break; }; metrics::counter!("gateway_ledger_write_failures_total").increment(batch.len() as u64); - if attempt == attempts { - committed = match spool { - Some(spool) => spool.append(batch).await, - None => false, - }; - if committed { - tracing::warn!(error = %e, rows = batch.len(), "billing ledger write failed; rows spooled for replay"); - } else { - tracing::error!(error = %e, rows = batch.len(), "billing ledger retries exhausted; rows dropped"); + if let Some(spool) = spool + && attempt >= LEDGER_SPOOL_AFTER_ATTEMPTS + { + match spool.append(batch).await { + Ok(()) => { + tracing::warn!(error = %e, rows = batch.len(), "billing ledger write failed; rows spooled for replay"); + committed = true; + break; + } + Err(spool_err) => { + tracing::error!(error = %spool_err, rows = batch.len(), "billing spool write failed") + } } + } + if attempt == LEDGER_RETRY_ATTEMPTS { + tracing::error!(error = %e, rows = batch.len(), "billing ledger retries exhausted; rows dropped"); break; } tracing::error!(error = %e, rows = batch.len(), "billing ledger write failed; retrying"); @@ -278,12 +279,10 @@ impl BillingLedger { } } -/// Ledger rows the store refused, appended to a local JSON-lines file and replayed in order. -#[derive(Debug)] struct LedgerSpool { path: PathBuf, replaying: PathBuf, - lock: tokio::sync::Mutex<()>, + lock: Mutex<()>, } impl LedgerSpool { @@ -293,93 +292,90 @@ impl LedgerSpool { Self { path, replaying, - lock: tokio::sync::Mutex::new(()), + lock: Mutex::new(()), } } - async fn append(&self, rows: &[BillingRecord]) -> bool { - let mut lines = Vec::new(); + async fn append(&self, rows: &[BillingRecord]) -> std::io::Result<()> { + // a leading newline ends a row a crash tore + let mut lines = vec![b'\n']; for r in rows { - if let Err(e) = serde_json::to_writer(&mut lines, r) { - tracing::error!(error = %e, request_id = r.request_id, "spool billing row"); - return false; - } + serde_json::to_writer(&mut lines, r)?; lines.push(b'\n'); } let path = self.path.clone(); let _held = self.lock.lock().await; - let written = blocking(move || { - let mut f = std::fs::OpenOptions::new() - .create(true) - .read(true) - .append(true) - .open(&path)?; - let len = f.metadata()?.len(); - let mut last = *b"\n"; - if len > 0 { - f.read_exact_at(&mut last, len - 1)?; - } - // a crash mid-append leaves a torn row; the next row starts on its own line - if last[0] != b'\n' { - f.write_all(b"\n")?; + blocking(move || { + if let Some(dir) = path.parent() { + std::fs::create_dir_all(dir)?; } + let mut f = OpenOptions::new().create(true).append(true).open(&path)?; f.write_all(&lines)?; f.sync_data() }) - .await; - if let Err(e) = &written { - tracing::error!(error = %e, path = %self.path.display(), "billing spool write failed"); - } - written.is_ok() + .await } /// `request_id` uniqueness makes a re-insert of an already accepted batch a no-op. async fn replay(&self, store: &Arc) { - let (path, replaying) = (self.path.clone(), self.replaying.clone()); - let read = { + let claimed = { + let (path, replaying) = (self.path.clone(), self.replaying.clone()); let _held = self.lock.lock().await; blocking(move || { - if !replaying.exists() { - if !path.exists() { - return Ok(None); - } - std::fs::rename(&path, &replaying)?; + if replaying.exists() { + return Ok(true); + } + if !path.exists() { + return Ok(false); } - std::fs::read(&replaying).map(Some) + std::fs::rename(&path, &replaying).map(|()| true) }) .await }; - let bytes = match read { - Ok(Some(bytes)) => bytes, - Ok(None) => return, + match claimed { + Ok(true) => {} + Ok(false) => return, Err(e) => { - tracing::error!(error = %e, path = %self.replaying.display(), "billing spool read failed"); + tracing::error!(error = %e, path = %self.path.display(), "billing spool claim failed"); return; } - }; - let rows: Vec = bytes - .split(|&b| b == b'\n') - .filter(|line| !line.is_empty()) - .filter_map(|line| { - serde_json::from_str(&String::from_utf8_lossy(line)) - .inspect_err( - |e| tracing::error!(error = %e, "skipping a torn billing spool line"), - ) - .ok() - }) - .collect(); - for chunk in rows.chunks(LEDGER_BATCH_MAX) { - if let Err(e) = store.ledger_add_batch(chunk).await { + } + let (mut offset, mut replayed, mut since) = (0, 0, i64::MAX); + loop { + let replaying = self.replaying.clone(); + let (rows, next) = match blocking(move || read_spool_batch(&replaying, offset)).await { + Ok(batch) => batch, + Err(e) => { + tracing::error!(error = %e, path = %self.replaying.display(), "billing spool read failed"); + return; + } + }; + if next == offset { + break; + } + if let Err(e) = store.ledger_add_batch(&rows).await { tracing::warn!(error = %e, rows = rows.len(), "billing spool replay deferred"); return; } + replayed += rows.len(); + since = rows + .iter() + .map(|r| r.created_at_epoch_secs) + .fold(since, i64::min); + offset = next; + } + if replayed > 0 + && let Err(e) = store.usage_rollup_since(since, crate::epoch_secs()).await + { + tracing::warn!(error = %e, rows = replayed, "billing spool rollup deferred"); + return; } let replaying = self.replaying.clone(); if let Err(e) = blocking(move || std::fs::remove_file(replaying)).await { tracing::error!(error = %e, path = %self.replaying.display(), "billing spool remove failed"); return; } - tracing::info!(rows = rows.len(), "billing spool replayed"); + tracing::info!(rows = replayed, "billing spool replayed"); } } @@ -893,6 +889,29 @@ fn month_prefix((y, m): (i64, u32)) -> String { format!("m:{y}{m:02}:") } +/// Up to `LEDGER_BATCH_MAX` rows from `offset`, and the offset after them; equal offsets mean the end. +fn read_spool_batch(path: &Path, offset: u64) -> std::io::Result<(Vec, u64)> { + let mut reader = BufReader::new(File::open(path)?); + reader.seek(SeekFrom::Start(offset))?; + let (mut rows, mut next, mut line) = (Vec::new(), offset, Vec::new()); + while rows.len() < LEDGER_BATCH_MAX { + line.clear(); + let n = reader.read_until(b'\n', &mut line)?; + if n == 0 { + break; + } + next += n as u64; + if line.trim_ascii().is_empty() { + continue; + } + match serde_json::from_slice(&line) { + Ok(row) => rows.push(row), + Err(e) => tracing::error!(error = %e, "skipping a torn billing spool line"), + } + } + Ok((rows, next)) +} + async fn blocking(f: F) -> std::io::Result where T: Send + 'static, @@ -1142,7 +1161,7 @@ mod tests { "a spooled row is not a dropped row" ); let spooled = std::fs::read_to_string(&path).unwrap(); - assert_eq!(spooled.lines().count(), 1); + assert_eq!(spooled.lines().filter(|l| !l.is_empty()).count(), 1); assert!(spooled.contains("req-spooled")); } @@ -1152,13 +1171,16 @@ mod tests { let spool = LedgerSpool::new(dir.path().join("ledger-spool.jsonl").to_str().unwrap()); let store = Arc::new(crate::MemoryStore::default()); let as_store: Arc = store.clone(); - assert!(spool.append(&[record("req-a"), record("req-b")]).await); + spool + .append(&[record("req-a"), record("req-b")]) + .await + .unwrap(); store.fail_next_ledger_writes(1); spool.replay(&as_store).await; assert_eq!(store.ledger_snapshot(usize::MAX).await.unwrap().0, 0); assert!(spool.replaying.exists(), "a refused replay keeps its rows"); - assert!(spool.append(&[record("req-c")]).await); + spool.append(&[record("req-c")]).await.unwrap(); spool.replay(&as_store).await; spool.replay(&as_store).await; @@ -1174,7 +1196,7 @@ mod tests { let dir = tempfile::tempdir().unwrap(); let spool = LedgerSpool::new(dir.path().join("ledger-spool.jsonl").to_str().unwrap()); std::fs::write(&spool.path, b"{\"ak\":\"torn").unwrap(); - assert!(spool.append(&[record("req-after-crash")]).await); + spool.append(&[record("req-after-crash")]).await.unwrap(); let store = Arc::new(crate::MemoryStore::default()); let as_store: Arc = store.clone(); @@ -1185,6 +1207,92 @@ mod tests { assert_eq!(rows[0].request_id, "req-after-crash"); } + #[tokio::test] + async fn an_unwritable_spool_keeps_retrying_the_store() { + let dir = tempfile::tempdir().unwrap(); + let blocker = dir.path().join("not-a-dir"); + std::fs::write(&blocker, b"").unwrap(); + let store = Arc::new(crate::MemoryStore::default()); + store.fail_next_ledger_writes(LEDGER_SPOOL_AFTER_ATTEMPTS + 1); + let path = blocker.join("ledger-spool.jsonl"); + let ledger = BillingLedger::repairing(store.clone(), path.to_str().unwrap()); + + ledger.write(record("req-unspooled")).await; + + assert_eq!(ledger.flush().await, 0); + let (count, rows) = store.ledger_snapshot(usize::MAX).await.unwrap(); + assert_eq!((count, rows[0].request_id.as_str()), (1, "req-unspooled")); + } + + #[tokio::test] + async fn a_spool_dir_that_does_not_exist_yet_is_created() { + let dir = tempfile::tempdir().unwrap(); + let spool = LedgerSpool::new(dir.path().join("gw/ledger/spool.jsonl").to_str().unwrap()); + spool.append(&[record("req-nested")]).await.unwrap(); + assert!( + std::fs::read_to_string(&spool.path) + .unwrap() + .contains("req-nested") + ); + } + + #[tokio::test] + async fn replay_streams_batches_and_rolls_up_rows_billed_before_the_watermark() { + let dir = tempfile::tempdir().unwrap(); + let spool = LedgerSpool::new(dir.path().join("ledger-spool.jsonl").to_str().unwrap()); + let store = Arc::new(crate::MemoryStore::default()); + let as_store: Arc = store.clone(); + let now = crate::epoch_secs(); + let mut late = record("req-current"); + late.created_at_epoch_secs = now - 600; + store.ledger_add(&late).await.unwrap(); + store.usage_rollup_advance(now).await.unwrap(); + let billed_at = now - 7_200; + let rows: Vec<_> = (0..=LEDGER_BATCH_MAX) + .map(|i| BillingRecord { + created_at_epoch_secs: billed_at, + ..record(format!("req-late-{i}")) + }) + .collect(); + spool.append(&rows).await.unwrap(); + + spool.replay(&as_store).await; + + assert_eq!( + store.ledger_snapshot(usize::MAX).await.unwrap().0, + LEDGER_BATCH_MAX + 2 + ); + let usage = store + .usage_by_user(None, None, billed_at, billed_at + 60) + .await + .unwrap(); + assert_eq!( + usage.iter().map(|u| u.requests).sum::(), + LEDGER_BATCH_MAX as i64 + 1, + "{usage:?}" + ); + assert!(!spool.replaying.exists()); + } + + #[tokio::test] + async fn a_spool_row_from_an_older_release_still_replays() { + let dir = tempfile::tempdir().unwrap(); + let spool = LedgerSpool::new(dir.path().join("ledger-spool.jsonl").to_str().unwrap()); + std::fs::write( + &spool.path, + br#"{"ak":"ak","tenant":"default","request_id":"req-old","created_at_epoch_secs":1,"model":"m","total_tokens":2}"#, + ) + .unwrap(); + let store = Arc::new(crate::MemoryStore::default()); + let as_store: Arc = store.clone(); + + spool.replay(&as_store).await; + + let (count, rows) = store.ledger_snapshot(usize::MAX).await.unwrap(); + assert_eq!((count, rows[0].request_id.as_str()), (1, "req-old")); + assert_eq!(rows[0].total_tokens, 2); + } + #[test] fn civil_month_follows_the_gregorian_calendar() { assert_eq!(civil_month(0), (1970, 1)); diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index f62c5d9..d1a59f8 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -4,6 +4,7 @@ //! [`SqliteStore`] (`storage.sqlite_path`, one durable node), and //! [`PostgresStore`] (`storage.postgres_url`, shared across a fleet). +use std::borrow::Cow; use std::collections::{BTreeMap, HashMap, HashSet}; use std::sync::Mutex; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -60,7 +61,8 @@ const BATCH_ITEM_CHUNK: usize = 13_000; const PG_MAX_CONNECTIONS: u32 = 10; /// One billing entry (recorded locally only; no reporting upstream). -#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)] +#[serde(default)] pub struct BillingRecord { pub ak: String, pub product: String, @@ -386,9 +388,12 @@ pub trait Store: Send + Sync + std::fmt::Debug { until: i64, bucket_secs: i64, ) -> GResult>; - /// Recompute and upsert every rollup bucket in the trailing backfill window; buckets are never - /// deleted, so a missed tick self-heals. Returns the buckets written. - async fn usage_rollup_advance(&self, now_epoch_secs: i64) -> GResult; + async fn usage_rollup_advance(&self, now_epoch_secs: i64) -> GResult { + self.usage_rollup_since(i64::MAX, now_epoch_secs).await + } + /// Recompute and upsert every rollup bucket in the trailing backfill window, or back to `since` + /// for rows billed late; buckets are never deleted, so a missed tick self-heals. Returns the buckets written. + async fn usage_rollup_since(&self, since_epoch_secs: i64, now_epoch_secs: i64) -> GResult; /// Delete rolled billing rows past `ledger_max_rows`; returns rows deleted. async fn ledger_prune(&self) -> GResult { Ok(0) @@ -735,11 +740,13 @@ impl Store for MemoryStore { Ok(map.into_iter().collect()) } - async fn usage_rollup_advance(&self, now: i64) -> GResult { + async fn usage_rollup_since(&self, since: i64, now: i64) -> GResult { let hi = bucket_floor(now - ROLLUP_SETTLE_SECS); let mut rollup = lock(&self.rollup); let watermark = rollup_watermark(&rollup); - let lo = (hi - ROLLUP_BACKFILL_SECS).min(watermark); + let lo = (hi - ROLLUP_BACKFILL_SECS) + .min(watermark) + .min(bucket_floor(since)); let mut fresh = BTreeMap::new(); { let ledger = lock(&self.ledger); @@ -1579,21 +1586,21 @@ macro_rules! sql_store_impl { created_at_epoch_secs, estimated, billed_units) ", ); qb.push_values(rows, |mut v, r| { - v.push_bind(&r.ak) - .push_bind(&r.product) - .push_bind(&r.tenant) - .push_bind(&r.model) - .push_bind(&r.served_model) - .push_bind(&r.protocol) - .push_bind(&r.account) + v.push_bind(nul_free(&r.ak)) + .push_bind(nul_free(&r.product)) + .push_bind(nul_free(&r.tenant)) + .push_bind(nul_free(&r.model)) + .push_bind(nul_free(&r.served_model)) + .push_bind(nul_free(&r.protocol)) + .push_bind(nul_free(&r.account)) .push_bind(r.prompt_tokens) .push_bind(r.completion_tokens) .push_bind(r.total_tokens) .push_bind(r.cost_micros) .push_bind(r.vendor_cost_micros) .push_bind(r.ptu_spillover) - .push_bind(&r.user_id) - .push_bind(&r.request_id) + .push_bind(nul_free(&r.user_id)) + .push_bind(nul_free(&r.request_id)) .push_bind(r.created_at_epoch_secs) .push_bind(r.estimated) .push_bind(r.billed_units); @@ -1906,14 +1913,18 @@ macro_rules! sql_store_impl { } sql_store_impl!(SqliteStore, sqlite, { - async fn usage_rollup_advance(&self, now: i64) -> GResult { + async fn usage_rollup_since(&self, since: i64, now: i64) -> GResult { let hi = bucket_floor(now - ROLLUP_SETTLE_SECS); let watermark: i64 = sqlx::query_scalar(ROLLUP_WATERMARK_SQL) .fetch_one(&self.pool) .await .map_err(|e| crate::sqlx_err("read rollup watermark", e))?; let r = sqlx::query(dialect_sql!(sqlite, q_rollup_insert!(sqlite))) - .bind((hi - ROLLUP_BACKFILL_SECS).min(watermark)) + .bind( + (hi - ROLLUP_BACKFILL_SECS) + .min(watermark) + .min(bucket_floor(since)), + ) .bind(hi) .execute(&self.pool) .await @@ -2217,27 +2228,39 @@ impl PostgresStore { } sql_store_impl!(PostgresStore, postgres, { - async fn usage_rollup_advance(&self, now: i64) -> GResult { + async fn usage_rollup_since(&self, since: i64, now: i64) -> GResult { let hi = bucket_floor(now - ROLLUP_SETTLE_SECS); let mut tx = self .pool .begin() .await .map_err(|e| crate::sqlx_err("begin rollup", e))?; - let locked: bool = sqlx::query_scalar("SELECT pg_try_advisory_xact_lock($1)") - .bind(ROLLUP_LOCK_KEY) - .fetch_one(&mut *tx) - .await - .map_err(|e| crate::sqlx_err("acquire rollup lock", e))?; - if !locked { - return Ok(0); + if since == i64::MAX { + let locked: bool = sqlx::query_scalar("SELECT pg_try_advisory_xact_lock($1)") + .bind(ROLLUP_LOCK_KEY) + .fetch_one(&mut *tx) + .await + .map_err(|e| crate::sqlx_err("acquire rollup lock", e))?; + if !locked { + return Ok(0); + } + } else { + sqlx::query("SELECT pg_advisory_xact_lock($1)") + .bind(ROLLUP_LOCK_KEY) + .execute(&mut *tx) + .await + .map_err(|e| crate::sqlx_err("wait for rollup lock", e))?; } let watermark: i64 = sqlx::query_scalar(ROLLUP_WATERMARK_SQL) .fetch_one(&mut *tx) .await .map_err(|e| crate::sqlx_err("read rollup watermark", e))?; let r = sqlx::query(dialect_sql!(postgres, q_rollup_insert!(postgres))) - .bind((hi - ROLLUP_BACKFILL_SECS).min(watermark)) + .bind( + (hi - ROLLUP_BACKFILL_SECS) + .min(watermark) + .min(bucket_floor(since)), + ) .bind(hi) .execute(&mut *tx) .await @@ -2757,6 +2780,14 @@ where ) } +fn nul_free(s: &str) -> Cow<'_, str> { + if s.contains('\0') { + Cow::Owned(s.replace('\0', "")) + } else { + Cow::Borrowed(s) + } +} + fn bucket_floor(ts: i64) -> i64 { ts - ts.rem_euclid(ROLLUP_BUCKET_SECS) } @@ -3450,6 +3481,49 @@ mod tests { vec![(0, 2), (1_200, 1)], "rolled buckets and the raw tail group by bucket start" ); + + let mut late = record("m1"); + late.tenant = tenant.clone(); + late.user_id = "dave".into(); + late.created_at_epoch_secs = 100; + store.ledger_add(&late).await.unwrap(); + let billed_late = |rows: Vec| rows.iter().any(|r| r.user_id == "dave"); + store.usage_rollup_advance(1_500).await.unwrap(); + assert!( + !billed_late( + store + .usage_by_user(Some(&tenant), None, 0, 1_000) + .await + .unwrap() + ), + "a row billed below the watermark misses the trailing rollup" + ); + store.usage_rollup_since(100, 1_500).await.unwrap(); + assert!( + billed_late( + store + .usage_by_user(Some(&tenant), None, 0, 1_000) + .await + .unwrap() + ), + "a rollup since its billing time counts the late row" + ); + } + + async fn exercise_nul_text(store: &dyn Store) { + let mut r = record("m1"); + r.user_id = "nul\0user".into(); + r.request_id = format!("{}\0", r.request_id); + store + .ledger_add_batch(&[r.clone(), record("m1")]) + .await + .unwrap(); + let (_, rows) = store.ledger_snapshot(usize::MAX).await.unwrap(); + let stored = rows + .iter() + .find(|row| row.request_id == r.request_id.replace('\0', "")) + .unwrap(); + assert_eq!(stored.user_id, "nuluser"); } #[tokio::test] @@ -3464,6 +3538,7 @@ mod tests { .await .unwrap(); exercise_rollup(&store, "").await; + exercise_nul_text(&store).await; } #[tokio::test] @@ -3838,6 +3913,7 @@ mod tests { let store = PostgresStore::connect(&own_url).await.expect("pg connect"); let ns = format!("-{nonce}"); exercise_rollup(&store, &ns).await; + exercise_nul_text(&store).await; exercise_erase(&store, &ns).await; exercise_content_list(&store, &ns).await; exercise_terminal_idempotence(&store, &ns).await; diff --git a/docs/observability.md b/docs/observability.md index 3c94672..6a13ea4 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -96,10 +96,12 @@ Writes are idempotent by `request_id`; a transient store failure is held in a bounded repair queue and retried with backoff — eight attempts, about 13 s — before the batch is dropped and counted in `gateway_ledger_write_failures_total`. With `storage.ledger_spool_path` set, the batch instead goes to that local -JSON-lines file after three attempts (about 0.3 s), and a replay every 30 s -inserts it once the store accepts writes again (`billing spool replayed` in the -log). A row replayed more than 20 minutes after it was billed is in the ledger -but not in the usage rollup that `GET /admin/usage/users` reads. +JSON-lines file after three attempts (about 0.3 s; the directory is created on +first use, and a spool that cannot be written falls back to the remaining +retries), and a replay every 30 s inserts it in batches once the store accepts +writes again, then re-rolls the usage rollup back to the oldest replayed row +(`billing spool replayed` in the log). The spool path is per-instance local +disk; replicas must not share one file. A full queue applies backpressure to settlement until the repair worker makes room, so overload alone never discards an accepted row. Per-user usage additionally rolls into durable minute buckets every minute, so From 4ee2a12751800c4a1df11895ff9b046c31c095c8 Mon Sep 17 00:00:00 2001 From: CMGS Date: Sat, 17 Oct 2026 19:31:09 +0800 Subject: [PATCH 06/22] budgets: price the hold only when a day budget applies - BudgetGate priced the estimate (tenant price lookup, cost math) on every call, even a key with no budget; reserve_budgets now takes the estimate as a closure and evaluates it only when a day budget is configured. - The settle's trailing refund ran on an empty hold: settle reuses admission's config snapshot, user and override, so its day keys always match the held ones. - A settle whose actual equals its hold skipped the budget_exhausted alert; it now reads the counter for the cap check. - The two refund sites fold into DagContext::refund_reserves, which returns the quota/TPM reservation and the budget hold concurrently. - quota_settle's docs now match both backends (0 when nothing was applied), and the governance doc no longer claims a cap overshoots by at most one call. --- crates/dag/src/context.rs | 13 +++++++ crates/dag/src/nodes.rs | 23 +++--------- crates/handler/src/lib.rs | 10 +---- crates/state/src/admission.rs | 69 +++++++++++++++++++++++++--------- crates/state/src/governance.rs | 7 +++- crates/views/src/lib.rs | 2 +- docs/governance.md | 10 +++-- 7 files changed, 84 insertions(+), 50 deletions(-) diff --git a/crates/dag/src/context.rs b/crates/dag/src/context.rs index f5f9f8f..b665539 100644 --- a/crates/dag/src/context.rs +++ b/crates/dag/src/context.rs @@ -90,6 +90,19 @@ impl DagContext { .ok_or_else(|| GatewayError::internal("model param missing after resolve_model")) } + pub async fn refund_reserves(&mut self) { + let gov = self.state.governance.as_ref(); + tokio::join!( + gov.refund_reserves( + &self.ak.ak_id, + self.quota_reserved.take().unwrap_or(0), + self.tpm_reserved.take(), + self.quota_at, + ), + std::mem::take(&mut self.budget_hold).refund(gov), + ); + } + /// The effective end user: the key's `owner`, else request metadata, else `""`. pub fn effective_user_id(&self) -> &str { self.ak diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index 21e7c0f..7814043 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -430,7 +430,7 @@ impl DagNode for AkTpmLimit { } } -/// model_access/budget: holds the call's estimate against the tenant's daily budgets for the key and user. +/// model_access/budget: holds the call's estimate against the daily budgets (tenant, key, user) and checks the monthly ones. pub struct BudgetGate; #[async_trait::async_trait] @@ -440,16 +440,13 @@ impl DagNode for BudgetGate { } async fn execute(&self, ctx: &mut DagContext) -> GResult<()> { let (prompt, max_out) = ctx.reserve_parts; - let served = served_model(ctx.request.model_param_v2.as_ref()); - let est = admission::BudgetSpend { - tokens: ctx - .quota_reserved - .filter(|&reserved| reserved > 0) - .unwrap_or_else(|| prompt.saturating_add(max_out).min(MAX_RESERVE)), + let (cfg, request, tenant) = (&ctx.cfg, &ctx.request, &ctx.ak.tenant); + let est = || admission::BudgetSpend { + tokens: prompt.saturating_add(max_out).min(MAX_RESERVE), cost_micros: gw_models::cost_micros( prompt, max_out, - ctx.cfg.prices_for_tenant(&ctx.ak.tenant, served), + cfg.prices_for_tenant(tenant, served_model(request.model_param_v2.as_ref())), ), }; let (over, hold) = admission::reserve_budgets( @@ -765,15 +762,7 @@ pub async fn settle_deferred_stream(ctx: &mut DagContext, delivery: StreamDelive bill_usage(ctx).await } StreamDelivery::None => { - let gov = ctx.state.governance.as_ref(); - gov.refund_reserves( - &ctx.ak.ak_id, - ctx.quota_reserved.take().unwrap_or(0), - ctx.tpm_reserved.take(), - ctx.quota_at, - ) - .await; - std::mem::take(&mut ctx.budget_hold).refund(gov).await; + ctx.refund_reserves().await; if let Some(outcome) = ctx.outcome.as_mut() { outcome.response.prompt_tokens = 0; blank_stream_response(&mut outcome.response); diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index dcfe587..d8e28f9 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -270,15 +270,7 @@ impl OnlineHandler { break; }; // a failed pipeline refunds its reservations whole, on the reserve's day bucket - let gov = ctx.state.governance.as_ref(); - gov.refund_reserves( - &ctx.ak.ak_id, - ctx.quota_reserved.take().unwrap_or(0), - ctx.tpm_reserved.take(), - ctx.quota_at, - ) - .await; - std::mem::take(&mut ctx.budget_hold).refund(gov).await; + ctx.refund_reserves().await; // signed thinking replays only against the model that produced it if is_upstream_fault(&e) && !ctx.request.replays_reasoning_output() diff --git a/crates/state/src/admission.rs b/crates/state/src/admission.rs index eba7c15..a7ac420 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -479,18 +479,23 @@ pub async fn flush_billing(state: &GatewayState) { } } -/// Hold `est` against the day budgets, admitting while spent plus held is under each cap; month budgets and a zero `est` only check. +/// Hold the estimate `est` yields against the day budgets, admitting while spent plus held is under each cap; month budgets and a zero estimate only check. pub async fn reserve_budgets( state: &GatewayState, cfg: &GatewayConfig, ak: &AkInfo, user: &str, - est: BudgetSpend, + est: impl FnOnce() -> BudgetSpend, at: i64, ) -> GResult> { let over = user_budget(state, ak, user).await?; let gov = state.governance.as_ref(); let budgets = budgets(gov, cfg, ak, user, over).await; + let est = if budgets.iter().any(|b| b.window == Window::Day) { + est() + } else { + BudgetSpend::default() + }; let admitted = futures::future::join_all(budgets.iter().map(|b| async move { let amount = est.of(b.scope); match b.window { @@ -509,15 +514,16 @@ pub async fn reserve_budgets( match admitted { Some(held) if held > 0 => hold.held.push((b.key, held)), Some(_) => {} - None if denied.is_none() => { - denied = Some(format!( - "{} {} budget exhausted for {}", - b.window.label(), - b.scope.unit(), - b.scope.subject(ak, user) - )); + None => { + denied.get_or_insert_with(|| { + format!( + "{} {} budget exhausted for {}", + b.window.label(), + b.scope.unit(), + b.scope.subject(ak, user) + ) + }); } - None => {} } } let Some(denied) = denied else { @@ -756,7 +762,7 @@ async fn consume_budgets( async move { let amount = spent.of(b.scope); match (b.window, held) { - (Window::Day, Some(held)) if held == amount => None, + (Window::Day, Some(held)) if held == amount => Some(gov.quota_used(&b.key).await), (Window::Day, Some(held)) => { Some(gov.quota_settle(&b.key, amount - held, at).await) } @@ -769,8 +775,6 @@ async fn consume_budgets( } })) .await; - // a reload between admission and settle may have dropped a held budget - hold.refund(gov).await; for (b, used) in budgets.iter().zip(used) { if let Some(used) = used && used >= b.limit @@ -1375,7 +1379,7 @@ mod tests { cost_micros: 60, }; let now = crate::epoch_secs(); - let reserve = || reserve_budgets(&state, &cfg, &ak, "", est, now); + let reserve = || reserve_budgets(&state, &cfg, &ak, "", || est, now); let (_, first) = reserve().await.unwrap().unwrap(); let (_, second) = reserve().await.unwrap().unwrap(); @@ -1404,11 +1408,42 @@ mod tests { ); } + #[tokio::test] + async fn a_settle_that_matches_its_hold_still_alerts_at_the_cap() { + let yaml = "listen: {host: h, port: 1}\ntenants: [{name: t1, key_daily_cost_quota_micros: 100}]\naccess_keys: [{ak: k1, tenant: t1, product: p, qps: 1, daily_token_quota: 1000}]"; + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let mut alerts = state.alerts.take_receiver().unwrap(); + let ak = state.auth.get(&access_key_id("k1")).await.unwrap(); + let spend = BudgetSpend { + tokens: 1, + cost_micros: 100, + }; + let (_, hold) = reserve_budgets(&state, &cfg, &ak, "", || spend, crate::epoch_secs()) + .await + .unwrap() + .unwrap(); + + consume_budgets(&state, &cfg, &ak, "", None, hold, spend).await; + + let alert = alerts.try_recv().expect("budget_exhausted alert"); + assert_eq!(alert.kind, "budget_exhausted"); + } + + #[tokio::test] + async fn without_a_day_budget_admission_never_prices_the_estimate() { + let (cfg, state, ak) = monthly_fixture(false).await; + reserve_budgets(&state, &cfg, &ak, "", || unreachable!("no day budget"), 0) + .await + .unwrap() + .unwrap(); + } + #[tokio::test] async fn monthly_budget_denies_at_the_cap_without_touching_daily_counters() { let (cfg, state, ak) = monthly_fixture(false).await; let gov = state.governance.as_ref(); - reserve_budgets(&state, &cfg, &ak, "", BudgetSpend::default(), 0) + reserve_budgets(&state, &cfg, &ak, "", BudgetSpend::default, 0) .await .unwrap() .unwrap(); @@ -1425,7 +1460,7 @@ mod tests { }, ) .await; - let err = reserve_budgets(&state, &cfg, &ak, "", BudgetSpend::default(), 0) + let err = reserve_budgets(&state, &cfg, &ak, "", BudgetSpend::default, 0) .await .unwrap() .unwrap_err(); @@ -1436,7 +1471,7 @@ mod tests { assert_eq!(gov.quota_used(&format!("cb:ak:{}", ak.ak_id)).await, 0); gov.quota_reset_all().await; assert!( - reserve_budgets(&state, &cfg, &ak, "", BudgetSpend::default(), 0) + reserve_budgets(&state, &cfg, &ak, "", BudgetSpend::default, 0) .await .unwrap() .is_err(), diff --git a/crates/state/src/governance.rs b/crates/state/src/governance.rs index 8b9aaf4..22b625c 100644 --- a/crates/state/src/governance.rs +++ b/crates/state/src/governance.rs @@ -36,7 +36,7 @@ pub trait Governance: Send + Sync + std::fmt::Debug { at_epoch_secs: i64, ) -> Option; /// Apply the settle delta (actual - reserved; negative refunds) to the day - /// bucket the paired reserve used (`at_epoch_secs`); returns the new total, or 0 for a zero delta. + /// bucket the paired reserve used (`at_epoch_secs`); returns the new total, or 0 when nothing was applied. async fn quota_settle(&self, key: &str, delta: i64, at_epoch_secs: i64) -> i64; /// Tokens spent today under `key`. async fn quota_used(&self, key: &str) -> i64; @@ -106,6 +106,9 @@ impl Governance for MemoryGovernance { self.quota.reserve(key, amount, limit).then_some(amount) } async fn quota_settle(&self, key: &str, delta: i64, _at: i64) -> i64 { + if delta == 0 { + return 0; + } self.quota.settle(key, delta) } async fn quota_used(&self, key: &str) -> i64 { @@ -348,7 +351,7 @@ fn quota_key(key: &str) -> String { /// Apply a settle delta and floor at 0 atomically (a rollover between reserve /// and settle must not plant a negative); keeps an existing TTL, arms `window` otherwise. -/// Returns the new total, 0 when the round-trip failed. +/// Returns the new total, or 0 when the window moved on or the round-trip failed. async fn settle_floored( conn: &redis::aio::ConnectionManager, key: &str, diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index 29af601..55385be 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -590,7 +590,7 @@ async fn realtime_gate( cfg, &ak, ak.attributed_user(hint), - admission::BudgetSpend::default(), + admission::BudgetSpend::default, gw_state::epoch_secs(), ) .await diff --git a/docs/governance.md b/docs/governance.md index f299648..3c573fe 100644 --- a/docs/governance.md +++ b/docs/governance.md @@ -269,10 +269,12 @@ enforced with a hard `400 service_quota_exceeded_exception` once exceeded. On the REST pipeline every daily budget holds the call's estimate at admission (the token estimate for the token cap; for the cost caps that estimate's prompt and `max_tokens` sides at the tenant's price), admitting while spent plus held -is under the cap, so concurrent calls overshoot by at most one; billing settles -the hold to the charged amount and a failed call returns it. A call without -`max_tokens` holds 256 output tokens, so a long completion can still carry one -call past the cap. Realtime turns and the monthly caps check before the call +is under the cap, so concurrent holds pass it by at most one hold; billing +settles the hold to the charged amount and a failed call returns it. A call that +costs more than it held still carries the total past the cap by the difference: +a call without `max_tokens` holds 256 output tokens, and unit- or +vendor-priced surfaces (images, video, decisions) hold no cost, so the cost caps +only check them. Realtime turns and the monthly caps check before the call and accrue after it ("soft": concurrent turns can overshoot by one each). The per-(key, model) quota is soft the same way: it is a routing trigger (the fallback-model degrade), while the reserved per-key daily quota hard-caps spend. From 4a56d6ccddc9e743156c790bedb2c25b4e8330d0 Mon Sep 17 00:00:00 2001 From: CMGS Date: Sat, 17 Oct 2026 19:37:07 +0800 Subject: [PATCH 07/22] billing: settle an abandoned stream before shutdown, bounded and off the latency - The generation lookup slept up to ~38 s in a detached task that graceful shutdown never waited for, so a restart dropped the row of every OpenRouter stream abandoned in that window. Abort billing now holds a SettleGuard on a ledger TaskTracker; flush_billing cancels pending lookups, which bill the estimate at once, and waits for those settles before the final flush. - Each lookup is capped at 10 s, so a hung /generation call cannot stretch a settle to four account timeouts. - The access log's latency_ms stops when the stream ends (DagContext::served_at) instead of counting the post-stream lookup. - The polling loop moves into engines::generation, leaving the DAG settle provider-neutral; the account check uses wire_kind(), so a raw accounts: entry with provider openrouter is looked up too. --- Cargo.lock | 2 + Cargo.toml | 1 + crates/dag/src/context.rs | 2 + crates/dag/src/nodes.rs | 47 ++++++---------- crates/engines/Cargo.toml | 1 + crates/engines/src/generation.rs | 93 +++++++++++++++++++++++++++++--- crates/handler/src/lib.rs | 38 +++++++++++++ crates/state/Cargo.toml | 1 + crates/state/src/admission.rs | 33 +++++++++++- crates/state/src/lib.rs | 4 ++ crates/views/src/lib.rs | 4 +- docs/observability.md | 8 +-- 12 files changed, 189 insertions(+), 45 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 2e77991..da0db2f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1018,6 +1018,7 @@ dependencies = [ "sqlx", "tempfile", "tokio", + "tokio-util", "tracing", ] @@ -2847,6 +2848,7 @@ dependencies = [ "bytes", "futures-core", "futures-sink", + "futures-util", "pin-project-lite", "tokio", ] diff --git a/Cargo.toml b/Cargo.toml index f759d7f..888fc00 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -35,6 +35,7 @@ undocumented_unsafe_blocks = "deny" [workspace.dependencies] # --- async runtime --- tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "sync", "net", "signal", "io-util"] } +tokio-util = { version = "0.7", features = ["rt"] } async-trait = "0.1" futures = "0.3" # --- web (server) --- diff --git a/crates/dag/src/context.rs b/crates/dag/src/context.rs index b665539..4eeff60 100644 --- a/crates/dag/src/context.rs +++ b/crates/dag/src/context.rs @@ -45,6 +45,7 @@ pub struct DagContext { pub budget_hold: BudgetHold, /// Outbound DLP buffered this stream; billing waits for the view's delivery result. pub billing_deferred: bool, + pub served_at: Option, } impl DagContext { @@ -75,6 +76,7 @@ impl DagContext { user_budget: None, budget_hold: BudgetHold::default(), billing_deferred: false, + served_at: None, } } diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index 7814043..3d02347 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -13,13 +13,6 @@ use crate::executor::{DagNode, Layer}; const DEFAULT_COMPLETION_RESERVE: i64 = 256; /// Reservation cap, so a hostile `max_tokens` cannot overflow the estimate or the counter. const MAX_RESERVE: i64 = 1_000_000; -/// Waits before each lookup of an abandoned stream's recorded usage; OpenRouter records one in about 5 s. -const GENERATION_LOOKUP_DELAYS: [std::time::Duration; 4] = [ - std::time::Duration::from_secs(3), - std::time::Duration::from_secs(5), - std::time::Duration::from_secs(10), - std::time::Duration::from_secs(20), -]; /// preprocess/model_quota: per-(AK, model) daily cap (AK override, else tenant default); /// over it degrades to the tenant fallback model or denies, so it runs before resolve_model. @@ -736,7 +729,7 @@ impl BillTokens { } } -/// Close a buffered stream's billing: refund before delivery; after a partial one bill the upstream's usage, or estimate without it. +/// Close a buffered stream's billing: refund before delivery, else bill its usage or an estimate. pub async fn settle_deferred_stream(ctx: &mut DagContext, delivery: StreamDelivery) -> GResult<()> { if !ctx.billing_deferred { return Ok(()); @@ -777,7 +770,21 @@ async fn bill_aborted_stream( ctx: &mut DagContext, delivered_completion: Option, ) -> GResult<()> { - if let Some(usage) = recorded_usage(ctx).await { + let settle = ctx.state.begin_settle(); + ctx.served_at = Some(std::time::Instant::now()); + let usage = match (ctx.request.account.as_deref(), ctx.outcome.as_ref()) { + (Some(account), Some(outcome)) => { + gw_engines::generation::recorded_usage( + ctx.transport.as_ref(), + account, + &outcome.response.upstream_id, + settle.closing(), + ) + .await + } + _ => None, + }; + if let Some(usage) = usage { let outcome = ctx .outcome .as_mut() @@ -822,28 +829,6 @@ async fn bill_aborted_stream( bill(ctx, BillTokens::weighted(prompt, completion, &rate), true).await } -/// OpenRouter's record of an abandoned stream's usage, polled until it lands. -async fn recorded_usage(ctx: &DagContext) -> Option { - let account = ctx.request.account.as_deref()?; - let id = ctx.outcome.as_ref()?.response.upstream_id.as_str(); - if !gw_engines::generation::records_generation(account, id) { - return None; - } - for delay in GENERATION_LOOKUP_DELAYS { - tokio::time::sleep(delay).await; - match gw_engines::generation::openrouter_usage(ctx.transport.as_ref(), account, id).await { - Ok(Some(usage)) => return Some(usage), - Ok(None) => {} - Err(e) => tracing::warn!(error = %e, id, "generation usage lookup failed"), - } - } - tracing::warn!( - id, - "no recorded usage for an aborted stream; billing the estimate" - ); - None -} - fn blank_stream_response(response: &mut gw_models::GatewayResponse) { response.message.clear(); response.completion_tokens = 0; diff --git a/crates/engines/Cargo.toml b/crates/engines/Cargo.toml index 724c26e..b09c5c6 100644 --- a/crates/engines/Cargo.toml +++ b/crates/engines/Cargo.toml @@ -29,6 +29,7 @@ tokio = { workspace = true } [dev-dependencies] # real-socket wire test for HttpTransport (loopback only, never egress) axum = { workspace = true } +tokio = { workspace = true, features = ["test-util"] } [lints] workspace = true diff --git a/crates/engines/src/generation.rs b/crates/engines/src/generation.rs index 2ea3d0b..5e46b4a 100644 --- a/crates/engines/src/generation.rs +++ b/crates/engines/src/generation.rs @@ -1,5 +1,7 @@ //! Post-hoc usage of an OpenRouter generation, which OpenRouter records even for a cancelled stream. +use std::time::Duration; + use gw_consts::Protocol; use gw_models::{Account, GResult, GatewayError}; use serde_json::{Value, json}; @@ -8,10 +10,52 @@ use crate::base::versioned_url; use crate::transport::{Transport, UpstreamBody, UpstreamRequest}; const OPENROUTER_KIND: &str = "openrouter"; +/// Waits before each lookup of an abandoned stream's recorded usage; OpenRouter records one in about 5 s. +const LOOKUP_DELAYS: [Duration; 4] = [ + Duration::from_secs(3), + Duration::from_secs(5), + Duration::from_secs(10), + Duration::from_secs(20), +]; +const LOOKUP_TIMEOUT: Duration = Duration::from_secs(10); + +pub async fn recorded_usage( + transport: &dyn Transport, + account: &Account, + id: &str, + closing: impl Future, +) -> Option { + if !records_generation(account, id) { + return None; + } + tokio::pin!(closing); + for delay in LOOKUP_DELAYS { + let lookup = async { + tokio::time::sleep(delay).await; + tokio::time::timeout(LOOKUP_TIMEOUT, openrouter_usage(transport, account, id)).await + }; + tokio::select! { + () = &mut closing => { + tracing::warn!(id, "shutting down before the generation record; billing the estimate"); + return None; + } + found = lookup => match found { + Ok(Ok(Some(usage))) => return Some(usage), + Ok(Ok(None)) => {} + Ok(Err(e)) => tracing::warn!(error = %e, id, "generation usage lookup failed"), + Err(_) => tracing::warn!(id, "generation usage lookup timed out"), + }, + } + } + tracing::warn!( + id, + "no recorded usage for an aborted stream; billing the estimate" + ); + None +} -/// Whether `account` records generation `id` for a later [`openrouter_usage`] lookup. -pub fn records_generation(account: &Account, id: &str) -> bool { - account.kind == OPENROUTER_KIND +fn records_generation(account: &Account, id: &str) -> bool { + account.wire_kind() == OPENROUTER_KIND && !account.endpoint.is_empty() && !account.api_key_env.is_empty() && !id.is_empty() @@ -21,7 +65,7 @@ pub fn records_generation(account: &Account, id: &str) -> bool { } /// The generation's native counts and cost as chat-wire usage; `None` until OpenRouter records it. -pub async fn openrouter_usage( +async fn openrouter_usage( transport: &dyn Transport, account: &Account, id: &str, @@ -33,10 +77,7 @@ pub async fn openrouter_usage( .send(UpstreamRequest { protocol: Protocol::OpenaiChat, method: "GET", - url: versioned_url( - account.endpoint.trim_end_matches('/'), - &format!("generation?id={id}"), - ), + url: versioned_url(account.base_url(""), &format!("generation?id={id}")), headers: vec![("authorization", format!("Bearer {key}"))], body: Vec::new(), stream: false, @@ -164,6 +205,36 @@ mod tests { ); } + #[derive(Debug, Default)] + struct Silent(std::sync::atomic::AtomicUsize); + + #[async_trait::async_trait] + impl Transport for Silent { + async fn send(&self, _: UpstreamRequest) -> GResult { + self.0.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + std::future::pending().await + } + } + + #[tokio::test(start_paused = true)] + async fn a_lookup_that_never_answers_times_out_each_try() { + let t = Silent::default(); + let usage = recorded_usage(&t, &openrouter(), "gen-1", std::future::pending()).await; + assert!(usage.is_none()); + assert_eq!( + t.0.load(std::sync::atomic::Ordering::SeqCst), + LOOKUP_DELAYS.len() + ); + } + + #[tokio::test(start_paused = true)] + async fn a_closing_gateway_skips_the_lookup() { + let t = Silent::default(); + let usage = recorded_usage(&t, &openrouter(), "gen-1", std::future::ready(())).await; + assert!(usage.is_none()); + assert_eq!(t.0.load(std::sync::atomic::Ordering::SeqCst), 0); + } + #[test] fn only_an_openrouter_account_with_a_plain_id_records_generations() { assert!(records_generation( @@ -175,6 +246,12 @@ mod tests { ..openrouter() }; assert!(!records_generation(&other, "gen-1")); + let raw = Account { + kind: String::new(), + provider: OPENROUTER_KIND.into(), + ..openrouter() + }; + assert!(records_generation(&raw, "gen-1")); assert!(!records_generation(&openrouter(), "gen/../x")); assert!(!records_generation(&openrouter(), "")); } diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index d8e28f9..cd3fda3 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -2777,6 +2777,7 @@ mod tests { .await .unwrap(); + assert!(ctx.served_at.is_some()); let response = ctx.outcome.expect("outcome").response; assert!(response.aborted); assert_eq!( @@ -2809,6 +2810,43 @@ mod tests { assert_eq!(row.vendor_cost_micros, 500, "{row:?}"); } + #[tokio::test] + async fn shutdown_bills_a_pending_aborted_stream_before_the_flush_returns() { + let yaml = "listen: {host: h, port: 1}\nproviders: [{name: openrouter, kind: openrouter, api_key_env: PATH}]\nmodels: [{name: m, provider: openrouter, input_price_per_1k_micros: 1000, output_price_per_1k_micros: 2000}]\ntenants: [{name: t1}]\naccess_keys: [{ak: or-ak, tenant: t1, product: p, qps: 100, daily_token_quota: 100000}]"; + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let transport = Arc::new(RecordedGenerationStream::default()); + let h = OnlineHandler::new( + gw_state::SharedConfig::new(cfg, state.clone()), + transport.clone(), + ); + let ak = h.state().auth.get(&access_key_id("or-ak")).await.unwrap(); + let running = tokio::spawn({ + let h = h.clone(); + async move { + h.run(drained_stream_req("m", "count to a thousand"), ak) + .await + } + }); + tokio::time::sleep(std::time::Duration::from_millis(200)).await; + + tokio::time::timeout( + std::time::Duration::from_secs(2), + gw_state::admission::flush_billing(&state), + ) + .await + .expect("the flush waits only for the cancelled lookup"); + + let (_, ledger) = state.store.ledger_snapshot(usize::MAX).await.unwrap(); + assert_eq!(ledger.len(), 1, "billed before the flush returned"); + assert!(ledger[0].estimated, "{:?}", ledger[0]); + assert_eq!( + transport.lookups.load(std::sync::atomic::Ordering::SeqCst), + 0 + ); + running.await.unwrap().unwrap(); + } + #[derive(Debug)] struct PiiStream; diff --git a/crates/state/Cargo.toml b/crates/state/Cargo.toml index 3e1fd8e..e371a32 100644 --- a/crates/state/Cargo.toml +++ b/crates/state/Cargo.toml @@ -21,6 +21,7 @@ async-trait = { workspace = true } serde_json = { workspace = true } serde = { workspace = true } tokio = { workspace = true } +tokio-util = { workspace = true } futures = { workspace = true } metrics = { workspace = true } chacha20poly1305 = { workspace = true } diff --git a/crates/state/src/admission.rs b/crates/state/src/admission.rs index a7ac420..5d63693 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -13,6 +13,8 @@ use std::time::Duration; use gw_config::{GatewayConfig, TenantConf}; use gw_models::GResult; use tokio::sync::{Mutex, mpsc, oneshot}; +use tokio_util::sync::{CancellationToken, WaitForCancellationFuture}; +use tokio_util::task::{TaskTracker, task_tracker::TaskTrackerToken}; use crate::store::{BillingInput, BillingRecord, Store, billing_record}; use crate::{AkInfo, GatewayState, Governance, UserBudget, clamp_tokens}; @@ -109,11 +111,25 @@ pub struct Settled { pub cost_micros: i64, } +/// A settle in progress: shutdown waits for it to drop and fires `closing` so it bills without waiting. +pub struct SettleGuard { + _tracked: TaskTrackerToken, + closing: CancellationToken, +} + +impl SettleGuard { + pub fn closing(&self) -> WaitForCancellationFuture<'_> { + self.closing.cancelled() + } +} + #[derive(Debug, Clone)] pub(crate) struct BillingLedger { store: Arc, queue: Option>, deferred: bool, + settles: TaskTracker, + closing: CancellationToken, } impl BillingLedger { @@ -122,6 +138,8 @@ impl BillingLedger { store, queue: None, deferred: false, + settles: TaskTracker::new(), + closing: CancellationToken::new(), } } @@ -175,6 +193,15 @@ impl BillingLedger { store, queue: Some(queue), deferred, + settles: TaskTracker::new(), + closing: CancellationToken::new(), + } + } + + pub(crate) fn begin_settle(&self) -> SettleGuard { + SettleGuard { + _tracked: self.settles.token(), + closing: self.closing.clone(), } } @@ -379,7 +406,6 @@ impl LedgerSpool { } } -/// One budget: its window, governance counter and cap. struct Budget { scope: BudgetScope, window: Window, @@ -471,8 +497,11 @@ impl LedgerWrite { } } -/// Commit billing rows the batching writer still holds, before the process exits. +/// Finish in-flight settles, then commit the rows the batching writer holds, before the process exits. pub async fn flush_billing(state: &GatewayState) { + state.billing.closing.cancel(); + state.billing.settles.close(); + state.billing.settles.wait().await; let dropped = state.billing.flush().await; if dropped > 0 { tracing::error!(rows = dropped, "billing rows dropped before shutdown"); diff --git a/crates/state/src/lib.rs b/crates/state/src/lib.rs index dcb97d7..51b04d1 100644 --- a/crates/state/src/lib.rs +++ b/crates/state/src/lib.rs @@ -893,6 +893,10 @@ impl GatewayState { streams: prev.streams.clone(), }) } + + pub fn begin_settle(&self) -> admission::SettleGuard { + self.billing.begin_settle() + } } /// A consistent (config, derived-state) pair. A request loads one snapshot and diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index 55385be..ea90ca0 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -1233,7 +1233,9 @@ fn log_access(surface: &str, ctx: &DagContext, started: Instant) { ) }) .unwrap_or_default(); - let latency = started.elapsed(); + let latency = ctx + .served_at + .map_or_else(|| started.elapsed(), |at| at.duration_since(started)); let user_id = ctx.effective_user_id(); let decisions = ctx.decisions_line(); let ak_id = &*ctx.ak.ak_id; diff --git a/docs/observability.md b/docs/observability.md index 6a13ea4..f59f79b 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -89,9 +89,11 @@ the unit-priced surfaces; also summed in `/admin/usage`), charged when counts came from an aborted stream rather than a vendor usage payload). An aborted stream served by an OpenRouter account is billed from OpenRouter's record of the generation instead (`GET /api/v1/generation`, polled for about -40 s; cancelled generations included), with cache reads at their weight and -OpenRouter's cost as `vendor_cost_micros`; only when no record arrives does it -fall back to the estimate. +40 s with each lookup capped at 10 s; cancelled generations included), with cache +reads at their weight and OpenRouter's cost as `vendor_cost_micros`; only when no +record arrives, or the gateway starts shutting down, does it fall back to the +estimate. Shutdown waits for these settles before the final ledger flush, and the +access log's `latency_ms` stops at the end of the stream, not of the lookup. Writes are idempotent by `request_id`; a transient store failure is held in a bounded repair queue and retried with backoff — eight attempts, about 13 s — before the batch is dropped and counted in `gateway_ledger_write_failures_total`. From cad7ec841a68e0d01c90605a0c3f00adcd8af67f Mon Sep 17 00:00:00 2001 From: CMGS Date: Sat, 17 Oct 2026 19:51:45 +0800 Subject: [PATCH 08/22] simplify: strip NUL at the client hint, one estimate total, leaner lookup - NUL is stripped where every client user enters (cap_user_hint) instead of at the SQL ledger insert: user_id is the only client-reachable billing column, and the same user also lands in security events, retained content and batch rows, which Postgres refused just the same. - reserve_total replaces three copies of the estimate cap; AkTpmLimit reads the parts QuotaCheck cached instead of re-walking the prompt. - usage_rollup_since takes Option, so the late-row backfill and its blocking rollup lock are explicit instead of keyed off i64::MAX. - The generation lookup resolves the key once, reuses parse_json_reply, and takes a SettleGuard only for streams that actually wait on a record. - The spool replay tick delays instead of bursting after a slow replay; BillingLedger::repairing builds on direct; shared test fixtures. --- crates/dag/src/nodes.rs | 39 ++++++++-------- crates/engines/src/generation.rs | 47 +++++++------------ crates/handler/src/lib.rs | 26 ++++++----- crates/state/src/admission.rs | 44 ++++++++++-------- crates/state/src/store.rs | 79 ++++++++++---------------------- crates/views/src/lib.rs | 17 +++++-- 6 files changed, 109 insertions(+), 143 deletions(-) diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index 3d02347..a19559c 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -240,8 +240,7 @@ impl Write for PromptBytes { } /// Admission estimate: ~bytes/4 prompt plus the requested max_tokens, saturating and capped. -fn reserve_estimate(req: &gw_models::GatewayRequest) -> i64 { - let (prompt, max_out) = reserve_parts(req); +fn reserve_total((prompt, max_out): (i64, i64)) -> i64 { prompt.saturating_add(max_out).min(MAX_RESERVE) } @@ -274,8 +273,7 @@ impl DagNode for QuotaCheck { } async fn execute(&self, ctx: &mut DagContext) -> GResult<()> { ctx.reserve_parts = reserve_parts(&ctx.request); - let (prompt, max_out) = ctx.reserve_parts; - let est = prompt.saturating_add(max_out).min(MAX_RESERVE); + let est = reserve_total(ctx.reserve_parts); let at = gw_state::epoch_secs(); let reserved = admission::reserve_daily(ctx.state.governance.as_ref(), &ctx.ak, est, at) .await @@ -412,10 +410,7 @@ impl DagNode for AkTpmLimit { "ak_tpm" } async fn execute(&self, ctx: &mut DagContext) -> GResult<()> { - let est = ctx - .quota_reserved - .filter(|&reserved| reserved > 0) - .unwrap_or_else(|| reserve_estimate(&ctx.request)); + let est = reserve_total(ctx.reserve_parts); ctx.tpm_reserved = admission::reserve_tpm(ctx.state.governance.as_ref(), &ctx.ak, est) .await .map_err(key_limit_denied)?; @@ -423,7 +418,7 @@ impl DagNode for AkTpmLimit { } } -/// model_access/budget: holds the call's estimate against the daily budgets (tenant, key, user) and checks the monthly ones. +/// model_access/budget: hold the estimate against the day budgets, check the month ones. pub struct BudgetGate; #[async_trait::async_trait] @@ -435,7 +430,7 @@ impl DagNode for BudgetGate { let (prompt, max_out) = ctx.reserve_parts; let (cfg, request, tenant) = (&ctx.cfg, &ctx.request, &ctx.ak.tenant); let est = || admission::BudgetSpend { - tokens: prompt.saturating_add(max_out).min(MAX_RESERVE), + tokens: reserve_total((prompt, max_out)), cost_micros: gw_models::cost_micros( prompt, max_out, @@ -770,25 +765,27 @@ async fn bill_aborted_stream( ctx: &mut DagContext, delivered_completion: Option, ) -> GResult<()> { - let settle = ctx.state.begin_settle(); ctx.served_at = Some(std::time::Instant::now()); + let mut settle = None; let usage = match (ctx.request.account.as_deref(), ctx.outcome.as_ref()) { - (Some(account), Some(outcome)) => { + (Some(account), Some(outcome)) + if gw_engines::generation::records_generation( + account, + &outcome.response.upstream_id, + ) => + { + let guard = settle.insert(ctx.state.begin_settle()); gw_engines::generation::recorded_usage( ctx.transport.as_ref(), account, &outcome.response.upstream_id, - settle.closing(), + guard.closing(), ) .await } _ => None, }; - if let Some(usage) = usage { - let outcome = ctx - .outcome - .as_mut() - .ok_or_else(|| GatewayError::internal("aborted stream without an outcome"))?; + if let (Some(usage), Some(outcome)) = (usage, ctx.outcome.as_mut()) { let common = gw_engines::extract_common_usage(&usage, false); if let Some(u) = &common { let r = &mut outcome.response; @@ -842,7 +839,7 @@ async fn bill_usage(ctx: &mut DagContext) -> GResult<()> { let resp = &ctx .outcome .as_ref() - .ok_or_else(|| GatewayError::internal("cost_calc without an engine outcome"))? + .ok_or_else(|| GatewayError::internal("billing without an engine outcome"))? .response; let rate = gw_state::model_token_rate(&ctx.cfg, served_model(ctx.request.model_param_v2.as_ref())); @@ -1156,7 +1153,7 @@ mod tests { } #[test] - fn reserve_estimate_saturates_on_hostile_max_tokens() { + fn reserve_total_saturates_on_hostile_max_tokens() { use gw_models::params::ChatParams; use gw_models::{ChatMsg, GatewayRequest, ModelParamV2, TypedParams}; let mut param = ModelParamV2::with_name(gw_consts::Protocol::OpenaiChat, "m"); @@ -1169,7 +1166,7 @@ mod tests { model_param_v2: Some(param), ..Default::default() }; - let est = super::reserve_estimate(&req); + let est = super::reserve_total(super::reserve_parts(&req)); assert!(est > 0 && est <= super::MAX_RESERVE); } } diff --git a/crates/engines/src/generation.rs b/crates/engines/src/generation.rs index 5e46b4a..0251051 100644 --- a/crates/engines/src/generation.rs +++ b/crates/engines/src/generation.rs @@ -3,11 +3,11 @@ use std::time::Duration; use gw_consts::Protocol; -use gw_models::{Account, GResult, GatewayError}; +use gw_models::{Account, GResult}; use serde_json::{Value, json}; -use crate::base::versioned_url; -use crate::transport::{Transport, UpstreamBody, UpstreamRequest}; +use crate::base::{parse_json_reply, versioned_url}; +use crate::transport::{Transport, UpstreamRequest}; const OPENROUTER_KIND: &str = "openrouter"; /// Waits before each lookup of an abandoned stream's recorded usage; OpenRouter records one in about 5 s. @@ -25,14 +25,16 @@ pub async fn recorded_usage( id: &str, closing: impl Future, ) -> Option { - if !records_generation(account, id) { - return None; - } + let key = account.api_key()?; tokio::pin!(closing); for delay in LOOKUP_DELAYS { let lookup = async { tokio::time::sleep(delay).await; - tokio::time::timeout(LOOKUP_TIMEOUT, openrouter_usage(transport, account, id)).await + tokio::time::timeout( + LOOKUP_TIMEOUT, + openrouter_usage(transport, account, id, &key), + ) + .await }; tokio::select! { () = &mut closing => { @@ -54,10 +56,9 @@ pub async fn recorded_usage( None } -fn records_generation(account: &Account, id: &str) -> bool { +pub fn records_generation(account: &Account, id: &str) -> bool { account.wire_kind() == OPENROUTER_KIND && !account.endpoint.is_empty() - && !account.api_key_env.is_empty() && !id.is_empty() && id .bytes() @@ -69,10 +70,8 @@ async fn openrouter_usage( transport: &dyn Transport, account: &Account, id: &str, + key: &str, ) -> GResult> { - let Some(key) = account.api_key() else { - return Ok(None); - }; let resp = transport .send(UpstreamRequest { protocol: Protocol::OpenaiChat, @@ -89,24 +88,10 @@ async fn openrouter_usage( if resp.status == 404 { return Ok(None); } - let UpstreamBody::Json(body) = resp.body else { - return Err(lookup_failed("generation lookup returned no JSON body")); - }; - if resp.status != 200 { - return Err(lookup_failed(format!( - "generation lookup returned {}", - resp.status - ))); - } - let doc: Value = serde_json::from_slice(&body) - .map_err(|e| lookup_failed(format!("generation lookup body: {e}")))?; + let (_, doc) = parse_json_reply(resp)?; Ok(generation_usage(&doc["data"])) } -fn lookup_failed(message: impl Into) -> GatewayError { - GatewayError::new(gw_consts::ErrCode::FED_RESP_RPC_FAILED, 502, message) -} - /// Map `/generation` data to chat-wire usage, reasoning inside completion as in `completion_tokens`. fn generation_usage(data: &Value) -> Option { let count = |k: &str| data[k].as_i64(); @@ -125,7 +110,7 @@ mod tests { use std::sync::Mutex; use super::*; - use crate::transport::UpstreamResponse; + use crate::transport::{UpstreamBody, UpstreamResponse}; #[derive(Debug)] struct Canned(u16, &'static str, Mutex>); @@ -162,7 +147,7 @@ mod tests { r#"{"data":{"cancelled":true,"native_tokens_prompt":11,"native_tokens_completion":73,"native_tokens_cached":4,"native_tokens_reasoning":73,"total_cost":9.09e-05}}"#, Mutex::default(), ); - let usage = openrouter_usage(&t, &openrouter(), "gen-1792232109-kUA52UZsI2") + let usage = openrouter_usage(&t, &openrouter(), "gen-1792232109-kUA52UZsI2", "k") .await .unwrap() .unwrap(); @@ -187,7 +172,7 @@ mod tests { async fn an_unrecorded_generation_yields_nothing() { let missing = Canned(404, r#"{"error":{"code":404}}"#, Mutex::default()); assert!( - openrouter_usage(&missing, &openrouter(), "gen-1") + openrouter_usage(&missing, &openrouter(), "gen-1", "k") .await .unwrap() .is_none() @@ -198,7 +183,7 @@ mod tests { Mutex::default(), ); assert!( - openrouter_usage(&pending, &openrouter(), "gen-1") + openrouter_usage(&pending, &openrouter(), "gen-1", "k") .await .unwrap() .is_none() diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index cd3fda3..99516b7 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -2763,14 +2763,23 @@ mod tests { } } - #[tokio::test(start_paused = true)] - async fn an_aborted_openrouter_stream_bills_the_recorded_generation() { + async fn openrouter_handler() -> ( + OnlineHandler, + Arc, + Arc, + ) { let yaml = "listen: {host: h, port: 1}\nproviders: [{name: openrouter, kind: openrouter, api_key_env: PATH}]\nmodels: [{name: m, provider: openrouter, input_price_per_1k_micros: 1000, output_price_per_1k_micros: 2000, token_rate: {read_cache: 0.1}}]\ntenants: [{name: t1}]\naccess_keys: [{ak: or-ak, tenant: t1, product: p, qps: 100, daily_token_quota: 100000}]"; let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); let state = Arc::new(GatewayState::from_config(&cfg)); let transport = Arc::new(RecordedGenerationStream::default()); let h = OnlineHandler::new(gw_state::SharedConfig::new(cfg, state), transport.clone()); let ak = h.state().auth.get(&access_key_id("or-ak")).await.unwrap(); + (h, transport, ak) + } + + #[tokio::test(start_paused = true)] + async fn an_aborted_openrouter_stream_bills_the_recorded_generation() { + let (h, transport, ak) = openrouter_handler().await; let ctx = h .run(drained_stream_req("m", "count to a thousand"), ak) @@ -2810,17 +2819,10 @@ mod tests { assert_eq!(row.vendor_cost_micros, 500, "{row:?}"); } - #[tokio::test] + #[tokio::test(start_paused = true)] async fn shutdown_bills_a_pending_aborted_stream_before_the_flush_returns() { - let yaml = "listen: {host: h, port: 1}\nproviders: [{name: openrouter, kind: openrouter, api_key_env: PATH}]\nmodels: [{name: m, provider: openrouter, input_price_per_1k_micros: 1000, output_price_per_1k_micros: 2000}]\ntenants: [{name: t1}]\naccess_keys: [{ak: or-ak, tenant: t1, product: p, qps: 100, daily_token_quota: 100000}]"; - let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); - let state = Arc::new(GatewayState::from_config(&cfg)); - let transport = Arc::new(RecordedGenerationStream::default()); - let h = OnlineHandler::new( - gw_state::SharedConfig::new(cfg, state.clone()), - transport.clone(), - ); - let ak = h.state().auth.get(&access_key_id("or-ak")).await.unwrap(); + let (h, transport, ak) = openrouter_handler().await; + let state = h.state(); let running = tokio::spawn({ let h = h.clone(); async move { diff --git a/crates/state/src/admission.rs b/crates/state/src/admission.rs index 5d63693..98e519f 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -154,6 +154,7 @@ impl BillingLedger { // detached: replays the spool for the life of the process tokio::spawn(async move { let mut tick = tokio::time::interval(LEDGER_REPLAY_INTERVAL); + tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); loop { tick.tick().await; spool.replay(&replay_store).await; @@ -190,11 +191,9 @@ impl BillingLedger { } }); Self { - store, queue: Some(queue), deferred, - settles: TaskTracker::new(), - closing: CancellationToken::new(), + ..Self::direct(store) } } @@ -332,7 +331,7 @@ impl LedgerSpool { } let path = self.path.clone(); let _held = self.lock.lock().await; - blocking(move || { + run_blocking(move || { if let Some(dir) = path.parent() { std::fs::create_dir_all(dir)?; } @@ -348,7 +347,7 @@ impl LedgerSpool { let claimed = { let (path, replaying) = (self.path.clone(), self.replaying.clone()); let _held = self.lock.lock().await; - blocking(move || { + run_blocking(move || { if replaying.exists() { return Ok(true); } @@ -367,10 +366,12 @@ impl LedgerSpool { return; } } - let (mut offset, mut replayed, mut since) = (0, 0, i64::MAX); + let (mut offset, mut replayed, mut since) = (0, 0, None); loop { let replaying = self.replaying.clone(); - let (rows, next) = match blocking(move || read_spool_batch(&replaying, offset)).await { + let (rows, next) = match run_blocking(move || read_spool_batch(&replaying, offset)) + .await + { Ok(batch) => batch, Err(e) => { tracing::error!(error = %e, path = %self.replaying.display(), "billing spool read failed"); @@ -388,17 +389,18 @@ impl LedgerSpool { since = rows .iter() .map(|r| r.created_at_epoch_secs) - .fold(since, i64::min); + .chain(since) + .min(); offset = next; } - if replayed > 0 + if since.is_some() && let Err(e) = store.usage_rollup_since(since, crate::epoch_secs()).await { tracing::warn!(error = %e, rows = replayed, "billing spool rollup deferred"); return; } let replaying = self.replaying.clone(); - if let Err(e) = blocking(move || std::fs::remove_file(replaying)).await { + if let Err(e) = run_blocking(move || std::fs::remove_file(replaying)).await { tracing::error!(error = %e, path = %self.replaying.display(), "billing spool remove failed"); return; } @@ -508,7 +510,7 @@ pub async fn flush_billing(state: &GatewayState) { } } -/// Hold the estimate `est` yields against the day budgets, admitting while spent plus held is under each cap; month budgets and a zero estimate only check. +/// Hold the call's estimate against each day budget while spent plus held stays under its cap; month budgets only check. pub async fn reserve_budgets( state: &GatewayState, cfg: &GatewayConfig, @@ -945,7 +947,7 @@ fn read_spool_batch(path: &Path, offset: u64) -> std::io::Result<(Vec(f: F) -> std::io::Result +async fn run_blocking(f: F) -> std::io::Result where T: Send + 'static, F: FnOnce() -> std::io::Result + Send + 'static, @@ -1339,6 +1341,14 @@ mod tests { assert_eq!(previous_month((2026, 9)), (2026, 8)); } + async fn daily_fixture() -> (Arc, Arc, Arc) { + let yaml = "listen: {host: h, port: 1}\ntenants: [{name: t1, key_daily_cost_quota_micros: 100}]\naccess_keys: [{ak: k1, tenant: t1, product: p, qps: 1, daily_token_quota: 1000}]"; + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let ak = state.auth.get(&access_key_id("k1")).await.unwrap(); + (cfg, state, ak) + } + async fn monthly_fixture( rollover: bool, ) -> (Arc, Arc, Arc) { @@ -1397,10 +1407,7 @@ mod tests { #[tokio::test] async fn daily_budgets_hold_in_flight_calls_until_they_settle() { - let yaml = "listen: {host: h, port: 1}\ntenants: [{name: t1, key_daily_cost_quota_micros: 100}]\naccess_keys: [{ak: k1, tenant: t1, product: p, qps: 1, daily_token_quota: 1000}]"; - let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); - let state = Arc::new(GatewayState::from_config(&cfg)); - let ak = state.auth.get(&access_key_id("k1")).await.unwrap(); + let (cfg, state, ak) = daily_fixture().await; let gov = state.governance.as_ref(); let key = format!("cb:ak:{}", ak.ak_id); let est = BudgetSpend { @@ -1439,11 +1446,8 @@ mod tests { #[tokio::test] async fn a_settle_that_matches_its_hold_still_alerts_at_the_cap() { - let yaml = "listen: {host: h, port: 1}\ntenants: [{name: t1, key_daily_cost_quota_micros: 100}]\naccess_keys: [{ak: k1, tenant: t1, product: p, qps: 1, daily_token_quota: 1000}]"; - let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); - let state = Arc::new(GatewayState::from_config(&cfg)); + let (cfg, state, ak) = daily_fixture().await; let mut alerts = state.alerts.take_receiver().unwrap(); - let ak = state.auth.get(&access_key_id("k1")).await.unwrap(); let spend = BudgetSpend { tokens: 1, cost_micros: 100, diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index d1a59f8..887f32a 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -4,7 +4,6 @@ //! [`SqliteStore`] (`storage.sqlite_path`, one durable node), and //! [`PostgresStore`] (`storage.postgres_url`, shared across a fleet). -use std::borrow::Cow; use std::collections::{BTreeMap, HashMap, HashSet}; use std::sync::Mutex; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -389,11 +388,11 @@ pub trait Store: Send + Sync + std::fmt::Debug { bucket_secs: i64, ) -> GResult>; async fn usage_rollup_advance(&self, now_epoch_secs: i64) -> GResult { - self.usage_rollup_since(i64::MAX, now_epoch_secs).await + self.usage_rollup_since(None, now_epoch_secs).await } /// Recompute and upsert every rollup bucket in the trailing backfill window, or back to `since` /// for rows billed late; buckets are never deleted, so a missed tick self-heals. Returns the buckets written. - async fn usage_rollup_since(&self, since_epoch_secs: i64, now_epoch_secs: i64) -> GResult; + async fn usage_rollup_since(&self, since: Option, now_epoch_secs: i64) -> GResult; /// Delete rolled billing rows past `ledger_max_rows`; returns rows deleted. async fn ledger_prune(&self) -> GResult { Ok(0) @@ -544,7 +543,7 @@ pub struct MemoryStore { #[cfg(test)] ledger_failures: AtomicUsize, /// Minute buckets keyed by (minute, tenant, user, model); see - /// [`Store::usage_rollup_advance`]. + /// [`Store::usage_rollup_since`]. rollup: Mutex>, rollup_max_buckets: Option, sec_events: Mutex>, @@ -740,13 +739,10 @@ impl Store for MemoryStore { Ok(map.into_iter().collect()) } - async fn usage_rollup_since(&self, since: i64, now: i64) -> GResult { + async fn usage_rollup_since(&self, since: Option, now: i64) -> GResult { let hi = bucket_floor(now - ROLLUP_SETTLE_SECS); let mut rollup = lock(&self.rollup); - let watermark = rollup_watermark(&rollup); - let lo = (hi - ROLLUP_BACKFILL_SECS) - .min(watermark) - .min(bucket_floor(since)); + let lo = rollup_floor(hi, rollup_watermark(&rollup), since); let mut fresh = BTreeMap::new(); { let ledger = lock(&self.ledger); @@ -1586,21 +1582,21 @@ macro_rules! sql_store_impl { created_at_epoch_secs, estimated, billed_units) ", ); qb.push_values(rows, |mut v, r| { - v.push_bind(nul_free(&r.ak)) - .push_bind(nul_free(&r.product)) - .push_bind(nul_free(&r.tenant)) - .push_bind(nul_free(&r.model)) - .push_bind(nul_free(&r.served_model)) - .push_bind(nul_free(&r.protocol)) - .push_bind(nul_free(&r.account)) + v.push_bind(&r.ak) + .push_bind(&r.product) + .push_bind(&r.tenant) + .push_bind(&r.model) + .push_bind(&r.served_model) + .push_bind(&r.protocol) + .push_bind(&r.account) .push_bind(r.prompt_tokens) .push_bind(r.completion_tokens) .push_bind(r.total_tokens) .push_bind(r.cost_micros) .push_bind(r.vendor_cost_micros) .push_bind(r.ptu_spillover) - .push_bind(nul_free(&r.user_id)) - .push_bind(nul_free(&r.request_id)) + .push_bind(&r.user_id) + .push_bind(&r.request_id) .push_bind(r.created_at_epoch_secs) .push_bind(r.estimated) .push_bind(r.billed_units); @@ -1913,18 +1909,14 @@ macro_rules! sql_store_impl { } sql_store_impl!(SqliteStore, sqlite, { - async fn usage_rollup_since(&self, since: i64, now: i64) -> GResult { + async fn usage_rollup_since(&self, since: Option, now: i64) -> GResult { let hi = bucket_floor(now - ROLLUP_SETTLE_SECS); let watermark: i64 = sqlx::query_scalar(ROLLUP_WATERMARK_SQL) .fetch_one(&self.pool) .await .map_err(|e| crate::sqlx_err("read rollup watermark", e))?; let r = sqlx::query(dialect_sql!(sqlite, q_rollup_insert!(sqlite))) - .bind( - (hi - ROLLUP_BACKFILL_SECS) - .min(watermark) - .min(bucket_floor(since)), - ) + .bind(rollup_floor(hi, watermark, since)) .bind(hi) .execute(&self.pool) .await @@ -2228,14 +2220,14 @@ impl PostgresStore { } sql_store_impl!(PostgresStore, postgres, { - async fn usage_rollup_since(&self, since: i64, now: i64) -> GResult { + async fn usage_rollup_since(&self, since: Option, now: i64) -> GResult { let hi = bucket_floor(now - ROLLUP_SETTLE_SECS); let mut tx = self .pool .begin() .await .map_err(|e| crate::sqlx_err("begin rollup", e))?; - if since == i64::MAX { + if since.is_none() { let locked: bool = sqlx::query_scalar("SELECT pg_try_advisory_xact_lock($1)") .bind(ROLLUP_LOCK_KEY) .fetch_one(&mut *tx) @@ -2256,11 +2248,7 @@ sql_store_impl!(PostgresStore, postgres, { .await .map_err(|e| crate::sqlx_err("read rollup watermark", e))?; let r = sqlx::query(dialect_sql!(postgres, q_rollup_insert!(postgres))) - .bind( - (hi - ROLLUP_BACKFILL_SECS) - .min(watermark) - .min(bucket_floor(since)), - ) + .bind(rollup_floor(hi, watermark, since)) .bind(hi) .execute(&mut *tx) .await @@ -2780,12 +2768,9 @@ where ) } -fn nul_free(s: &str) -> Cow<'_, str> { - if s.contains('\0') { - Cow::Owned(s.replace('\0', "")) - } else { - Cow::Borrowed(s) - } +fn rollup_floor(hi: i64, watermark: i64, since: Option) -> i64 { + let lo = (hi - ROLLUP_BACKFILL_SECS).min(watermark); + since.map_or(lo, |since| lo.min(bucket_floor(since))) } fn bucket_floor(ts: i64) -> i64 { @@ -3498,7 +3483,7 @@ mod tests { ), "a row billed below the watermark misses the trailing rollup" ); - store.usage_rollup_since(100, 1_500).await.unwrap(); + store.usage_rollup_since(Some(100), 1_500).await.unwrap(); assert!( billed_late( store @@ -3510,22 +3495,6 @@ mod tests { ); } - async fn exercise_nul_text(store: &dyn Store) { - let mut r = record("m1"); - r.user_id = "nul\0user".into(); - r.request_id = format!("{}\0", r.request_id); - store - .ledger_add_batch(&[r.clone(), record("m1")]) - .await - .unwrap(); - let (_, rows) = store.ledger_snapshot(usize::MAX).await.unwrap(); - let stored = rows - .iter() - .find(|row| row.request_id == r.request_id.replace('\0', "")) - .unwrap(); - assert_eq!(stored.user_id, "nuluser"); - } - #[tokio::test] async fn memory_rollup_roundtrip() { exercise_rollup(&MemoryStore::default(), "").await; @@ -3538,7 +3507,6 @@ mod tests { .await .unwrap(); exercise_rollup(&store, "").await; - exercise_nul_text(&store).await; } #[tokio::test] @@ -3913,7 +3881,6 @@ mod tests { let store = PostgresStore::connect(&own_url).await.expect("pg connect"); let ns = format!("-{nonce}"); exercise_rollup(&store, &ns).await; - exercise_nul_text(&store).await; exercise_erase(&store, &ns).await; exercise_content_list(&store, &ns).await; exercise_terminal_idempotence(&store, &ns).await; diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index ea90ca0..176adcd 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -1350,14 +1350,14 @@ fn user_hint(hint: Option, field: &Value) -> Option { hint.or_else(|| field.as_str().map(cap_user_hint)) } -/// Bound an attribution hint at `USER_HINT_MAX_BYTES`: it keys governance -/// counters, so an unbounded value from any surface would grow the keyspace. +/// Bound an attribution hint at `USER_HINT_MAX_BYTES` (it keys governance counters) +/// and drop NUL, which Postgres TEXT refuses in every row the user lands in. fn cap_user_hint(s: &str) -> String { let mut end = s.len().min(USER_HINT_MAX_BYTES); while !s.is_char_boundary(end) { end -= 1; } - s[..end].to_owned() + s[..end].replace('\0', "") } /// AK auth: `Authorization: Bearer ` or `x-api-key: `. The error is @@ -6731,6 +6731,17 @@ mod tests { ); } + #[test] + fn user_hints_are_capped_and_drop_nul() { + assert_eq!(cap_user_hint("nul\0user"), "nuluser"); + let long = "é".repeat(USER_HINT_MAX_BYTES); + assert!(cap_user_hint(&long).len() <= USER_HINT_MAX_BYTES); + assert_eq!( + user_hint(None, &serde_json::json!("\u{0}")), + Some(String::new()) + ); + } + #[tokio::test] async fn realtime_attributes_user_from_owner_then_header_hint() { let yaml = "listen: {host: h, port: 1}\nmodels: [{name: rt, protocol: realtime, input_price_per_1k_micros: 1000, output_price_per_1k_micros: 1000}]\naccess_keys: [{ak: k-shared, product: p, qps: 10, daily_token_quota: 100000}, {ak: k-owned, product: p, qps: 10, daily_token_quota: 100000, owner: bob}]"; From 0f3a5c17737e4f385769b6083fe5d1b1fab76e70 Mon Sep 17 00:00:00 2001 From: CMGS Date: Sat, 17 Oct 2026 19:58:50 +0800 Subject: [PATCH 09/22] cut: fold Settled into BudgetSpend, drop the generation id whitelist loc-justify on PR 158: - Settled duplicated BudgetSpend field for field, and settle_and_bill rebuilt one from the other to charge the budgets; it now returns the BudgetSpend. - The generation id comes from the account's own OpenRouter stream and only reaches a GET to that host with that key; a malformed id fails the lookup, which already falls back to the estimate, so the character whitelist defended nothing inside the deployment contract. --- crates/dag/src/nodes.rs | 2 +- crates/engines/src/generation.rs | 10 ++-------- crates/state/src/admission.rs | 17 ++++------------- 3 files changed, 7 insertions(+), 22 deletions(-) diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index a19559c..a229b03 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -950,7 +950,7 @@ async fn bill(ctx: &mut DagContext, mut tokens: BillTokens, estimated: bool) -> "cost_calc", format!( "tokens={} cost_micros={}", - settled.total_tokens, settled.cost_micros + settled.tokens, settled.cost_micros ), ); Ok(()) diff --git a/crates/engines/src/generation.rs b/crates/engines/src/generation.rs index 0251051..ded339f 100644 --- a/crates/engines/src/generation.rs +++ b/crates/engines/src/generation.rs @@ -57,12 +57,7 @@ pub async fn recorded_usage( } pub fn records_generation(account: &Account, id: &str) -> bool { - account.wire_kind() == OPENROUTER_KIND - && !account.endpoint.is_empty() - && !id.is_empty() - && id - .bytes() - .all(|b| b.is_ascii_alphanumeric() || b == b'-' || b == b'_') + account.wire_kind() == OPENROUTER_KIND && !account.endpoint.is_empty() && !id.is_empty() } /// The generation's native counts and cost as chat-wire usage; `None` until OpenRouter records it. @@ -221,7 +216,7 @@ mod tests { } #[test] - fn only_an_openrouter_account_with_a_plain_id_records_generations() { + fn only_an_openrouter_account_with_an_id_records_generations() { assert!(records_generation( &openrouter(), "gen-1792232109-kUA52UZsI2" @@ -237,7 +232,6 @@ mod tests { ..openrouter() }; assert!(records_generation(&raw, "gen-1")); - assert!(!records_generation(&openrouter(), "gen/../x")); assert!(!records_generation(&openrouter(), "")); } } diff --git a/crates/state/src/admission.rs b/crates/state/src/admission.rs index 98e519f..d32f324 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -105,12 +105,6 @@ pub struct SettleInput<'a> { pub budget_hold: BudgetHold, } -/// What a settled request cost, for the decision trail. -pub struct Settled { - pub total_tokens: i64, - pub cost_micros: i64, -} - /// A settle in progress: shutdown waits for it to drop and fires `closing` so it bills without waiting. pub struct SettleGuard { _tracked: TaskTrackerToken, @@ -713,12 +707,12 @@ pub async fn settle_and_bill( state: &GatewayState, cfg: &GatewayConfig, s: SettleInput<'_>, -) -> Settled { +) -> BudgetSpend { let gov = state.governance.as_ref(); let total = clamp_tokens(s.billing.total); let record = billing_record(cfg, &s.billing); - let settled = Settled { - total_tokens: record.total_tokens, + let settled = BudgetSpend { + tokens: record.total_tokens, cost_micros: record.cost_micros, }; let settle_daily = gov.quota_settle(s.billing.ak, total - s.reserved, s.reserved_at); @@ -741,10 +735,7 @@ pub async fn settle_and_bill( s.billing.user_id, s.user_budget, s.budget_hold, - BudgetSpend { - tokens: settled.total_tokens, - cost_micros: settled.cost_micros, - }, + settled, ); let write_ledger = state.billing.write(record); tokio::join!( From 9a10d9b054a4bee268248aa090a57de6be493568 Mon Sep 17 00:00:00 2001 From: CMGS Date: Sat, 17 Oct 2026 20:06:47 +0800 Subject: [PATCH 10/22] docs: an aborted OpenRouter stream bills a cache write at the prompt weight OpenRouter's generation record carries cache_discount in dollars but no cache-write token count (live: claude-haiku-4.5 on Bedrock wrote 8402 tokens, total_cost matched 1.25x, native_tokens_cached 0), so the gateway cannot weight the write; vendor_cost_micros stays exact. --- docs/observability.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/observability.md b/docs/observability.md index f59f79b..7cc2840 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -90,7 +90,8 @@ when counts came from an aborted stream rather than a vendor usage payload). An aborted stream served by an OpenRouter account is billed from OpenRouter's record of the generation instead (`GET /api/v1/generation`, polled for about 40 s with each lookup capped at 10 s; cancelled generations included), with cache -reads at their weight and OpenRouter's cost as `vendor_cost_micros`; only when no +reads at their weight (the record carries no cache-write count, so a cache write +bills at the prompt weight) and OpenRouter's cost as `vendor_cost_micros`; only when no record arrives, or the gateway starts shutting down, does it fall back to the estimate. Shutdown waits for these settles before the final ledger flush, and the access log's `latency_ms` stops at the end of the stream, not of the lookup. From 259bffac81d6cb5d3f0fbf942033531b0f297230 Mon Sep 17 00:00:00 2001 From: CMGS Date: Sat, 17 Oct 2026 20:38:48 +0800 Subject: [PATCH 11/22] perf: hold and settle every day budget in one Redis script An A/B on cocoon-test1 (Redis governance, key + tenant day budgets, mock upstream) showed the per-budget holds cost 8.7% (small) and 11.8% (big) throughput against main: each day budget swapped a GET at admission and an INCRBY script at settle for its own reserve and settle scripts, and Redis, single-threaded, is the bottleneck under that load. quota_reserve_all checks every held budget and increments them in one script, all or none, so a denial needs no refund round trip; quota_settle_all settles, refunds and reads the alert totals in one script. A request now costs two EVALSHA for its day budgets where main paid two GETs plus two scripts. quota_settle and settle_floored return to main's shape. --- crates/state/src/admission.rs | 141 +++++++++++++++++----------- crates/state/src/governance.rs | 162 +++++++++++++++++++++++++++++---- 2 files changed, 233 insertions(+), 70 deletions(-) diff --git a/crates/state/src/admission.rs b/crates/state/src/admission.rs index d32f324..f781528 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -73,11 +73,12 @@ pub struct BudgetHold { impl BudgetHold { /// Return every held amount, for a call that never reached billing. pub async fn refund(self, gov: &dyn Governance) { - let at = self.at; - futures::future::join_all(self.held.into_iter().map(|(key, amount)| async move { - gov.quota_settle(&key, -amount, at).await; - })) - .await; + let refunds: Vec<_> = self + .held + .iter() + .map(|(key, amount)| (key.as_str(), -amount)) + .collect(); + gov.quota_settle_all(&refunds, self.at).await; } fn take(&mut self, key: &str) -> Option { @@ -521,41 +522,59 @@ pub async fn reserve_budgets( } else { BudgetSpend::default() }; - let admitted = futures::future::join_all(budgets.iter().map(|b| async move { - let amount = est.of(b.scope); - match b.window { - Window::Day if amount > 0 => gov.quota_reserve(&b.key, amount, b.limit, at).await, - Window::Day => gov.quota_check(&b.key, b.limit).await.then_some(0), - Window::Month => (gov.counter_get(&b.key).await < b.limit).then_some(0), - } - })) - .await; - let mut hold = BudgetHold { - held: Vec::new(), - at, - }; - let mut denied = None; - for (b, admitted) in budgets.into_iter().zip(admitted) { - match admitted { - Some(held) if held > 0 => hold.held.push((b.key, held)), - Some(_) => {} - None => { - denied.get_or_insert_with(|| { - format!( - "{} {} budget exhausted for {}", - b.window.label(), - b.scope.unit(), - b.scope.subject(ak, user) - ) - }); + let holds: Vec<(usize, i64)> = budgets + .iter() + .enumerate() + .filter(|(_, b)| b.window == Window::Day) + .map(|(i, b)| (i, est.of(b.scope))) + .filter(|&(_, amount)| amount > 0) + .collect(); + let hold_args: Vec<_> = holds + .iter() + .map(|&(i, amount)| (budgets[i].key.as_str(), amount, budgets[i].limit)) + .collect(); + let (reserved, checked) = tokio::join!( + gov.quota_reserve_all(&hold_args, at), + futures::future::join_all(budgets.iter().map(|b| async move { + match b.window { + Window::Day if est.of(b.scope) > 0 => true, + Window::Day => gov.quota_check(&b.key, b.limit).await, + Window::Month => gov.counter_get(&b.key).await < b.limit, } - } - } - let Some(denied) = denied else { - return Ok(Ok((over, hold))); + })), + ); + let refused = reserved.err().map(|i| holds[i].0); + let Some(i) = checked + .iter() + .position(|ok| !ok) + .into_iter() + .chain(refused) + .min() + else { + let held = if reserved == Ok(true) { + holds + .iter() + .map(|&(i, amount)| (budgets[i].key.clone(), amount)) + .collect() + } else { + Vec::new() + }; + return Ok(Ok((over, BudgetHold { held, at }))); }; - hold.refund(gov).await; - Ok(Err(denied)) + if reserved == Ok(true) { + let refunds: Vec<_> = hold_args + .iter() + .map(|&(key, amount, _)| (key, -amount)) + .collect(); + gov.quota_settle_all(&refunds, at).await; + } + let b = &budgets[i]; + Ok(Err(format!( + "{} {} budget exhausted for {}", + b.window.label(), + b.scope.unit(), + b.scope.subject(ak, user) + ))) } pub fn model_quota_key(ak_id: &str, model: &str) -> String { @@ -778,25 +797,39 @@ async fn consume_budgets( }), }; let budgets = budgets(gov, cfg, ak, user, over).await; - let at = hold.at; - let used = futures::future::join_all(budgets.iter().map(|b| { - let held = hold.take(&b.key); - async move { + let (mut settles, mut accrues) = (Vec::new(), Vec::new()); + for (i, b) in budgets.iter().enumerate() { + let amount = spent.of(b.scope); + match (b.window, hold.take(&b.key)) { + (Window::Day, Some(held)) => settles.push((i, b.key.as_str(), amount - held)), + _ if amount <= 0 => {} + _ => accrues.push(i), + } + } + let deltas: Vec<_> = settles + .iter() + .map(|&(_, key, delta)| (key, delta)) + .collect(); + let (settled, accrued) = tokio::join!( + gov.quota_settle_all(&deltas, hold.at), + futures::future::join_all(accrues.iter().map(|&i| { + let b = &budgets[i]; let amount = spent.of(b.scope); - match (b.window, held) { - (Window::Day, Some(held)) if held == amount => Some(gov.quota_used(&b.key).await), - (Window::Day, Some(held)) => { - Some(gov.quota_settle(&b.key, amount - held, at).await) - } - _ if amount <= 0 => None, - (Window::Day, None) => Some(gov.quota_consume(&b.key, amount).await), - (Window::Month, _) => { - Some(gov.counter_add(&b.key, amount, MONTH_COUNTER_TTL).await) + async move { + match b.window { + Window::Day => gov.quota_consume(&b.key, amount).await, + Window::Month => gov.counter_add(&b.key, amount, MONTH_COUNTER_TTL).await, } } - } - })) - .await; + })), + ); + let mut used = vec![None; budgets.len()]; + for (&(i, _, _), total) in settles.iter().zip(settled) { + used[i] = Some(total); + } + for (&i, total) in accrues.iter().zip(accrued) { + used[i] = Some(total); + } for (b, used) in budgets.iter().zip(used) { if let Some(used) = used && used >= b.limit diff --git a/crates/state/src/governance.rs b/crates/state/src/governance.rs index 22b625c..10de5bc 100644 --- a/crates/state/src/governance.rs +++ b/crates/state/src/governance.rs @@ -35,9 +35,16 @@ pub trait Governance: Send + Sync + std::fmt::Debug { limit: i64, at_epoch_secs: i64, ) -> Option; + /// [`Self::quota_reserve`] over every hold at once, all or none: `Err(i)` = `holds[i]` is at its cap, `Ok(false)` = admitted unrecorded. + async fn quota_reserve_all( + &self, + holds: &[(&str, i64, i64)], + at_epoch_secs: i64, + ) -> Result; /// Apply the settle delta (actual - reserved; negative refunds) to the day - /// bucket the paired reserve used (`at_epoch_secs`); returns the new total, or 0 when nothing was applied. - async fn quota_settle(&self, key: &str, delta: i64, at_epoch_secs: i64) -> i64; + /// bucket the paired reserve used (`at_epoch_secs`). + async fn quota_settle(&self, key: &str, delta: i64, at_epoch_secs: i64); + async fn quota_settle_all(&self, deltas: &[(&str, i64)], at_epoch_secs: i64) -> Vec; /// Tokens spent today under `key`. async fn quota_used(&self, key: &str) -> i64; /// Add to `key`'s spent tokens; returns the day's new total. @@ -105,11 +112,25 @@ impl Governance for MemoryGovernance { async fn quota_reserve(&self, key: &str, amount: i64, limit: i64, _at: i64) -> Option { self.quota.reserve(key, amount, limit).then_some(amount) } - async fn quota_settle(&self, key: &str, delta: i64, _at: i64) -> i64 { - if delta == 0 { - return 0; + async fn quota_reserve_all(&self, holds: &[(&str, i64, i64)], _at: i64) -> Result { + for (i, &(key, amount, limit)) in holds.iter().enumerate() { + if !self.quota.reserve(key, amount, limit) { + for &(key, amount, _) in &holds[..i] { + self.quota.settle(key, -amount); + } + return Err(i); + } } - self.quota.settle(key, delta) + Ok(true) + } + async fn quota_settle(&self, key: &str, delta: i64, _at: i64) { + self.quota.settle(key, delta); + } + async fn quota_settle_all(&self, deltas: &[(&str, i64)], _at: i64) -> Vec { + deltas + .iter() + .map(|&(key, delta)| self.quota.settle(key, delta)) + .collect() } async fn quota_used(&self, key: &str) -> i64 { self.quota.used(key) @@ -278,9 +299,42 @@ impl Governance for RedisGovernance { } }) } - async fn quota_settle(&self, key: &str, delta: i64, at: i64) -> i64 { + async fn quota_reserve_all(&self, holds: &[(&str, i64, i64)], at: i64) -> Result { + static SCRIPT: LazyLock = LazyLock::new(|| { + redis::Script::new( + "for i = 1, #KEYS do + if tonumber(redis.call('GET', KEYS[i]) or '0') >= tonumber(ARGV[2 * i + 1]) then return i end + end + for i = 1, #KEYS do + local amount = tonumber(ARGV[2 * i]) + if redis.call('INCRBY', KEYS[i], amount) == amount then redis.call('PEXPIRE', KEYS[i], ARGV[1]) end + end + return 0", + ) + }); + if holds.is_empty() { + return Ok(true); + } + let mut invocation = SCRIPT.prepare_invoke(); + invocation.arg(QUOTA_TTL_MS); + for &(key, amount, limit) in holds { + invocation.key(quota_key_at(key, at)).arg(amount).arg(limit); + } + match invocation + .invoke_async::(&mut self.conn.clone()) + .await + { + Ok(0) => Ok(true), + Ok(i) => Err(i - 1), + Err(e) => { + tracing::warn!(error = %e, "redis budget reserve failed; admitting"); + Ok(false) + } + } + } + async fn quota_settle(&self, key: &str, delta: i64, at: i64) { if delta == 0 { - return 0; + return; } settle_floored( &self.conn, @@ -289,7 +343,36 @@ impl Governance for RedisGovernance { Duration::from_millis(QUOTA_TTL_MS as u64), None, ) - .await + .await; + } + async fn quota_settle_all(&self, deltas: &[(&str, i64)], at: i64) -> Vec { + static SCRIPT: LazyLock = LazyLock::new(|| { + redis::Script::new( + "local out = {} + for i = 1, #KEYS do + local v = redis.call('INCRBY', KEYS[i], ARGV[i + 1]) + if v < 0 then redis.call('SET', KEYS[i], 0, 'KEEPTTL'); v = 0 end + if redis.call('PTTL', KEYS[i]) < 0 then redis.call('PEXPIRE', KEYS[i], ARGV[1]) end + out[i] = v + end + return out", + ) + }); + if deltas.is_empty() { + return Vec::new(); + } + let mut invocation = SCRIPT.prepare_invoke(); + invocation.arg(QUOTA_TTL_MS); + for &(key, delta) in deltas { + invocation.key(quota_key_at(key, at)).arg(delta); + } + invocation + .invoke_async(&mut self.conn.clone()) + .await + .unwrap_or_else(|e| { + tracing::warn!(error = %e, "redis budget settle failed"); + vec![0; deltas.len()] + }) } async fn quota_consume(&self, key: &str, tokens: i64) -> i64 { self.incr_window( @@ -351,14 +434,13 @@ fn quota_key(key: &str) -> String { /// Apply a settle delta and floor at 0 atomically (a rollover between reserve /// and settle must not plant a negative); keeps an existing TTL, arms `window` otherwise. -/// Returns the new total, or 0 when the window moved on or the round-trip failed. async fn settle_floored( conn: &redis::aio::ConnectionManager, key: &str, delta: i64, window: Duration, id: Option, -) -> i64 { +) { static SCRIPT: LazyLock = LazyLock::new(|| { redis::Script::new( "local id = tonumber(ARGV[3]) @@ -372,17 +454,16 @@ async fn settle_floored( ) }); let mut conn = conn.clone(); - SCRIPT + if let Err(e) = SCRIPT .key(key) .arg(delta) .arg(window.as_millis() as i64) .arg(id.unwrap_or(-1)) .invoke_async::(&mut conn) .await - .unwrap_or_else(|e| { - tracing::warn!(error = %e, key, "redis settle failed"); - 0 - }) + { + tracing::warn!(error = %e, key, "redis settle failed"); + } } #[cfg(test)] @@ -414,6 +495,55 @@ mod tests { .unwrap(); } + async fn exercise_reserve_all(g: &dyn Governance, ns: &str) { + let (a, b, c) = (format!("ra{ns}"), format!("rb{ns}"), format!("rc{ns}")); + let now = crate::epoch_secs(); + let pair = [(a.as_str(), 60, 100), (b.as_str(), 30, 50)]; + assert_eq!(g.quota_reserve_all(&pair, now).await, Ok(true)); + assert_eq!(g.quota_reserve_all(&pair, now).await, Ok(true)); + assert_eq!((g.quota_used(&a).await, g.quota_used(&b).await), (120, 60)); + assert_eq!( + g.quota_reserve_all(&[(a.as_str(), 1, 100), (b.as_str(), 1, 50)], now) + .await, + Err(0) + ); + assert_eq!( + g.quota_reserve_all(&[(c.as_str(), 5, 100), (b.as_str(), 1, 50)], now) + .await, + Err(1) + ); + assert_eq!( + ( + g.quota_used(&a).await, + g.quota_used(&b).await, + g.quota_used(&c).await + ), + (120, 60, 0) + ); + assert_eq!( + g.quota_settle_all(&[(a.as_str(), -200), (b.as_str(), -10)], now) + .await, + vec![0, 50] + ); + assert_eq!(g.quota_settle_all(&[(b.as_str(), 0)], now).await, vec![50]); + assert_eq!(g.quota_reserve_all(&[], now).await, Ok(true)); + assert!(g.quota_settle_all(&[], now).await.is_empty()); + } + + #[tokio::test] + async fn memory_reserve_all_holds_all_or_none() { + exercise_reserve_all(&MemoryGovernance::default(), "").await; + } + + #[tokio::test] + async fn redis_reserve_all_holds_all_or_none() { + let Ok(url) = std::env::var("GW_TEST_REDIS_URL") else { + return; + }; + let g = RedisGovernance::connect(&url).await.expect("redis connect"); + exercise_reserve_all(&g, &format!("-{}", std::process::id())).await; + } + #[tokio::test] async fn redis_governance_enforces_limits() { let Ok(url) = std::env::var("GW_TEST_REDIS_URL") else { From 949017063a6ce37a201a268e25b44432244cfcab Mon Sep 17 00:00:00 2001 From: CMGS Date: Sat, 17 Oct 2026 20:52:33 +0800 Subject: [PATCH 12/22] perf: skip the budget scripts when nothing is held Governance is an async_trait, so calling quota_reserve_all / quota_settle_all boxed a future even for an empty input: a request with no budgets paid two heap allocations main did not. The callers now skip empty inputs (and the redundant guards inside the Redis impl go); held keys move out of the budget list instead of being cloned. --- crates/state/src/admission.rs | 32 ++++++++++++++++++++++++++++---- crates/state/src/governance.rs | 6 ------ 2 files changed, 28 insertions(+), 10 deletions(-) diff --git a/crates/state/src/admission.rs b/crates/state/src/admission.rs index f781528..1ec7f8b 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -73,6 +73,9 @@ pub struct BudgetHold { impl BudgetHold { /// Return every held amount, for a call that never reached billing. pub async fn refund(self, gov: &dyn Governance) { + if self.held.is_empty() { + return; + } let refunds: Vec<_> = self .held .iter() @@ -516,7 +519,16 @@ pub async fn reserve_budgets( ) -> GResult> { let over = user_budget(state, ak, user).await?; let gov = state.governance.as_ref(); - let budgets = budgets(gov, cfg, ak, user, over).await; + let mut budgets = budgets(gov, cfg, ak, user, over).await; + if budgets.is_empty() { + return Ok(Ok(( + over, + BudgetHold { + held: Vec::new(), + at, + }, + ))); + } let est = if budgets.iter().any(|b| b.window == Window::Day) { est() } else { @@ -534,7 +546,13 @@ pub async fn reserve_budgets( .map(|&(i, amount)| (budgets[i].key.as_str(), amount, budgets[i].limit)) .collect(); let (reserved, checked) = tokio::join!( - gov.quota_reserve_all(&hold_args, at), + async { + if hold_args.is_empty() { + Ok(true) + } else { + gov.quota_reserve_all(&hold_args, at).await + } + }, futures::future::join_all(budgets.iter().map(|b| async move { match b.window { Window::Day if est.of(b.scope) > 0 => true, @@ -554,7 +572,7 @@ pub async fn reserve_budgets( let held = if reserved == Ok(true) { holds .iter() - .map(|&(i, amount)| (budgets[i].key.clone(), amount)) + .map(|&(i, amount)| (std::mem::take(&mut budgets[i].key), amount)) .collect() } else { Vec::new() @@ -811,7 +829,13 @@ async fn consume_budgets( .map(|&(_, key, delta)| (key, delta)) .collect(); let (settled, accrued) = tokio::join!( - gov.quota_settle_all(&deltas, hold.at), + async { + if deltas.is_empty() { + Vec::new() + } else { + gov.quota_settle_all(&deltas, hold.at).await + } + }, futures::future::join_all(accrues.iter().map(|&i| { let b = &budgets[i]; let amount = spent.of(b.scope); diff --git a/crates/state/src/governance.rs b/crates/state/src/governance.rs index 10de5bc..db462d8 100644 --- a/crates/state/src/governance.rs +++ b/crates/state/src/governance.rs @@ -312,9 +312,6 @@ impl Governance for RedisGovernance { return 0", ) }); - if holds.is_empty() { - return Ok(true); - } let mut invocation = SCRIPT.prepare_invoke(); invocation.arg(QUOTA_TTL_MS); for &(key, amount, limit) in holds { @@ -358,9 +355,6 @@ impl Governance for RedisGovernance { return out", ) }); - if deltas.is_empty() { - return Vec::new(); - } let mut invocation = SCRIPT.prepare_invoke(); invocation.arg(QUOTA_TTL_MS); for &(key, delta) in deltas { From ddad28a85205f4c9d188437b42d2b5cf55c4ea43 Mon Sep 17 00:00:00 2001 From: CMGS Date: Sat, 17 Oct 2026 21:50:11 +0800 Subject: [PATCH 13/22] perf: cheaper budget scripts, matching main's Redis work per request perf stat on the 16379 redis-server (budget-redis config, 200k requests per arm) still showed 196k instructions per request against main's 146k: the reserve script read every key before incrementing it, and the settle script ran PTTL on every key. Reserve now increments first and rolls back on a refusal, the way reserve_capped does for one key; settle arms the TTL only on a key it created, the way incr_window does. The Redis test asserts the TTL and deletes its keys. --- crates/state/src/governance.rs | 34 ++++++++++++++++++++++++++++------ 1 file changed, 28 insertions(+), 6 deletions(-) diff --git a/crates/state/src/governance.rs b/crates/state/src/governance.rs index db462d8..84d52a6 100644 --- a/crates/state/src/governance.rs +++ b/crates/state/src/governance.rs @@ -303,11 +303,13 @@ impl Governance for RedisGovernance { static SCRIPT: LazyLock = LazyLock::new(|| { redis::Script::new( "for i = 1, #KEYS do - if tonumber(redis.call('GET', KEYS[i]) or '0') >= tonumber(ARGV[2 * i + 1]) then return i end - end - for i = 1, #KEYS do local amount = tonumber(ARGV[2 * i]) - if redis.call('INCRBY', KEYS[i], amount) == amount then redis.call('PEXPIRE', KEYS[i], ARGV[1]) end + local v = redis.call('INCRBY', KEYS[i], amount) + if v == amount then redis.call('PEXPIRE', KEYS[i], ARGV[1]) end + if v - amount >= tonumber(ARGV[2 * i + 1]) then + for j = 1, i do redis.call('DECRBY', KEYS[j], ARGV[2 * j]) end + return i + end end return 0", ) @@ -348,8 +350,8 @@ impl Governance for RedisGovernance { "local out = {} for i = 1, #KEYS do local v = redis.call('INCRBY', KEYS[i], ARGV[i + 1]) + if v == tonumber(ARGV[i + 1]) then redis.call('PEXPIRE', KEYS[i], ARGV[1]) end if v < 0 then redis.call('SET', KEYS[i], 0, 'KEEPTTL'); v = 0 end - if redis.call('PTTL', KEYS[i]) < 0 then redis.call('PEXPIRE', KEYS[i], ARGV[1]) end out[i] = v end return out", @@ -535,7 +537,27 @@ mod tests { return; }; let g = RedisGovernance::connect(&url).await.expect("redis connect"); - exercise_reserve_all(&g, &format!("-{}", std::process::id())).await; + let ns = format!("-{}", std::process::id()); + exercise_reserve_all(&g, &ns).await; + let now = crate::epoch_secs(); + let keys: Vec<_> = ["ra", "rb", "rc"] + .iter() + .map(|k| quota_key_at(&format!("{k}{ns}"), now)) + .collect(); + let mut conn = g.conn.clone(); + for key in &keys { + let ttl: i64 = redis::cmd("PTTL") + .arg(key) + .query_async(&mut conn) + .await + .unwrap(); + assert!(ttl > 0, "{key} ttl {ttl}"); + } + let _: i64 = redis::cmd("DEL") + .arg(&keys) + .query_async(&mut conn) + .await + .unwrap(); } #[tokio::test] From f71493dd49e797013aa1feba80429bcdd4c22ff1 Mon Sep 17 00:00:00 2001 From: CMGS Date: Sat, 17 Oct 2026 21:58:55 +0800 Subject: [PATCH 14/22] perf: hold and settle day budgets with pipelined INCRBY, not Lua Redis commandstats per request (budget-redis config) still showed the two multi-key scripts costing 2.5 us more than main's two GETs plus two single-key scripts: every EVALSHA pays the Lua VM. Both now send one pipeline of plain INCRBY and read the totals; a refusal sends one pipeline of DECRBY, a key the increment created gets its TTL, and a negative total is floored, each only when it happens. A refusal is no longer atomic: for one round trip other admissions see the refused hold, which can only refuse them, never overcharge. --- crates/state/src/governance.rs | 118 +++++++++++++++++++-------------- 1 file changed, 69 insertions(+), 49 deletions(-) diff --git a/crates/state/src/governance.rs b/crates/state/src/governance.rs index 84d52a6..fa82e93 100644 --- a/crates/state/src/governance.rs +++ b/crates/state/src/governance.rs @@ -240,6 +240,30 @@ impl RedisGovernance { } } + async fn incr_all(&self, incrs: &[(String, i64)]) -> redis::RedisResult> { + if incrs.is_empty() { + return Ok(Vec::new()); + } + let mut pipe = redis::pipe(); + for (key, by) in incrs { + pipe.cmd("INCRBY").arg(key).arg(*by); + } + let mut conn = self.conn.clone(); + let totals: Vec = pipe.query_async(&mut conn).await?; + let mut arm = redis::pipe(); + let mut fresh = false; + for ((key, by), total) in incrs.iter().zip(&totals) { + if total == by { + arm.cmd("PEXPIRE").arg(key).arg(QUOTA_TTL_MS).ignore(); + fresh = true; + } + } + if fresh { + arm.query_async::<()>(&mut conn).await?; + } + Ok(totals) + } + /// Increment `key` and set its TTL on first use; returns the post-increment /// count. A failed round-trip returns 0 so limits fail open rather than /// wedging the gateway on a Redis blip. @@ -300,36 +324,32 @@ impl Governance for RedisGovernance { }) } async fn quota_reserve_all(&self, holds: &[(&str, i64, i64)], at: i64) -> Result { - static SCRIPT: LazyLock = LazyLock::new(|| { - redis::Script::new( - "for i = 1, #KEYS do - local amount = tonumber(ARGV[2 * i]) - local v = redis.call('INCRBY', KEYS[i], amount) - if v == amount then redis.call('PEXPIRE', KEYS[i], ARGV[1]) end - if v - amount >= tonumber(ARGV[2 * i + 1]) then - for j = 1, i do redis.call('DECRBY', KEYS[j], ARGV[2 * j]) end - return i - end - end - return 0", - ) - }); - let mut invocation = SCRIPT.prepare_invoke(); - invocation.arg(QUOTA_TTL_MS); - for &(key, amount, limit) in holds { - invocation.key(quota_key_at(key, at)).arg(amount).arg(limit); - } - match invocation - .invoke_async::(&mut self.conn.clone()) - .await - { - Ok(0) => Ok(true), - Ok(i) => Err(i - 1), + let incrs: Vec<_> = holds + .iter() + .map(|&(key, amount, _)| (quota_key_at(key, at), amount)) + .collect(); + let totals = match self.incr_all(&incrs).await { + Ok(totals) => totals, Err(e) => { tracing::warn!(error = %e, "redis budget reserve failed; admitting"); - Ok(false) + return Ok(false); } + }; + let Some(i) = holds + .iter() + .zip(&totals) + .position(|(&(_, amount, limit), &v)| v - amount >= limit) + else { + return Ok(true); + }; + let mut undo = redis::pipe(); + for (key, amount) in &incrs { + undo.cmd("DECRBY").arg(key).arg(*amount).ignore(); + } + if let Err(e) = undo.query_async::<()>(&mut self.conn.clone()).await { + tracing::warn!(error = %e, "redis budget rollback failed"); } + Err(i) } async fn quota_settle(&self, key: &str, delta: i64, at: i64) { if delta == 0 { @@ -345,30 +365,30 @@ impl Governance for RedisGovernance { .await; } async fn quota_settle_all(&self, deltas: &[(&str, i64)], at: i64) -> Vec { - static SCRIPT: LazyLock = LazyLock::new(|| { - redis::Script::new( - "local out = {} - for i = 1, #KEYS do - local v = redis.call('INCRBY', KEYS[i], ARGV[i + 1]) - if v == tonumber(ARGV[i + 1]) then redis.call('PEXPIRE', KEYS[i], ARGV[1]) end - if v < 0 then redis.call('SET', KEYS[i], 0, 'KEEPTTL'); v = 0 end - out[i] = v - end - return out", - ) - }); - let mut invocation = SCRIPT.prepare_invoke(); - invocation.arg(QUOTA_TTL_MS); - for &(key, delta) in deltas { - invocation.key(quota_key_at(key, at)).arg(delta); - } - invocation - .invoke_async(&mut self.conn.clone()) - .await - .unwrap_or_else(|e| { + let incrs: Vec<_> = deltas + .iter() + .map(|&(key, delta)| (quota_key_at(key, at), delta)) + .collect(); + let mut totals = match self.incr_all(&incrs).await { + Ok(totals) => totals, + Err(e) => { tracing::warn!(error = %e, "redis budget settle failed"); - vec![0; deltas.len()] - }) + return vec![0; deltas.len()]; + } + }; + let mut floor = redis::pipe(); + let mut floored = false; + for ((key, _), total) in incrs.iter().zip(totals.iter_mut()) { + if *total < 0 { + floor.cmd("SET").arg(key).arg(0).arg("KEEPTTL").ignore(); + *total = 0; + floored = true; + } + } + if floored && let Err(e) = floor.query_async::<()>(&mut self.conn.clone()).await { + tracing::warn!(error = %e, "redis budget floor failed"); + } + totals } async fn quota_consume(&self, key: &str, tokens: i64) -> i64 { self.incr_window( From a9425894e0d5dc8200e447ed5aee0c87e08b4c23 Mon Sep 17 00:00:00 2001 From: CMGS Date: Sat, 17 Oct 2026 22:33:23 +0800 Subject: [PATCH 15/22] perf: no budget work when no budget applies, fewer allocations when one does perf stat (instructions per request, 200k requests per arm) showed a request with no budgets paying ~480 instructions more than main: consume_budgets still built empty vectors and joined empty futures. It now returns as soon as no budget applies, as reserve_budgets already did. With budgets, the admission index vector and the settle-side totals vector go: the refused budget is found by position among the held ones, and alerts are checked as the totals arrive. --- crates/state/src/admission.rs | 59 +++++++++++++++++------------------ 1 file changed, 28 insertions(+), 31 deletions(-) diff --git a/crates/state/src/admission.rs b/crates/state/src/admission.rs index 1ec7f8b..3dd8176 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -534,16 +534,11 @@ pub async fn reserve_budgets( } else { BudgetSpend::default() }; - let holds: Vec<(usize, i64)> = budgets + let holds = |b: &Budget| b.window == Window::Day && est.of(b.scope) > 0; + let hold_args: Vec<_> = budgets .iter() - .enumerate() - .filter(|(_, b)| b.window == Window::Day) - .map(|(i, b)| (i, est.of(b.scope))) - .filter(|&(_, amount)| amount > 0) - .collect(); - let hold_args: Vec<_> = holds - .iter() - .map(|&(i, amount)| (budgets[i].key.as_str(), amount, budgets[i].limit)) + .filter(|b| holds(b)) + .map(|b| (b.key.as_str(), est.of(b.scope), b.limit)) .collect(); let (reserved, checked) = tokio::join!( async { @@ -555,13 +550,16 @@ pub async fn reserve_budgets( }, futures::future::join_all(budgets.iter().map(|b| async move { match b.window { - Window::Day if est.of(b.scope) > 0 => true, + _ if holds(b) => true, Window::Day => gov.quota_check(&b.key, b.limit).await, Window::Month => gov.counter_get(&b.key).await < b.limit, } })), ); - let refused = reserved.err().map(|i| holds[i].0); + let refused = reserved + .err() + .and_then(|n| budgets.iter().enumerate().filter(|(_, b)| holds(b)).nth(n)) + .map(|(i, _)| i); let Some(i) = checked .iter() .position(|ok| !ok) @@ -570,9 +568,10 @@ pub async fn reserve_budgets( .min() else { let held = if reserved == Ok(true) { - holds - .iter() - .map(|&(i, amount)| (std::mem::take(&mut budgets[i].key), amount)) + budgets + .iter_mut() + .filter(|b| holds(b)) + .map(|b| (std::mem::take(&mut b.key), est.of(b.scope))) .collect() } else { Vec::new() @@ -815,19 +814,21 @@ async fn consume_budgets( }), }; let budgets = budgets(gov, cfg, ak, user, over).await; - let (mut settles, mut accrues) = (Vec::new(), Vec::new()); + if budgets.is_empty() { + return; + } + let (mut settles, mut deltas, mut accrues) = (Vec::new(), Vec::new(), Vec::new()); for (i, b) in budgets.iter().enumerate() { let amount = spent.of(b.scope); match (b.window, hold.take(&b.key)) { - (Window::Day, Some(held)) => settles.push((i, b.key.as_str(), amount - held)), + (Window::Day, Some(held)) => { + settles.push(i); + deltas.push((b.key.as_str(), amount - held)); + } _ if amount <= 0 => {} _ => accrues.push(i), } } - let deltas: Vec<_> = settles - .iter() - .map(|&(_, key, delta)| (key, delta)) - .collect(); let (settled, accrued) = tokio::join!( async { if deltas.is_empty() { @@ -847,17 +848,13 @@ async fn consume_budgets( } })), ); - let mut used = vec![None; budgets.len()]; - for (&(i, _, _), total) in settles.iter().zip(settled) { - used[i] = Some(total); - } - for (&i, total) in accrues.iter().zip(accrued) { - used[i] = Some(total); - } - for (b, used) in budgets.iter().zip(used) { - if let Some(used) = used - && used >= b.limit - { + let totals = settles + .iter() + .zip(settled) + .chain(accrues.iter().zip(accrued)); + for (&i, used) in totals { + let b = &budgets[i]; + if used >= b.limit { state.alerts.emit( "budget_exhausted", b.scope.subject(ak, user), From 04c343ca208874c3a35a22a62771d8fccd3fac1d Mon Sep 17 00:00:00 2001 From: CMGS Date: Sat, 17 Oct 2026 22:54:13 +0800 Subject: [PATCH 16/22] budgets: hold the unit price of units known at admission BudgetGate priced the hold from token prices only, so a unit-priced call held nothing and the cost caps fell back to check-only: with a 250-micro day budget, two concurrent 250-micro searches both passed and spent 500. The hold now adds the unit price times the units admission already knows: one search, `n` images, the input characters of speech, one unit for transcription and rerank. Video bills at its poll, after the submit settled its hold, so it holds none. --- crates/dag/src/nodes.rs | 58 +++++++++++++++++++++++++++---- crates/handler/src/lib.rs | 73 +++++++++++++++++++++++++++++++++++++++ docs/governance.md | 8 +++-- 3 files changed, 129 insertions(+), 10 deletions(-) diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index a229b03..a148a3e 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -263,6 +263,15 @@ fn reserve_parts(req: &gw_models::GatewayRequest) -> (i64, i64) { ((prompt as i64 / 4).clamp(1, MAX_RESERVE), max_out) } +fn reserve_units(req: &gw_models::GatewayRequest) -> i64 { + match req.model_param_v2.as_ref().and_then(|p| p.typed.as_ref()) { + Some(gw_models::TypedParams::Image(p)) => p.n.max(1), + Some(gw_models::TypedParams::AudioTts(p)) => p.input.chars().count() as i64, + Some(gw_models::TypedParams::Video(_)) => 0, + _ => 1, + } +} + /// preprocess/quota_check: reserve the estimate against the AK daily quota; billing settles it. pub struct QuotaCheck; @@ -429,13 +438,19 @@ impl DagNode for BudgetGate { async fn execute(&self, ctx: &mut DagContext) -> GResult<()> { let (prompt, max_out) = ctx.reserve_parts; let (cfg, request, tenant) = (&ctx.cfg, &ctx.request, &ctx.ak.tenant); - let est = || admission::BudgetSpend { - tokens: reserve_total((prompt, max_out)), - cost_micros: gw_models::cost_micros( - prompt, - max_out, - cfg.prices_for_tenant(tenant, served_model(request.model_param_v2.as_ref())), - ), + let est = || { + let served = served_model(request.model_param_v2.as_ref()); + let units = + reserve_units(request).saturating_mul(cfg.unit_price_for_tenant(tenant, served)); + admission::BudgetSpend { + tokens: reserve_total((prompt, max_out)), + cost_micros: gw_models::cost_micros( + prompt, + max_out, + cfg.prices_for_tenant(tenant, served), + ) + .saturating_add(units), + } }; let (over, hold) = admission::reserve_budgets( &ctx.state, @@ -1152,6 +1167,35 @@ mod tests { assert_eq!(t.total(), 100); } + #[test] + fn reserve_units_counts_what_admission_already_knows() { + use gw_models::params::{ImageParams, TtsParams, VideoParams}; + use gw_models::{GatewayRequest, ModelParamV2, TypedParams}; + let units = |typed: TypedParams| { + let mut param = ModelParamV2::with_name(gw_consts::Protocol::OpenaiChat, "m"); + param.typed = Some(typed); + super::reserve_units(&GatewayRequest { + model_param_v2: Some(param), + ..Default::default() + }) + }; + assert_eq!( + units(TypedParams::Image(ImageParams { + n: 3, + ..Default::default() + })), + 3 + ); + assert_eq!( + units(TypedParams::AudioTts(TtsParams { + input: "héllo".into(), + ..Default::default() + })), + 5 + ); + assert_eq!(units(TypedParams::Video(VideoParams::default())), 0); + } + #[test] fn reserve_total_saturates_on_hostile_max_tokens() { use gw_models::params::ChatParams; diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index 99516b7..6cd55a1 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -2849,6 +2849,79 @@ mod tests { running.await.unwrap().unwrap(); } + #[derive(Debug)] + struct GatedSearch(tokio::sync::Semaphore); + + #[async_trait::async_trait] + impl gw_engines::transport::Transport for GatedSearch { + async fn send( + &self, + _: gw_engines::transport::UpstreamRequest, + ) -> GResult { + self.0.acquire().await.unwrap().forget(); + Ok(gw_engines::transport::UpstreamResponse { + status: 200, + body: gw_engines::transport::UpstreamBody::Json(bytes::Bytes::from_static( + br#"{"items":[{"title":"t","url":"https://u"}],"metadata":{"query":"q","requestId":"r","latencyMs":1}}"#, + )), + headers: Default::default(), + }) + } + } + + #[tokio::test] + async fn an_in_flight_unit_priced_call_holds_the_daily_cost_budget() { + let yaml = "listen: {host: h, port: 1}\nmodels: [{name: web, protocol: search, provider: cloudflare, unit_price_micros: 250}]\naccounts: [{name: cf, provider: cloudflare, protocols: [search]}]\ntenants: [{name: t1, key_daily_cost_quota_micros: 250}]\naccess_keys: [{ak: s-ak, tenant: t1, product: p, qps: 100, daily_token_quota: 100000}]"; + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let transport = Arc::new(GatedSearch(tokio::sync::Semaphore::new(0))); + let h = OnlineHandler::new(gw_state::SharedConfig::new(cfg, state), transport.clone()); + let ak = h.state().auth.get(&access_key_id("s-ak")).await.unwrap(); + let search = || { + let mut param = ModelParamV2::with_name(Protocol::Search, "web"); + param.typed = Some(gw_models::TypedParams::Search( + gw_models::params::SearchParams { + query: "q".into(), + count: 1, + }, + )); + GatewayRequest { + is_online: true, + model_param_v2: Some(param), + ..Default::default() + } + }; + let first = tokio::spawn({ + let (h, ak) = (h.clone(), ak.clone()); + async move { h.run(search(), ak).await } + }); + let key = format!("cb:ak:{}", ak.ak_id); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while h.state().governance.quota_used(&key).await == 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("the first call holds its unit price"); + + let second = h.run(search(), ak.clone()).await; + + let err = second.err().expect("the in-flight hold leaves no budget"); + assert!( + err.message.contains("daily cost budget exhausted"), + "{}", + err.message + ); + transport.0.add_permits(1); + first.await.unwrap().unwrap(); + let (_, ledger) = h.state().store.ledger_snapshot(usize::MAX).await.unwrap(); + assert_eq!( + (ledger.len(), ledger[0].cost_micros), + (1, 250), + "{ledger:?}" + ); + } + #[derive(Debug)] struct PiiStream; diff --git a/docs/governance.md b/docs/governance.md index 3c573fe..003bf5f 100644 --- a/docs/governance.md +++ b/docs/governance.md @@ -272,9 +272,11 @@ and `max_tokens` sides at the tenant's price), admitting while spent plus held is under the cap, so concurrent holds pass it by at most one hold; billing settles the hold to the charged amount and a failed call returns it. A call that costs more than it held still carries the total past the cap by the difference: -a call without `max_tokens` holds 256 output tokens, and unit- or -vendor-priced surfaces (images, video, decisions) hold no cost, so the cost caps -only check them. Realtime turns and the monthly caps check before the call +a call without `max_tokens` holds 256 output tokens. Units known at admission +hold at the unit price (a search, `n` images, the characters of a speech +request; one unit for transcription and rerank); video bills at its poll and +vendor-priced decisions report their cost after the call, so for those the cost +caps only check. Realtime turns and the monthly caps check before the call and accrue after it ("soft": concurrent turns can overshoot by one each). The per-(key, model) quota is soft the same way: it is a routing trigger (the fallback-model degrade), while the reserved per-key daily quota hard-caps spend. From bcdd970afe406bbd1f38456b80ea6828aa804b56 Mon Sep 17 00:00:00 2001 From: CMGS Date: Sat, 17 Oct 2026 22:59:02 +0800 Subject: [PATCH 17/22] billing: add replayed rows to pruned rollup buckets atomically The late-row re-roll recomputed a bucket from its raw rows and merged it with MAX: with ledger_max_rows on, a bucket whose rows were rolled and pruned (A) before a spooled row (B) replayed came out MAX(A, B), not A + B, and no retry converged. One instance cut off from Postgres while its peers keep writing reaches this in minutes. Replay now goes through Store::ledger_replay_batch: one transaction inserts the rows (ON CONFLICT DO NOTHING RETURNING request_id) and adds only the rows it inserted to buckets below the watermark, so a repeated replay adds nothing and buckets above the watermark stay with the regular tick. usage_rollup_since is gone; the rollup tick is back to main's shape. --- crates/state/src/admission.rs | 15 +- crates/state/src/store.rs | 304 ++++++++++++++++++++++++---------- docs/observability.md | 4 +- 3 files changed, 217 insertions(+), 106 deletions(-) diff --git a/crates/state/src/admission.rs b/crates/state/src/admission.rs index 3dd8176..f893f4a 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -364,7 +364,7 @@ impl LedgerSpool { return; } } - let (mut offset, mut replayed, mut since) = (0, 0, None); + let (mut offset, mut replayed) = (0, 0); loop { let replaying = self.replaying.clone(); let (rows, next) = match run_blocking(move || read_spool_batch(&replaying, offset)) @@ -379,24 +379,13 @@ impl LedgerSpool { if next == offset { break; } - if let Err(e) = store.ledger_add_batch(&rows).await { + if let Err(e) = store.ledger_replay_batch(&rows).await { tracing::warn!(error = %e, rows = rows.len(), "billing spool replay deferred"); return; } replayed += rows.len(); - since = rows - .iter() - .map(|r| r.created_at_epoch_secs) - .chain(since) - .min(); offset = next; } - if since.is_some() - && let Err(e) = store.usage_rollup_since(since, crate::epoch_secs()).await - { - tracing::warn!(error = %e, rows = replayed, "billing spool rollup deferred"); - return; - } let replaying = self.replaying.clone(); if let Err(e) = run_blocking(move || std::fs::remove_file(replaying)).await { tracing::error!(error = %e, path = %self.replaying.display(), "billing spool remove failed"); diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index 887f32a..bcedf0d 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -360,6 +360,8 @@ pub trait Store: Send + Sync + std::fmt::Debug { } Ok(()) } + /// Insert late rows and atomically add the ones inserted to rollup buckets the watermark passed. + async fn ledger_replay_batch(&self, rows: &[BillingRecord]) -> GResult<()>; /// Whether rows may batch off the request path; in-process backends stay synchronous. fn defers_ledger_writes(&self) -> bool { false @@ -387,12 +389,8 @@ pub trait Store: Send + Sync + std::fmt::Debug { until: i64, bucket_secs: i64, ) -> GResult>; - async fn usage_rollup_advance(&self, now_epoch_secs: i64) -> GResult { - self.usage_rollup_since(None, now_epoch_secs).await - } - /// Recompute and upsert every rollup bucket in the trailing backfill window, or back to `since` - /// for rows billed late; buckets are never deleted, so a missed tick self-heals. Returns the buckets written. - async fn usage_rollup_since(&self, since: Option, now_epoch_secs: i64) -> GResult; + /// Recompute the trailing backfill window's rollup buckets (never deleted: a missed tick self-heals); returns those written. + async fn usage_rollup_advance(&self, now_epoch_secs: i64) -> GResult; /// Delete rolled billing rows past `ledger_max_rows`; returns rows deleted. async fn ledger_prune(&self) -> GResult { Ok(0) @@ -543,7 +541,7 @@ pub struct MemoryStore { #[cfg(test)] ledger_failures: AtomicUsize, /// Minute buckets keyed by (minute, tenant, user, model); see - /// [`Store::usage_rollup_since`]. + /// [`Store::usage_rollup_advance`]. rollup: Mutex>, rollup_max_buckets: Option, sec_events: Mutex>, @@ -580,11 +578,8 @@ impl MemoryStore { pub(crate) fn fail_next_ledger_writes(&self, count: usize) { self.ledger_failures.store(count, Ordering::Relaxed); } -} -#[async_trait::async_trait] -impl Store for MemoryStore { - async fn ledger_add(&self, r: &BillingRecord) -> GResult<()> { + fn injected_ledger_failure(&self) -> GResult<()> { #[cfg(test)] if self .ledger_failures @@ -597,6 +592,14 @@ impl Store for MemoryStore { "injected ledger write failure", )); } + Ok(()) + } +} + +#[async_trait::async_trait] +impl Store for MemoryStore { + async fn ledger_add(&self, r: &BillingRecord) -> GResult<()> { + self.injected_ledger_failure()?; // watermark first: rollup-then-records is the lock order advance uses let watermark = (self.ledger_max_rows > 0).then(|| rollup_watermark(&lock(&self.rollup))); let mut ledger = lock(&self.ledger); @@ -629,6 +632,35 @@ impl Store for MemoryStore { Ok(()) } + async fn ledger_replay_batch(&self, rows: &[BillingRecord]) -> GResult<()> { + self.injected_ledger_failure()?; + let mut rollup = lock(&self.rollup); + let watermark = rollup_watermark(&rollup); + let mut ledger = lock(&self.ledger); + for r in rows { + if !ledger.request_ids.insert(r.request_id.clone()) { + continue; + } + ledger.rows.push(r.clone()); + if r.created_at_epoch_secs < watermark { + rollup + .entry(( + bucket_floor(r.created_at_epoch_secs), + r.tenant.clone(), + r.user_id.clone(), + r.model.clone(), + )) + .or_insert_with(|| UserUsageRow { + user_id: r.user_id.clone(), + model: r.model.clone(), + ..Default::default() + }) + .add_record(r); + } + } + Ok(()) + } + async fn ledger_snapshot(&self, limit: usize) -> GResult<(usize, Vec)> { let ledger = lock(&self.ledger); let records = &ledger.rows; @@ -739,10 +771,11 @@ impl Store for MemoryStore { Ok(map.into_iter().collect()) } - async fn usage_rollup_since(&self, since: Option, now: i64) -> GResult { + async fn usage_rollup_advance(&self, now: i64) -> GResult { let hi = bucket_floor(now - ROLLUP_SETTLE_SECS); let mut rollup = lock(&self.rollup); - let lo = rollup_floor(hi, rollup_watermark(&rollup), since); + let watermark = rollup_watermark(&rollup); + let lo = (hi - ROLLUP_BACKFILL_SECS).min(watermark); let mut fresh = BTreeMap::new(); { let ledger = lock(&self.ledger); @@ -1381,6 +1414,40 @@ macro_rules! q_series_raw { }; } +macro_rules! billing_insert { + ($rows:expr) => {{ + // the QueryBuilder numbers placeholders per dialect, so no dialect_sql! here + let mut qb = sqlx::QueryBuilder::new( + "INSERT INTO billing (ak, product, tenant, model, served_model, protocol, + account, prompt_tokens, completion_tokens, total_tokens, cost_micros, + vendor_cost_micros, ptu_spillover, user_id, request_id, + created_at_epoch_secs, estimated, billed_units) ", + ); + qb.push_values($rows, |mut v, r| { + v.push_bind(&r.ak) + .push_bind(&r.product) + .push_bind(&r.tenant) + .push_bind(&r.model) + .push_bind(&r.served_model) + .push_bind(&r.protocol) + .push_bind(&r.account) + .push_bind(r.prompt_tokens) + .push_bind(r.completion_tokens) + .push_bind(r.total_tokens) + .push_bind(r.cost_micros) + .push_bind(r.vendor_cost_micros) + .push_bind(r.ptu_spillover) + .push_bind(&r.user_id) + .push_bind(&r.request_id) + .push_bind(r.created_at_epoch_secs) + .push_bind(r.estimated) + .push_bind(r.billed_units); + }); + qb.push(" ON CONFLICT (request_id) DO NOTHING"); + qb + }}; +} + macro_rules! q_rollup_insert { ($d:ident) => { concat!( @@ -1574,34 +1641,7 @@ macro_rules! sql_store_impl { if rows.is_empty() { return Ok(()); } - // the QueryBuilder numbers placeholders per dialect, so no dialect_sql! here - let mut qb = sqlx::QueryBuilder::new( - "INSERT INTO billing (ak, product, tenant, model, served_model, protocol, - account, prompt_tokens, completion_tokens, total_tokens, cost_micros, - vendor_cost_micros, ptu_spillover, user_id, request_id, - created_at_epoch_secs, estimated, billed_units) ", - ); - qb.push_values(rows, |mut v, r| { - v.push_bind(&r.ak) - .push_bind(&r.product) - .push_bind(&r.tenant) - .push_bind(&r.model) - .push_bind(&r.served_model) - .push_bind(&r.protocol) - .push_bind(&r.account) - .push_bind(r.prompt_tokens) - .push_bind(r.completion_tokens) - .push_bind(r.total_tokens) - .push_bind(r.cost_micros) - .push_bind(r.vendor_cost_micros) - .push_bind(r.ptu_spillover) - .push_bind(&r.user_id) - .push_bind(&r.request_id) - .push_bind(r.created_at_epoch_secs) - .push_bind(r.estimated) - .push_bind(r.billed_units); - }); - qb.push(" ON CONFLICT (request_id) DO NOTHING") + billing_insert!(rows) .build() .execute(&self.pool) .await @@ -1609,6 +1649,55 @@ macro_rules! sql_store_impl { Ok(()) } + async fn ledger_replay_batch(&self, rows: &[BillingRecord]) -> GResult<()> { + if rows.is_empty() { + return Ok(()); + } + let mut tx = self + .pool + .begin() + .await + .map_err(|e| crate::sqlx_err("begin replay", e))?; + let mut insert = billing_insert!(rows); + let inserted: Vec = insert + .push(" RETURNING request_id") + .build_query_scalar() + .fetch_all(&mut *tx) + .await + .map_err(|e| crate::sqlx_err("replay billing records", e))?; + if !inserted.is_empty() { + let mut add = sqlx::QueryBuilder::new(concat!( + "INSERT INTO usage_rollup (minute_epoch, tenant, user_id, model, requests, \ + prompt_tokens, completion_tokens, total_tokens, cost_micros, vendor_cost_micros, \ + billed_units) SELECT (created_at_epoch_secs/60)*60, tenant, user_id, model, COUNT(*), ", + usage_sums!($dialect), + " FROM billing WHERE created_at_epoch_secs < (", + )); + add.push(ROLLUP_WATERMARK_SQL).push(") AND request_id IN ("); + let mut ids = add.separated(", "); + for id in &inserted { + ids.push_bind(id); + } + add.push( + ") GROUP BY 1, 2, 3, 4 ON CONFLICT (minute_epoch, tenant, user_id, model) DO UPDATE SET \ + requests = usage_rollup.requests + excluded.requests, \ + prompt_tokens = usage_rollup.prompt_tokens + excluded.prompt_tokens, \ + completion_tokens = usage_rollup.completion_tokens + excluded.completion_tokens, \ + total_tokens = usage_rollup.total_tokens + excluded.total_tokens, \ + cost_micros = usage_rollup.cost_micros + excluded.cost_micros, \ + vendor_cost_micros = usage_rollup.vendor_cost_micros + excluded.vendor_cost_micros, \ + billed_units = usage_rollup.billed_units + excluded.billed_units", + ) + .build() + .execute(&mut *tx) + .await + .map_err(|e| crate::sqlx_err("roll up replayed records", e))?; + } + tx.commit() + .await + .map_err(|e| crate::sqlx_err("commit replay", e)) + } + fn defers_ledger_writes(&self) -> bool { true } @@ -1909,14 +1998,14 @@ macro_rules! sql_store_impl { } sql_store_impl!(SqliteStore, sqlite, { - async fn usage_rollup_since(&self, since: Option, now: i64) -> GResult { + async fn usage_rollup_advance(&self, now: i64) -> GResult { let hi = bucket_floor(now - ROLLUP_SETTLE_SECS); let watermark: i64 = sqlx::query_scalar(ROLLUP_WATERMARK_SQL) .fetch_one(&self.pool) .await .map_err(|e| crate::sqlx_err("read rollup watermark", e))?; let r = sqlx::query(dialect_sql!(sqlite, q_rollup_insert!(sqlite))) - .bind(rollup_floor(hi, watermark, since)) + .bind((hi - ROLLUP_BACKFILL_SECS).min(watermark)) .bind(hi) .execute(&self.pool) .await @@ -2220,35 +2309,27 @@ impl PostgresStore { } sql_store_impl!(PostgresStore, postgres, { - async fn usage_rollup_since(&self, since: Option, now: i64) -> GResult { + async fn usage_rollup_advance(&self, now: i64) -> GResult { let hi = bucket_floor(now - ROLLUP_SETTLE_SECS); let mut tx = self .pool .begin() .await .map_err(|e| crate::sqlx_err("begin rollup", e))?; - if since.is_none() { - let locked: bool = sqlx::query_scalar("SELECT pg_try_advisory_xact_lock($1)") - .bind(ROLLUP_LOCK_KEY) - .fetch_one(&mut *tx) - .await - .map_err(|e| crate::sqlx_err("acquire rollup lock", e))?; - if !locked { - return Ok(0); - } - } else { - sqlx::query("SELECT pg_advisory_xact_lock($1)") - .bind(ROLLUP_LOCK_KEY) - .execute(&mut *tx) - .await - .map_err(|e| crate::sqlx_err("wait for rollup lock", e))?; + let locked: bool = sqlx::query_scalar("SELECT pg_try_advisory_xact_lock($1)") + .bind(ROLLUP_LOCK_KEY) + .fetch_one(&mut *tx) + .await + .map_err(|e| crate::sqlx_err("acquire rollup lock", e))?; + if !locked { + return Ok(0); } let watermark: i64 = sqlx::query_scalar(ROLLUP_WATERMARK_SQL) .fetch_one(&mut *tx) .await .map_err(|e| crate::sqlx_err("read rollup watermark", e))?; let r = sqlx::query(dialect_sql!(postgres, q_rollup_insert!(postgres))) - .bind(rollup_floor(hi, watermark, since)) + .bind((hi - ROLLUP_BACKFILL_SECS).min(watermark)) .bind(hi) .execute(&mut *tx) .await @@ -2768,11 +2849,6 @@ where ) } -fn rollup_floor(hi: i64, watermark: i64, since: Option) -> i64 { - let lo = (hi - ROLLUP_BACKFILL_SECS).min(watermark); - since.map_or(lo, |since| lo.min(bucket_floor(since))) -} - fn bucket_floor(ts: i64) -> i64 { ts - ts.rem_euclid(ROLLUP_BUCKET_SECS) } @@ -3466,33 +3542,6 @@ mod tests { vec![(0, 2), (1_200, 1)], "rolled buckets and the raw tail group by bucket start" ); - - let mut late = record("m1"); - late.tenant = tenant.clone(); - late.user_id = "dave".into(); - late.created_at_epoch_secs = 100; - store.ledger_add(&late).await.unwrap(); - let billed_late = |rows: Vec| rows.iter().any(|r| r.user_id == "dave"); - store.usage_rollup_advance(1_500).await.unwrap(); - assert!( - !billed_late( - store - .usage_by_user(Some(&tenant), None, 0, 1_000) - .await - .unwrap() - ), - "a row billed below the watermark misses the trailing rollup" - ); - store.usage_rollup_since(Some(100), 1_500).await.unwrap(); - assert!( - billed_late( - store - .usage_by_user(Some(&tenant), None, 0, 1_000) - .await - .unwrap() - ), - "a rollup since its billing time counts the late row" - ); } #[tokio::test] @@ -3509,6 +3558,69 @@ mod tests { exercise_rollup(&store, "").await; } + async fn exercise_replay(store: &dyn Store, ns: &str, prune: impl AsyncFnOnce(&str)) { + let tenant = format!("rp{ns}"); + let billed = |at: i64| BillingRecord { + tenant: tenant.clone(), + user_id: "late".into(), + created_at_epoch_secs: at, + ..record("m1") + }; + let rolled = billed(400); + store.ledger_add(&rolled).await.unwrap(); + store.usage_rollup_advance(1_500).await.unwrap(); + prune(&rolled.request_id).await; + let late = billed(410); + store + .ledger_replay_batch(std::slice::from_ref(&late)) + .await + .unwrap(); + store + .ledger_replay_batch(std::slice::from_ref(&late)) + .await + .unwrap(); + store.usage_rollup_advance(1_500).await.unwrap(); + let usage = store + .usage_by_user(Some(&tenant), Some("late"), 0, 1_000) + .await + .unwrap(); + assert_eq!( + usage + .iter() + .map(|u| (u.requests, u.cost_micros)) + .collect::>(), + vec![(2, 2 * rolled.cost_micros)], + "{usage:?}" + ); + } + + #[tokio::test] + async fn memory_replay_adds_late_rows_to_a_pruned_bucket() { + let store = MemoryStore::default(); + exercise_replay(&store, "", async |id: &str| { + let mut ledger = lock(&store.ledger); + ledger.rows.retain(|r| r.request_id != id); + ledger.request_ids.remove(id); + }) + .await; + } + + #[tokio::test] + async fn sqlite_replay_adds_late_rows_to_a_pruned_bucket() { + let dir = tempfile::tempdir().unwrap(); + let store = SqliteStore::open(dir.path().join("replay.db").to_str().unwrap()) + .await + .unwrap(); + exercise_replay(&store, "", async |id: &str| { + sqlx::query("DELETE FROM billing WHERE request_id = ?1") + .bind(id) + .execute(&store.pool) + .await + .unwrap(); + }) + .await; + } + #[tokio::test] async fn rollup_survives_ledger_prune() { let store = MemoryStore::with_ledger_cap(1); @@ -3881,6 +3993,14 @@ mod tests { let store = PostgresStore::connect(&own_url).await.expect("pg connect"); let ns = format!("-{nonce}"); exercise_rollup(&store, &ns).await; + exercise_replay(&store, &ns, async |id: &str| { + sqlx::query("DELETE FROM billing WHERE request_id = $1") + .bind(id) + .execute(&store.pool) + .await + .unwrap(); + }) + .await; exercise_erase(&store, &ns).await; exercise_content_list(&store, &ns).await; exercise_terminal_idempotence(&store, &ns).await; diff --git a/docs/observability.md b/docs/observability.md index 7cc2840..1bfd6b3 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -102,7 +102,9 @@ With `storage.ledger_spool_path` set, the batch instead goes to that local JSON-lines file after three attempts (about 0.3 s; the directory is created on first use, and a spool that cannot be written falls back to the remaining retries), and a replay every 30 s inserts it in batches once the store accepts -writes again, then re-rolls the usage rollup back to the oldest replayed row +writes again; in the same transaction it adds the rows it inserted to any usage +rollup bucket the watermark already passed, so a late row counts even after its +bucket's other rows were pruned, and a repeated replay adds nothing (`billing spool replayed` in the log). The spool path is per-instance local disk; replicas must not share one file. A full queue applies backpressure to settlement until the repair worker makes From edc7077c228bf789a12cb28519924b10df2717bf Mon Sep 17 00:00:00 2001 From: CMGS Date: Sat, 17 Oct 2026 23:03:17 +0800 Subject: [PATCH 18/22] perf: compute usage tokens synchronously, bill at the call site bill_usage was an async wrapper around a pure computation plus bill(), which added a future state machine to every settle; perf stat on a request with no budgets measured ~420 instructions over main. usage_tokens now only computes the BillTokens and each caller awaits bill() itself: +131 instructions per request against main (0.08%, inside main's own run-to-run spread of 361). --- crates/dag/src/nodes.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index a148a3e..57d4f3a 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -691,7 +691,7 @@ impl DagNode for CostCalc { if resp.aborted && resp.completion_tokens == 0 { return bill_aborted_stream(ctx, None).await; } - bill_usage(ctx).await + bill(ctx, usage_tokens(ctx)?, false).await } } @@ -762,7 +762,7 @@ pub async fn settle_deferred_stream(ctx: &mut DagContext, delivery: StreamDelive } outcome.response.message.clear(); outcome.response.aborted = true; - bill_usage(ctx).await + bill(ctx, usage_tokens(ctx)?, false).await } StreamDelivery::None => { ctx.refund_reserves().await; @@ -810,7 +810,7 @@ async fn bill_aborted_stream( outcome.response.common_usage = common; outcome.response.raw_usage = Some(usage); ctx.decide("cost_calc", "aborted stream, billed the provider's record"); - return bill_usage(ctx).await; + return bill(ctx, usage_tokens(ctx)?, false).await; } let response = &ctx .outcome @@ -850,7 +850,7 @@ fn blank_stream_response(response: &mut gw_models::GatewayResponse) { response.aborted = true; } -async fn bill_usage(ctx: &mut DagContext) -> GResult<()> { +fn usage_tokens(ctx: &DagContext) -> GResult { let resp = &ctx .outcome .as_ref() @@ -886,7 +886,7 @@ async fn bill_usage(ctx: &mut DagContext) -> GResult<()> { None => BillTokens::weighted(resp.prompt_tokens, resp.completion_tokens, &rate), }; tokens.units = resp.billed_units; - bill(ctx, tokens, false).await + Ok(tokens) } /// Settle reserves and write the ledger; `estimated` marks an aborted stream's estimated counts. From c0146b9d6f791b312ea526bba3def6b2f486cd42 Mon Sep 17 00:00:00 2001 From: CMGS Date: Sat, 17 Oct 2026 23:14:35 +0800 Subject: [PATCH 19/22] billing: bound the image unit hold, resume a replay where it stopped - reserve_units took a client image `n` unbounded: n * unit price saturated the hold at i64::MAX, which in the memory governance pinned the counter at MAX and the settle then wiped the day's earlier spend. `n` is clamped to MAX_RESERVE like the token estimate. - Replay restarted at offset 0 after a failed batch, so rows a committed batch inserted could be pruned before the retry, re-insert, and be added to the rollup twice. Each committed batch now saves the next offset beside the .replaying file and a replay resumes there; a fresh claim clears it. - Two rollup comments that said a rolled minute's rows only ever shrink now name the replay; the rollup INSERT head is shared through rollup_select!; the unit-hold test bounds its second call too. --- crates/dag/src/nodes.rs | 9 ++++- crates/handler/src/lib.rs | 7 +++- crates/state/src/admission.rs | 74 +++++++++++++++++++++++++++++------ crates/state/src/store.rs | 29 ++++++++------ 4 files changed, 92 insertions(+), 27 deletions(-) diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index 57d4f3a..58056c8 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -265,7 +265,7 @@ fn reserve_parts(req: &gw_models::GatewayRequest) -> (i64, i64) { fn reserve_units(req: &gw_models::GatewayRequest) -> i64 { match req.model_param_v2.as_ref().and_then(|p| p.typed.as_ref()) { - Some(gw_models::TypedParams::Image(p)) => p.n.max(1), + Some(gw_models::TypedParams::Image(p)) => p.n.clamp(1, MAX_RESERVE), Some(gw_models::TypedParams::AudioTts(p)) => p.input.chars().count() as i64, Some(gw_models::TypedParams::Video(_)) => 0, _ => 1, @@ -1194,6 +1194,13 @@ mod tests { 5 ); assert_eq!(units(TypedParams::Video(VideoParams::default())), 0); + assert_eq!( + units(TypedParams::Image(ImageParams { + n: i64::MAX, + ..Default::default() + })), + super::MAX_RESERVE + ); } #[test] diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index 6cd55a1..1e627c7 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -2904,7 +2904,12 @@ mod tests { .await .expect("the first call holds its unit price"); - let second = h.run(search(), ak.clone()).await; + let second = tokio::time::timeout( + std::time::Duration::from_secs(5), + h.run(search(), ak.clone()), + ) + .await + .expect("a denial returns without waiting on the upstream"); let err = second.err().expect("the in-flight hold leaves no budget"); assert!( diff --git a/crates/state/src/admission.rs b/crates/state/src/admission.rs index f893f4a..ff824dd 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -306,16 +306,21 @@ impl BillingLedger { struct LedgerSpool { path: PathBuf, replaying: PathBuf, + replayed: PathBuf, lock: Mutex<()>, } impl LedgerSpool { fn new(path: &str) -> Self { let path = PathBuf::from(path); - let replaying = path.with_extension("replaying"); + let (replaying, replayed) = ( + path.with_extension("replaying"), + path.with_extension("replayed"), + ); Self { path, replaying, + replayed, lock: Mutex::new(()), } } @@ -340,31 +345,37 @@ impl LedgerSpool { .await } - /// `request_id` uniqueness makes a re-insert of an already accepted batch a no-op. + /// Resumes after the last batch it saved; `request_id` uniqueness covers one committed but not saved. async fn replay(&self, store: &Arc) { let claimed = { - let (path, replaying) = (self.path.clone(), self.replaying.clone()); + let (path, replaying, replayed) = ( + self.path.clone(), + self.replaying.clone(), + self.replayed.clone(), + ); let _held = self.lock.lock().await; run_blocking(move || { if replaying.exists() { - return Ok(true); + let saved = std::fs::read_to_string(&replayed).unwrap_or_default(); + return Ok(Some(saved.trim().parse().unwrap_or(0))); } if !path.exists() { - return Ok(false); + return Ok(None); } - std::fs::rename(&path, &replaying).map(|()| true) + remove_if_present(&replayed)?; + std::fs::rename(&path, &replaying).map(|()| Some(0)) }) .await }; - match claimed { - Ok(true) => {} - Ok(false) => return, + let mut offset = match claimed { + Ok(Some(offset)) => offset, + Ok(None) => return, Err(e) => { tracing::error!(error = %e, path = %self.path.display(), "billing spool claim failed"); return; } - } - let (mut offset, mut replayed) = (0, 0); + }; + let mut replayed = 0; loop { let replaying = self.replaying.clone(); let (rows, next) = match run_blocking(move || read_spool_batch(&replaying, offset)) @@ -385,9 +396,17 @@ impl LedgerSpool { } replayed += rows.len(); offset = next; + let saved = self.replayed.clone(); + if let Err(e) = run_blocking(move || std::fs::write(saved, offset.to_string())).await { + tracing::warn!(error = %e, "billing spool offset save failed"); + } } - let replaying = self.replaying.clone(); - if let Err(e) = run_blocking(move || std::fs::remove_file(replaying)).await { + let (replaying, saved) = (self.replaying.clone(), self.replayed.clone()); + if let Err(e) = run_blocking(move || { + std::fs::remove_file(replaying).and_then(|()| remove_if_present(&saved)) + }) + .await + { tracing::error!(error = %e, path = %self.replaying.display(), "billing spool remove failed"); return; } @@ -981,6 +1000,13 @@ fn read_spool_batch(path: &Path, offset: u64) -> std::io::Result<(Vec std::io::Result<()> { + match std::fs::remove_file(path) { + Err(e) if e.kind() != std::io::ErrorKind::NotFound => Err(e), + _ => Ok(()), + } +} + async fn run_blocking(f: F) -> std::io::Result where T: Send + 'static, @@ -1343,6 +1369,28 @@ mod tests { assert!(!spool.replaying.exists()); } + #[tokio::test] + async fn replay_resumes_after_the_last_saved_batch() { + let dir = tempfile::tempdir().unwrap(); + let spool = LedgerSpool::new(dir.path().join("ledger-spool.jsonl").to_str().unwrap()); + let mut lines = Vec::new(); + for id in ["req-committed", "req-pending"] { + serde_json::to_writer(&mut lines, &record(id)).unwrap(); + lines.push(b'\n'); + } + let committed = lines.iter().position(|&b| b == b'\n').unwrap() + 1; + std::fs::write(&spool.replaying, &lines).unwrap(); + std::fs::write(&spool.replayed, committed.to_string()).unwrap(); + let store = Arc::new(crate::MemoryStore::default()); + let as_store: Arc = store.clone(); + + spool.replay(&as_store).await; + + let (count, rows) = store.ledger_snapshot(usize::MAX).await.unwrap(); + assert_eq!((count, rows[0].request_id.as_str()), (1, "req-pending")); + assert!(!spool.replaying.exists() && !spool.replayed.exists()); + } + #[tokio::test] async fn a_spool_row_from_an_older_release_still_replays() { let dir = tempfile::tempdir().unwrap(); diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index bcedf0d..53447e7 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -35,8 +35,8 @@ const ROLLUP_BUCKET_SECS: i64 = 60; const ROLLUP_BACKFILL_SECS: i64 = 20 * 60; /// A minute is rolled only once closed for this long, so no in-flight write or -/// trailing replica lands a row in a rolled minute — a rolled minute's source -/// set can only shrink, which keeps the max-upsert sound. +/// trailing replica lands a row in a rolled minute; a spool replay that does adds +/// the row to the bucket in the same transaction, which keeps the max-upsert sound. const ROLLUP_SETTLE_SECS: i64 = ROLLUP_BUCKET_SECS; /// In-process rollup retention: the monthly-budget window, then a bucket cap. const ROLLUP_RETENTION_SECS: i64 = 62 * 86_400; @@ -291,10 +291,10 @@ impl UserUsageRow { self.billed_units = self.billed_units.saturating_add(r.billed_units); } - /// Keep the larger of each counter. A bucket's source rows are append-only - /// and only ever shrink via ledger pruning, so per-column max always picks - /// the more complete snapshot — a recompute from a partially pruned ledger - /// can never shrink a bucket. + /// Keep the larger of each counter. A bucket's source rows only shrink via ledger + /// pruning, and a replayed late row is added to its bucket as it lands, so per-column + /// max always picks the more complete snapshot — a recompute from a partially pruned + /// ledger can never shrink a bucket. fn keep_max(&mut self, o: &UserUsageRow) { self.requests = self.requests.max(o.requests); self.prompt_tokens = self.prompt_tokens.max(o.prompt_tokens); @@ -1448,13 +1448,21 @@ macro_rules! billing_insert { }}; } -macro_rules! q_rollup_insert { +macro_rules! rollup_select { ($d:ident) => { concat!( "INSERT INTO usage_rollup (minute_epoch, tenant, user_id, model, requests, \ prompt_tokens, completion_tokens, total_tokens, cost_micros, vendor_cost_micros, \ billed_units) SELECT (created_at_epoch_secs/60)*60, tenant, user_id, model, COUNT(*), ", - usage_sums!($d), + usage_sums!($d) + ) + }; +} + +macro_rules! q_rollup_insert { + ($d:ident) => { + concat!( + rollup_select!($d), " FROM billing WHERE created_at_epoch_secs >= ?1 AND created_at_epoch_secs < ?2 \ GROUP BY 1, 2, 3, 4 ON CONFLICT (minute_epoch, tenant, user_id, model) DO UPDATE SET requests = ", sql_greatest!($d), "(usage_rollup.requests, excluded.requests), prompt_tokens = ", @@ -1667,10 +1675,7 @@ macro_rules! sql_store_impl { .map_err(|e| crate::sqlx_err("replay billing records", e))?; if !inserted.is_empty() { let mut add = sqlx::QueryBuilder::new(concat!( - "INSERT INTO usage_rollup (minute_epoch, tenant, user_id, model, requests, \ - prompt_tokens, completion_tokens, total_tokens, cost_micros, vendor_cost_micros, \ - billed_units) SELECT (created_at_epoch_secs/60)*60, tenant, user_id, model, COUNT(*), ", - usage_sums!($dialect), + rollup_select!($dialect), " FROM billing WHERE created_at_epoch_secs < (", )); add.push(ROLLUP_WATERMARK_SQL).push(") AND request_id IN ("); From d93c54a4668d2d26afd19876f71f833f89999837 Mon Sep 17 00:00:00 2001 From: CMGS Date: Sat, 17 Oct 2026 23:20:06 +0800 Subject: [PATCH 20/22] review: name the replay offset file and row count apart The replay used replayed and saved for both a path and a value; a fix commit's naming nit from the c0146b9d review, no behavior change. --- crates/state/src/admission.rs | 28 +++++++++++++--------------- 1 file changed, 13 insertions(+), 15 deletions(-) diff --git a/crates/state/src/admission.rs b/crates/state/src/admission.rs index ff824dd..87ac028 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -313,14 +313,10 @@ struct LedgerSpool { impl LedgerSpool { fn new(path: &str) -> Self { let path = PathBuf::from(path); - let (replaying, replayed) = ( - path.with_extension("replaying"), - path.with_extension("replayed"), - ); Self { + replaying: path.with_extension("replaying"), + replayed: path.with_extension("replayed"), path, - replaying, - replayed, lock: Mutex::new(()), } } @@ -356,8 +352,8 @@ impl LedgerSpool { let _held = self.lock.lock().await; run_blocking(move || { if replaying.exists() { - let saved = std::fs::read_to_string(&replayed).unwrap_or_default(); - return Ok(Some(saved.trim().parse().unwrap_or(0))); + let text = std::fs::read_to_string(&replayed).unwrap_or_default(); + return Ok(Some(text.trim().parse().unwrap_or(0))); } if !path.exists() { return Ok(None); @@ -375,7 +371,7 @@ impl LedgerSpool { return; } }; - let mut replayed = 0; + let mut total = 0; loop { let replaying = self.replaying.clone(); let (rows, next) = match run_blocking(move || read_spool_batch(&replaying, offset)) @@ -394,23 +390,25 @@ impl LedgerSpool { tracing::warn!(error = %e, rows = rows.len(), "billing spool replay deferred"); return; } - replayed += rows.len(); + total += rows.len(); offset = next; - let saved = self.replayed.clone(); - if let Err(e) = run_blocking(move || std::fs::write(saved, offset.to_string())).await { + let offset_file = self.replayed.clone(); + if let Err(e) = + run_blocking(move || std::fs::write(offset_file, offset.to_string())).await + { tracing::warn!(error = %e, "billing spool offset save failed"); } } - let (replaying, saved) = (self.replaying.clone(), self.replayed.clone()); + let (replaying, offset_file) = (self.replaying.clone(), self.replayed.clone()); if let Err(e) = run_blocking(move || { - std::fs::remove_file(replaying).and_then(|()| remove_if_present(&saved)) + std::fs::remove_file(replaying).and_then(|()| remove_if_present(&offset_file)) }) .await { tracing::error!(error = %e, path = %self.replaying.display(), "billing spool remove failed"); return; } - tracing::info!(rows = replayed, "billing spool replayed"); + tracing::info!(rows = total, "billing spool replayed"); } } From 5f6cb2a2551eb701c99c0fb6cde1e47d93c15dec Mon Sep 17 00:00:00 2001 From: CMGS Date: Sun, 18 Oct 2026 00:28:26 +0800 Subject: [PATCH 21/22] billing: serialize replay and retain pruned request ids Use the rollup transaction lock for replay and pruning. Retain pruned request ids atomically so an unacknowledged commit cannot be counted again after raw ledger retention. Keep normal billing inserts unchanged. Cover rollup overlap and replay after actual pruning on PostgreSQL and SQLite. Make parallel test database names unique. --- crates/state/src/lib.rs | 7 +- crates/state/src/store.rs | 234 ++++++++++++++++++++++++++++++++++---- docs/deployment.md | 3 +- 3 files changed, 221 insertions(+), 23 deletions(-) diff --git a/crates/state/src/lib.rs b/crates/state/src/lib.rs index 51b04d1..04cf34f 100644 --- a/crates/state/src/lib.rs +++ b/crates/state/src/lib.rs @@ -1082,11 +1082,16 @@ fn settle_on(counter: &mut i64, delta: i64) { #[cfg(test)] pub(crate) async fn scratch_pg(url: &str) -> (String, u128) { + static NEXT: AtomicUsize = AtomicUsize::new(0); let nonce = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .unwrap() .as_nanos(); - let db = format!("gwtest_{nonce}"); + let db = format!( + "gwtest_{nonce}_{}_{}", + std::process::id(), + NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + ); let admin = sqlx::PgPool::connect(url).await.expect("pg admin"); sqlx::query(sqlx::AssertSqlSafe(format!("CREATE DATABASE {db}"))) .execute(&admin) diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index 53447e7..4af5d33 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -42,8 +42,6 @@ const ROLLUP_SETTLE_SECS: i64 = ROLLUP_BUCKET_SECS; const ROLLUP_RETENTION_SECS: i64 = 62 * 86_400; const ROLLUP_MAX_BUCKETS: usize = 1_000_000; -/// Postgres advisory-lock key serializing the fleet's rollup: one replica -/// advances per tick (the upsert is idempotent; the lock only avoids repeated scans). const ROLLUP_LOCK_KEY: i64 = 0x6777_726f_6c6c; /// One minute past the newest rolled bucket — where the raw ledger tail @@ -360,7 +358,7 @@ pub trait Store: Send + Sync + std::fmt::Debug { } Ok(()) } - /// Insert late rows and atomically add the ones inserted to rollup buckets the watermark passed. + /// Insert late rows and atomically roll up requests not already counted before pruning. async fn ledger_replay_batch(&self, rows: &[BillingRecord]) -> GResult<()>; /// Whether rows may batch off the request path; in-process backends stay synchronous. fn defers_ledger_writes(&self) -> bool { @@ -1170,6 +1168,7 @@ impl SqliteStore { billed_units INTEGER NOT NULL DEFAULT 0)", "CREATE INDEX IF NOT EXISTS billing_created_idx ON billing (created_at_epoch_secs)", "CREATE UNIQUE INDEX IF NOT EXISTS billing_request_uidx ON billing (request_id)", + "CREATE TABLE IF NOT EXISTS billing_pruned (request_id TEXT PRIMARY KEY)", "CREATE TABLE IF NOT EXISTS erasures ( tenant TEXT NOT NULL, user_id TEXT NOT NULL, erased_at_epoch_millis INTEGER NOT NULL, @@ -1448,6 +1447,25 @@ macro_rules! billing_insert { }}; } +macro_rules! lock_rollup { + (sqlite, $tx:ident) => {}; + (postgres, $tx:ident) => { + sqlx::query("SELECT pg_advisory_xact_lock($1)") + .bind(ROLLUP_LOCK_KEY) + .execute(&mut *$tx) + .await + .map_err(|e| crate::sqlx_err("acquire rollup lock", e))?; + }; +} + +macro_rules! prunable_billing { + () => { + " FROM billing WHERE n <= (SELECT MAX(n) FROM billing) - ?1 + AND created_at_epoch_secs < + (SELECT COALESCE(MAX(minute_epoch), -60) + 60 FROM usage_rollup)" + }; +} + macro_rules! rollup_select { ($d:ident) => { concat!( @@ -1666,6 +1684,7 @@ macro_rules! sql_store_impl { .begin() .await .map_err(|e| crate::sqlx_err("begin replay", e))?; + lock_rollup!($dialect, tx); let mut insert = billing_insert!(rows); let inserted: Vec = insert .push(" RETURNING request_id") @@ -1684,7 +1703,8 @@ macro_rules! sql_store_impl { ids.push_bind(id); } add.push( - ") GROUP BY 1, 2, 3, 4 ON CONFLICT (minute_epoch, tenant, user_id, model) DO UPDATE SET \ + ") AND NOT EXISTS (SELECT 1 FROM billing_pruned p WHERE p.request_id = billing.request_id) \ + GROUP BY 1, 2, 3, 4 ON CONFLICT (minute_epoch, tenant, user_id, model) DO UPDATE SET \ requests = usage_rollup.requests + excluded.requests, \ prompt_tokens = usage_rollup.prompt_tokens + excluded.prompt_tokens, \ completion_tokens = usage_rollup.completion_tokens + excluded.completion_tokens, \ @@ -1707,23 +1727,6 @@ macro_rules! sql_store_impl { true } - async fn ledger_prune(&self) -> GResult { - if self.ledger_max_rows == 0 { - return Ok(0); - } - let done = sqlx::query(dialect_sql!( - $dialect, - "DELETE FROM billing WHERE n <= (SELECT MAX(n) FROM billing) - ? - AND created_at_epoch_secs < - (SELECT COALESCE(MAX(minute_epoch), -60) + 60 FROM usage_rollup)" - )) - .bind(self.ledger_max_rows as i64) - .execute(&self.pool) - .await - .map_err(|e| crate::sqlx_err("prune billing records", e))?; - Ok(done.rows_affected()) - } - async fn ledger_snapshot(&self, limit: usize) -> GResult<(usize, Vec)> { let total: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM billing") .fetch_one(&self.pool) @@ -2003,6 +2006,35 @@ macro_rules! sql_store_impl { } sql_store_impl!(SqliteStore, sqlite, { + async fn ledger_prune(&self) -> GResult { + if self.ledger_max_rows == 0 { + return Ok(0); + } + let mut tx = self + .pool + .begin() + .await + .map_err(|e| crate::sqlx_err("begin prune", e))?; + sqlx::query(concat!( + "INSERT INTO billing_pruned SELECT request_id", + prunable_billing!(), + " ON CONFLICT (request_id) DO NOTHING" + )) + .bind(self.ledger_max_rows as i64) + .execute(&mut *tx) + .await + .map_err(|e| crate::sqlx_err("retain pruned billing ids", e))?; + let done = sqlx::query(concat!("DELETE", prunable_billing!())) + .bind(self.ledger_max_rows as i64) + .execute(&mut *tx) + .await + .map_err(|e| crate::sqlx_err("prune billing records", e))?; + tx.commit() + .await + .map_err(|e| crate::sqlx_err("commit prune", e))?; + Ok(done.rows_affected()) + } + async fn usage_rollup_advance(&self, now: i64) -> GResult { let hi = bucket_floor(now - ROLLUP_SETTLE_SECS); let watermark: i64 = sqlx::query_scalar(ROLLUP_WATERMARK_SQL) @@ -2223,6 +2255,7 @@ impl PostgresStore { billed_units BIGINT NOT NULL DEFAULT 0)", "CREATE INDEX IF NOT EXISTS billing_created_idx ON billing (created_at_epoch_secs)", "CREATE UNIQUE INDEX IF NOT EXISTS billing_request_uidx ON billing (request_id)", + "CREATE TABLE IF NOT EXISTS billing_pruned (request_id TEXT PRIMARY KEY)", "CREATE TABLE IF NOT EXISTS usage_rollup ( minute_epoch BIGINT NOT NULL, tenant TEXT NOT NULL, user_id TEXT NOT NULL, model TEXT NOT NULL, requests BIGINT NOT NULL, @@ -2314,6 +2347,36 @@ impl PostgresStore { } sql_store_impl!(PostgresStore, postgres, { + async fn ledger_prune(&self) -> GResult { + if self.ledger_max_rows == 0 { + return Ok(0); + } + let mut tx = self + .pool + .begin() + .await + .map_err(|e| crate::sqlx_err("begin prune", e))?; + lock_rollup!(postgres, tx); + let count: i64 = sqlx::query_scalar(dialect_sql!( + postgres, + concat!( + "WITH pruned AS (DELETE", + prunable_billing!(), + " RETURNING request_id), + retained AS (INSERT INTO billing_pruned SELECT request_id FROM pruned + ON CONFLICT (request_id) DO NOTHING) SELECT COUNT(*) FROM pruned" + ) + )) + .bind(self.ledger_max_rows as i64) + .fetch_one(&mut *tx) + .await + .map_err(|e| crate::sqlx_err("prune billing records", e))?; + tx.commit() + .await + .map_err(|e| crate::sqlx_err("commit prune", e))?; + Ok(count as u64) + } + async fn usage_rollup_advance(&self, now: i64) -> GResult { let hi = bucket_floor(now - ROLLUP_SETTLE_SECS); let mut tx = self @@ -3626,6 +3689,135 @@ mod tests { .await; } + async fn replay_after_prune_is_idempotent(store: &dyn Store) { + for (user, replay) in [("normal", false), ("replay", true)] { + let late = BillingRecord { + user_id: user.into(), + created_at_epoch_secs: 400, + ..record("late") + }; + if !replay { + store.ledger_add(&late).await.unwrap(); + } + let newer = |at| BillingRecord { + created_at_epoch_secs: at, + ..record("newer") + }; + store.ledger_add(&newer(7_000)).await.unwrap(); + store.usage_rollup_advance(10_000).await.unwrap(); + if replay { + store + .ledger_replay_batch(std::slice::from_ref(&late)) + .await + .unwrap(); + } + store.ledger_add(&newer(8_000)).await.unwrap(); + store.usage_rollup_advance(10_000).await.unwrap(); + assert!(store.ledger_prune().await.unwrap() > 0); + assert!( + store + .ledger_snapshot(usize::MAX) + .await + .unwrap() + .1 + .iter() + .all(|r| r.request_id != late.request_id) + ); + for _ in 0..2 { + store + .ledger_replay_batch(std::slice::from_ref(&late)) + .await + .unwrap(); + store.ledger_add(&newer(9_000)).await.unwrap(); + store.usage_rollup_advance(10_000).await.unwrap(); + assert!(store.ledger_prune().await.unwrap() > 0); + let usage = store + .usage_by_user(None, Some(user), 0, 1_000) + .await + .unwrap(); + assert_eq!(usage.len(), 1); + assert_eq!( + (usage[0].requests, usage[0].cost_micros), + (1, late.cost_micros) + ); + } + } + } + + #[tokio::test] + async fn sqlite_replay_does_not_rebill_pruned_records() { + let dir = tempfile::tempdir().unwrap(); + let store = SqliteStore::open_with_cap(dir.path().join("replay.db").to_str().unwrap(), 1) + .await + .unwrap(); + replay_after_prune_is_idempotent(&store).await; + } + + #[tokio::test] + async fn postgres_replay_does_not_rebill_pruned_records() { + let Ok(url) = std::env::var("GW_TEST_PG_URL") else { + return; + }; + let (url, _) = crate::scratch_pg(&url).await; + let store = PostgresStore::connect_with_cap(&url, 1, 0).await.unwrap(); + replay_after_prune_is_idempotent(&store).await; + } + + #[tokio::test] + async fn postgres_replay_waits_for_rollup_before_classifying_late_rows() { + let Ok(url) = std::env::var("GW_TEST_PG_URL") else { + return; + }; + let (url, _) = crate::scratch_pg(&url).await; + let store = std::sync::Arc::new(PostgresStore::connect(&url).await.unwrap()); + let billed = |at| BillingRecord { + created_at_epoch_secs: at, + ..record("late") + }; + store.ledger_add(&billed(300)).await.unwrap(); + store.usage_rollup_advance(420).await.unwrap(); + store.ledger_add(&billed(390)).await.unwrap(); + let mut rollup = store.pool.begin().await.unwrap(); + sqlx::query("SELECT pg_advisory_xact_lock($1)") + .bind(ROLLUP_LOCK_KEY) + .execute(&mut *rollup) + .await + .unwrap(); + let replay_store = store.clone(); + let late = billed(400); + let replay = tokio::spawn(async move { + replay_store.ledger_replay_batch(&[late]).await.unwrap(); + }); + crate::eventually("replay waiting for rollup", async || { + assert!( + !replay.is_finished(), + "replay bypassed the rollup transaction" + ); + sqlx::query_scalar::<_, bool>( + "SELECT EXISTS (SELECT 1 FROM pg_locks WHERE locktype = 'advisory' + AND database = (SELECT oid FROM pg_database WHERE datname = current_database()) + AND objid::bigint = $1 AND NOT granted)", + ) + .bind(ROLLUP_LOCK_KEY & 0xffff_ffff) + .fetch_one(&store.pool) + .await + .unwrap() + }) + .await; + sqlx::query(dialect_sql!(postgres, q_rollup_insert!(postgres))) + .bind(360_i64) + .bind(2_940_i64) + .execute(&mut *rollup) + .await + .unwrap(); + rollup.commit().await.unwrap(); + replay.await.unwrap(); + store.usage_rollup_advance(3_060).await.unwrap(); + let usage = store.usage_by_user(None, None, 360, 420).await.unwrap(); + assert_eq!(usage.len(), 1); + assert_eq!(usage[0].requests, 2); + } + #[tokio::test] async fn rollup_survives_ledger_prune() { let store = MemoryStore::with_ledger_cap(1); diff --git a/docs/deployment.md b/docs/deployment.md index b2aae4f..6e42232 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -93,7 +93,8 @@ storage: warning until it returns ([Governance](governance.md#limits)). - `ledger_max_rows` is not a hard cap: pruning spares rows not yet folded into the usage rollup, so the table can briefly exceed the cap under - rollup lag. + rollup lag. Pruned request IDs remain as compact deduplication records for + late spool replay. - With Postgres + Redis, the batching billing writer shares the connection pool with request-path reads; at the default `postgres_max_connections: 10` a single-node bench showed the big-body lane going bimodal (p99 9 ms vs From ecf05c22cea0cc176b1e4c1a834d1c264df41729 Mon Sep 17 00:00:00 2001 From: CMGS Date: Sun, 18 Oct 2026 00:38:09 +0800 Subject: [PATCH 22/22] Revert billing replay retention and locking Revert 5f6cb2a2551eb701c99c0fb6cde1e47d93c15dec except the scratch_pg test database naming fix. Remove permanent pruned-request storage, added replay/prune locks, and their tests and documentation. --- crates/state/src/store.rs | 234 ++++---------------------------------- docs/deployment.md | 3 +- 2 files changed, 22 insertions(+), 215 deletions(-) diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index 4af5d33..53447e7 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -42,6 +42,8 @@ const ROLLUP_SETTLE_SECS: i64 = ROLLUP_BUCKET_SECS; const ROLLUP_RETENTION_SECS: i64 = 62 * 86_400; const ROLLUP_MAX_BUCKETS: usize = 1_000_000; +/// Postgres advisory-lock key serializing the fleet's rollup: one replica +/// advances per tick (the upsert is idempotent; the lock only avoids repeated scans). const ROLLUP_LOCK_KEY: i64 = 0x6777_726f_6c6c; /// One minute past the newest rolled bucket — where the raw ledger tail @@ -358,7 +360,7 @@ pub trait Store: Send + Sync + std::fmt::Debug { } Ok(()) } - /// Insert late rows and atomically roll up requests not already counted before pruning. + /// Insert late rows and atomically add the ones inserted to rollup buckets the watermark passed. async fn ledger_replay_batch(&self, rows: &[BillingRecord]) -> GResult<()>; /// Whether rows may batch off the request path; in-process backends stay synchronous. fn defers_ledger_writes(&self) -> bool { @@ -1168,7 +1170,6 @@ impl SqliteStore { billed_units INTEGER NOT NULL DEFAULT 0)", "CREATE INDEX IF NOT EXISTS billing_created_idx ON billing (created_at_epoch_secs)", "CREATE UNIQUE INDEX IF NOT EXISTS billing_request_uidx ON billing (request_id)", - "CREATE TABLE IF NOT EXISTS billing_pruned (request_id TEXT PRIMARY KEY)", "CREATE TABLE IF NOT EXISTS erasures ( tenant TEXT NOT NULL, user_id TEXT NOT NULL, erased_at_epoch_millis INTEGER NOT NULL, @@ -1447,25 +1448,6 @@ macro_rules! billing_insert { }}; } -macro_rules! lock_rollup { - (sqlite, $tx:ident) => {}; - (postgres, $tx:ident) => { - sqlx::query("SELECT pg_advisory_xact_lock($1)") - .bind(ROLLUP_LOCK_KEY) - .execute(&mut *$tx) - .await - .map_err(|e| crate::sqlx_err("acquire rollup lock", e))?; - }; -} - -macro_rules! prunable_billing { - () => { - " FROM billing WHERE n <= (SELECT MAX(n) FROM billing) - ?1 - AND created_at_epoch_secs < - (SELECT COALESCE(MAX(minute_epoch), -60) + 60 FROM usage_rollup)" - }; -} - macro_rules! rollup_select { ($d:ident) => { concat!( @@ -1684,7 +1666,6 @@ macro_rules! sql_store_impl { .begin() .await .map_err(|e| crate::sqlx_err("begin replay", e))?; - lock_rollup!($dialect, tx); let mut insert = billing_insert!(rows); let inserted: Vec = insert .push(" RETURNING request_id") @@ -1703,8 +1684,7 @@ macro_rules! sql_store_impl { ids.push_bind(id); } add.push( - ") AND NOT EXISTS (SELECT 1 FROM billing_pruned p WHERE p.request_id = billing.request_id) \ - GROUP BY 1, 2, 3, 4 ON CONFLICT (minute_epoch, tenant, user_id, model) DO UPDATE SET \ + ") GROUP BY 1, 2, 3, 4 ON CONFLICT (minute_epoch, tenant, user_id, model) DO UPDATE SET \ requests = usage_rollup.requests + excluded.requests, \ prompt_tokens = usage_rollup.prompt_tokens + excluded.prompt_tokens, \ completion_tokens = usage_rollup.completion_tokens + excluded.completion_tokens, \ @@ -1727,6 +1707,23 @@ macro_rules! sql_store_impl { true } + async fn ledger_prune(&self) -> GResult { + if self.ledger_max_rows == 0 { + return Ok(0); + } + let done = sqlx::query(dialect_sql!( + $dialect, + "DELETE FROM billing WHERE n <= (SELECT MAX(n) FROM billing) - ? + AND created_at_epoch_secs < + (SELECT COALESCE(MAX(minute_epoch), -60) + 60 FROM usage_rollup)" + )) + .bind(self.ledger_max_rows as i64) + .execute(&self.pool) + .await + .map_err(|e| crate::sqlx_err("prune billing records", e))?; + Ok(done.rows_affected()) + } + async fn ledger_snapshot(&self, limit: usize) -> GResult<(usize, Vec)> { let total: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM billing") .fetch_one(&self.pool) @@ -2006,35 +2003,6 @@ macro_rules! sql_store_impl { } sql_store_impl!(SqliteStore, sqlite, { - async fn ledger_prune(&self) -> GResult { - if self.ledger_max_rows == 0 { - return Ok(0); - } - let mut tx = self - .pool - .begin() - .await - .map_err(|e| crate::sqlx_err("begin prune", e))?; - sqlx::query(concat!( - "INSERT INTO billing_pruned SELECT request_id", - prunable_billing!(), - " ON CONFLICT (request_id) DO NOTHING" - )) - .bind(self.ledger_max_rows as i64) - .execute(&mut *tx) - .await - .map_err(|e| crate::sqlx_err("retain pruned billing ids", e))?; - let done = sqlx::query(concat!("DELETE", prunable_billing!())) - .bind(self.ledger_max_rows as i64) - .execute(&mut *tx) - .await - .map_err(|e| crate::sqlx_err("prune billing records", e))?; - tx.commit() - .await - .map_err(|e| crate::sqlx_err("commit prune", e))?; - Ok(done.rows_affected()) - } - async fn usage_rollup_advance(&self, now: i64) -> GResult { let hi = bucket_floor(now - ROLLUP_SETTLE_SECS); let watermark: i64 = sqlx::query_scalar(ROLLUP_WATERMARK_SQL) @@ -2255,7 +2223,6 @@ impl PostgresStore { billed_units BIGINT NOT NULL DEFAULT 0)", "CREATE INDEX IF NOT EXISTS billing_created_idx ON billing (created_at_epoch_secs)", "CREATE UNIQUE INDEX IF NOT EXISTS billing_request_uidx ON billing (request_id)", - "CREATE TABLE IF NOT EXISTS billing_pruned (request_id TEXT PRIMARY KEY)", "CREATE TABLE IF NOT EXISTS usage_rollup ( minute_epoch BIGINT NOT NULL, tenant TEXT NOT NULL, user_id TEXT NOT NULL, model TEXT NOT NULL, requests BIGINT NOT NULL, @@ -2347,36 +2314,6 @@ impl PostgresStore { } sql_store_impl!(PostgresStore, postgres, { - async fn ledger_prune(&self) -> GResult { - if self.ledger_max_rows == 0 { - return Ok(0); - } - let mut tx = self - .pool - .begin() - .await - .map_err(|e| crate::sqlx_err("begin prune", e))?; - lock_rollup!(postgres, tx); - let count: i64 = sqlx::query_scalar(dialect_sql!( - postgres, - concat!( - "WITH pruned AS (DELETE", - prunable_billing!(), - " RETURNING request_id), - retained AS (INSERT INTO billing_pruned SELECT request_id FROM pruned - ON CONFLICT (request_id) DO NOTHING) SELECT COUNT(*) FROM pruned" - ) - )) - .bind(self.ledger_max_rows as i64) - .fetch_one(&mut *tx) - .await - .map_err(|e| crate::sqlx_err("prune billing records", e))?; - tx.commit() - .await - .map_err(|e| crate::sqlx_err("commit prune", e))?; - Ok(count as u64) - } - async fn usage_rollup_advance(&self, now: i64) -> GResult { let hi = bucket_floor(now - ROLLUP_SETTLE_SECS); let mut tx = self @@ -3689,135 +3626,6 @@ mod tests { .await; } - async fn replay_after_prune_is_idempotent(store: &dyn Store) { - for (user, replay) in [("normal", false), ("replay", true)] { - let late = BillingRecord { - user_id: user.into(), - created_at_epoch_secs: 400, - ..record("late") - }; - if !replay { - store.ledger_add(&late).await.unwrap(); - } - let newer = |at| BillingRecord { - created_at_epoch_secs: at, - ..record("newer") - }; - store.ledger_add(&newer(7_000)).await.unwrap(); - store.usage_rollup_advance(10_000).await.unwrap(); - if replay { - store - .ledger_replay_batch(std::slice::from_ref(&late)) - .await - .unwrap(); - } - store.ledger_add(&newer(8_000)).await.unwrap(); - store.usage_rollup_advance(10_000).await.unwrap(); - assert!(store.ledger_prune().await.unwrap() > 0); - assert!( - store - .ledger_snapshot(usize::MAX) - .await - .unwrap() - .1 - .iter() - .all(|r| r.request_id != late.request_id) - ); - for _ in 0..2 { - store - .ledger_replay_batch(std::slice::from_ref(&late)) - .await - .unwrap(); - store.ledger_add(&newer(9_000)).await.unwrap(); - store.usage_rollup_advance(10_000).await.unwrap(); - assert!(store.ledger_prune().await.unwrap() > 0); - let usage = store - .usage_by_user(None, Some(user), 0, 1_000) - .await - .unwrap(); - assert_eq!(usage.len(), 1); - assert_eq!( - (usage[0].requests, usage[0].cost_micros), - (1, late.cost_micros) - ); - } - } - } - - #[tokio::test] - async fn sqlite_replay_does_not_rebill_pruned_records() { - let dir = tempfile::tempdir().unwrap(); - let store = SqliteStore::open_with_cap(dir.path().join("replay.db").to_str().unwrap(), 1) - .await - .unwrap(); - replay_after_prune_is_idempotent(&store).await; - } - - #[tokio::test] - async fn postgres_replay_does_not_rebill_pruned_records() { - let Ok(url) = std::env::var("GW_TEST_PG_URL") else { - return; - }; - let (url, _) = crate::scratch_pg(&url).await; - let store = PostgresStore::connect_with_cap(&url, 1, 0).await.unwrap(); - replay_after_prune_is_idempotent(&store).await; - } - - #[tokio::test] - async fn postgres_replay_waits_for_rollup_before_classifying_late_rows() { - let Ok(url) = std::env::var("GW_TEST_PG_URL") else { - return; - }; - let (url, _) = crate::scratch_pg(&url).await; - let store = std::sync::Arc::new(PostgresStore::connect(&url).await.unwrap()); - let billed = |at| BillingRecord { - created_at_epoch_secs: at, - ..record("late") - }; - store.ledger_add(&billed(300)).await.unwrap(); - store.usage_rollup_advance(420).await.unwrap(); - store.ledger_add(&billed(390)).await.unwrap(); - let mut rollup = store.pool.begin().await.unwrap(); - sqlx::query("SELECT pg_advisory_xact_lock($1)") - .bind(ROLLUP_LOCK_KEY) - .execute(&mut *rollup) - .await - .unwrap(); - let replay_store = store.clone(); - let late = billed(400); - let replay = tokio::spawn(async move { - replay_store.ledger_replay_batch(&[late]).await.unwrap(); - }); - crate::eventually("replay waiting for rollup", async || { - assert!( - !replay.is_finished(), - "replay bypassed the rollup transaction" - ); - sqlx::query_scalar::<_, bool>( - "SELECT EXISTS (SELECT 1 FROM pg_locks WHERE locktype = 'advisory' - AND database = (SELECT oid FROM pg_database WHERE datname = current_database()) - AND objid::bigint = $1 AND NOT granted)", - ) - .bind(ROLLUP_LOCK_KEY & 0xffff_ffff) - .fetch_one(&store.pool) - .await - .unwrap() - }) - .await; - sqlx::query(dialect_sql!(postgres, q_rollup_insert!(postgres))) - .bind(360_i64) - .bind(2_940_i64) - .execute(&mut *rollup) - .await - .unwrap(); - rollup.commit().await.unwrap(); - replay.await.unwrap(); - store.usage_rollup_advance(3_060).await.unwrap(); - let usage = store.usage_by_user(None, None, 360, 420).await.unwrap(); - assert_eq!(usage.len(), 1); - assert_eq!(usage[0].requests, 2); - } - #[tokio::test] async fn rollup_survives_ledger_prune() { let store = MemoryStore::with_ledger_cap(1); diff --git a/docs/deployment.md b/docs/deployment.md index 6e42232..b2aae4f 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -93,8 +93,7 @@ storage: warning until it returns ([Governance](governance.md#limits)). - `ledger_max_rows` is not a hard cap: pruning spares rows not yet folded into the usage rollup, so the table can briefly exceed the cap under - rollup lag. Pruned request IDs remain as compact deduplication records for - late spool replay. + rollup lag. - With Postgres + Redis, the batching billing writer shares the connection pool with request-path reads; at the default `postgres_max_connections: 10` a single-node bench showed the big-body lane going bimodal (p99 9 ms vs