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>
102 lines
2.9 KiB
Go
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
|
|
}
|