Files
llmgw/internal/accounting/budget.go
T
Ilker ManapandClaude Opus 5 0b93fb2883 Enforce budgets and rate limits, and write usage off the request path
The second half of phase 5. A key now has a cost, a quota and a bill.

Rate limits use a sliding window rather than a fixed one. A fixed window lets
a caller send its whole minute's allowance in the last second of one window
and again in the first second of the next, so the observed rate is twice the
configured limit exactly when the system is least able to take it. Checking
and charging are one operation, or two concurrent requests both see room and
both take it.

Token limits are checked against an estimate, because the real count only
exists once the request has run; the difference is charged afterwards, so a
caller who consistently underestimates is still bounded, one window late. A
request with no max_tokens gets a modest assumption rather than zero —
assuming zero would let an unbounded request through every token limit.

Limits are checked before admission, not after: a request over its quota
should be refused rather than queued, where it would occupy a band's capacity
to no purpose. All refusals are 429, but a spent budget carries no Retry-After
because waiting does not help until the period rolls over, and saying
otherwise turns one client into a retry loop.

Usage goes to a buffered channel and is written in batches. Record never
blocks: if the buffer fills — a database stalled long enough to lose eight
thousand records — records are dropped and counted, because losing accounting
is bad and turning a stalled database into an outage is worse. A clean
shutdown drains what is queued; those are requests that were already served.
Failures are logged at error, since silently losing billing data is the kind
of thing discovered at the end of a month.

Requests are accounted whatever their outcome. A request that used a GPU and
then failed still used it, and a key that could spend nothing by failing would
be a way around a budget.

Budgets are counted in memory and restored from the store at startup, so a
restart does not hand every key a fresh allowance. The budget period travels
with the reservation rather than being re-derived, because a monthly budget
credited to a daily counter is a budget that never binds.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-09-16 08:52:36 +02:00

102 lines
2.9 KiB
Go

package accounting
import (
"sync"
"time"
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)
// PeriodStart is when the current budget period began.
//
// Periods are anchored to UTC rather than to a per-key start date, so every
// key's day and month boundary is the same instant. A key whose period started
// when it was created would make "spend this month" mean something different
// for every key, and make the restore query unanswerable in one pass.
func PeriodStart(p store.Period, now time.Time) time.Time {
u := now.UTC()
switch p {
case store.PeriodDay:
return time.Date(u.Year(), u.Month(), u.Day(), 0, 0, 0, 0, time.UTC)
case store.PeriodMonth:
return time.Date(u.Year(), u.Month(), 1, 0, 0, 0, 0, time.UTC)
default:
return time.Time{}
}
}
// Budgets tracks spend per key in memory, so a budget check costs no query.
//
// The database catches up through the batch writer. That means the in-memory
// figure leads the stored one, which is the right direction: a budget that only
// counted what had already been written would let a burst overspend by however
// much the writer was behind.
type Budgets struct {
now func() time.Time
mu sync.Mutex
spend map[string]counter
}
type counter struct {
micros int64
since time.Time
}
// NewBudgets builds an empty tracker.
func NewBudgets() *Budgets {
return &Budgets{now: time.Now, spend: map[string]counter{}}
}
// Restore seeds the counters from what the store already recorded.
//
// A restart must not reset a budget, which it would if the counter lived only
// in memory. The caller passes the spend since each period's start.
func (b *Budgets) Restore(period store.Period, spend map[string]int64) {
start := PeriodStart(period, b.now())
b.mu.Lock()
defer b.mu.Unlock()
for key, micros := range spend {
b.spend[key] = counter{micros: micros, since: start}
}
}
// Spent returns a key's spend in the current period, resetting it if the period
// has rolled over.
func (b *Budgets) Spent(keyID string, period store.Period) int64 {
start := PeriodStart(period, b.now())
b.mu.Lock()
defer b.mu.Unlock()
return b.currentLocked(keyID, start).micros
}
// Add records spend.
func (b *Budgets) Add(keyID string, period store.Period, micros int64) {
if micros == 0 {
return
}
start := PeriodStart(period, b.now())
b.mu.Lock()
defer b.mu.Unlock()
c := b.currentLocked(keyID, start)
c.micros += micros
b.spend[keyID] = c
}
// Exhausted reports whether a key has spent its budget.
func (b *Budgets) Exhausted(keyID string, period store.Period, budget int64) bool {
return b.Spent(keyID, period) >= budget
}
// currentLocked returns the counter, rolled over if the period changed.
func (b *Budgets) currentLocked(keyID string, start time.Time) counter {
c, ok := b.spend[keyID]
if !ok || c.since.Before(start) {
// Either never seen, or the period rolled over while the process was
// running. Both start from zero.
c = counter{since: start}
b.spend[keyID] = c
}
return c
}