Ilker ManapandClaude Opus 5.5
8b7eefef16
CI / test (push) Successful in 9m35s
A standby follows the leader's Ready replicas over the agent port, every two seconds, and reports back the replicas its router used and the on-demand models it was asked for. No engine keys travel: every control plane derives them from the cluster secret. Budgets are raised to the store's every 30 s. Unheard from the leader for a minute, a standby passes /v1 on as before. Off by default. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
115 lines
3.3 KiB
Go
115 lines
3.3 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}
|
|
}
|
|
}
|
|
|
|
// Raise lifts counters to the spend given, where it is more.
|
|
func (b *Budgets) Raise(period store.Period, spend map[string]int64) {
|
|
start := PeriodStart(period, b.now())
|
|
b.mu.Lock()
|
|
defer b.mu.Unlock()
|
|
for key, micros := range spend {
|
|
if c := b.currentLocked(key, start); micros > c.micros {
|
|
c.micros = micros
|
|
b.spend[key] = c
|
|
}
|
|
}
|
|
}
|
|
|
|
// 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
|
|
}
|