Repository navigation
billing: ledger spool, budget holds at admission, abandoned streams billed from the upstream's usage - #158
Merged
Conversation
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.
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.
…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.
BudgetGate re-derived the estimate parts, serializing a decisions request's fields a second time on every call.
- 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.
- 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.
…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.
…okup - 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<i64>, 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.
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.
…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.
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.
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.
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.
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.
…ne 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.
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.
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.
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).
- 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.
The replay used replayed and saved for both a path and a value; a fix commit's naming nit from the c0146b9 review, no behavior change.
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.
Revert 5f6cb2a except the scratch_pg test database naming fix. Remove permanent pruned-request storage, added replay/prune locks, and their tests and documentation.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Three billing fixes before paying users. They are independent; a rebase-merge keeps them as separate commits, followed by the review round's fixes.
storage.ledger_spool_pathset, a batch goes to a local JSON-lines file after three attempts (the directory is created on first use; a spool that cannot be written falls back to the remaining retries). A 30 s replay reads the spool in batches and resumes after the last batch it committed. Each batch goes throughStore::ledger_replay_batch, which in one transaction inserts the rows (ON CONFLICT (request_id) DO NOTHING RETURNING request_id) and adds only the inserted ones to rollup buckets the watermark already passed. A late row therefore counts even after its bucket's other rows were pruned, and a repeated replay adds nothing.BillingRecorddeserializes with defaults, so a spool written by an older release replays after an upgrade. Clientuserhints drop NUL: Postgres TEXT refuses it, and on main a request whoseuserholds a NUL character already fails admission with a 503 from the user-budget lookup.max_tokenssides at the tenant price, and the unit price of units known at admission: one search,nimages, the characters of speech, one unit for transcription and rerank. Billing settles the hold and a failed call returns it. The hold is priced only when a day budget applies. On Redis, every day budget is held with one pipeline ofINCRBYand settled with another, with no Lua. Monthly caps, realtime turns and video (billed at its poll) still check only.GET /api/v1/generation(3/5/10/20 s, each lookup capped at 10 s) and bills its native counts and cost, falling back to the estimate when no record lands. Completion includes reasoning: live, DeepInfra'stotal_costmatched 73 prompt + 34 completion tokens whilenative_tokens_reasoningread 37. The record carries no cache-write count, so a cache write bills at the prompt weight. These settles are tracked: shutdown cancels pending lookups, which bill the estimate, and waits for them before the final ledger flush. The access log'slatency_msstops at the end of the stream.Hot path
Measured on one host (cocoon-test1), main
385ed379vs this branch: mock upstream, arms interleaved with order swapped, andperf statinstruction counts per request (200k requests per arm), which do not depend on host load.Redis commands per request with two day budgets: 3 EVALSHA (main: 5), 0 GET (main: 2), 7 INCRBY (main: 5).
Verification
cargo fmt --check,cargo clippy --workspace --all-targets -- -D warnings,cargo test --workspace: 764 passed, 0 failed, 4 ignored. This ran on macOS and on Linux (rust 1.99.0) withGW_TEST_PG_URL(Postgres 17) andGW_TEST_REDIS_URL(Redis 7) set; the Postgres-gated tests created their scratch databases.userbills as stripped.billingaway spools rows; renaming it back replays them within 30 s.estimated = false).