quota: z.ai credit-bucket back-pressure + resource gate + usage accounting (Redmine 490+491)
The loop now consults quota and host state before every dispatch and
DEFERS gated work with a logged reason instead of letting turns die at
the provider (the 2026-08-28 19:00 quota-wall failure mode, replayed as
a test). Adds internal/quota: 5h/weekly credit buckets (provider poll
when z.ai ships an endpoint - fake-server tested - else locally
estimated from the documented credit formula), TZ-aware peak window
(default 01:00-05:00 America/Chicago weekdays, matching the documented
z.ai peak Mon-Fri 14:00-18:00 Singapore), block/defer thresholds, a
read-only load/mem/disk/IO-PSI monitor, an optional redis shared-state
hop (stdlib RESP2 mini-client) so all instances of an account
coordinate, per-class token+credit accounting in loop.jsonl, and
`harness quota status|probe|gate`. Config: [quota] + [resources]
sections; README runbook covers the redis container and deploy-time
cgroup enforcement.
💘 Generated with Crush
Assisted-by: Crush:glm-5.2
This commit is contained in:
@@ -0,0 +1,288 @@
|
||||
package quota
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"ukrrs.com/mopac/harness/internal/config"
|
||||
)
|
||||
|
||||
// Decision is the back-pressure verdict for one candidate dispatch.
|
||||
type Decision struct {
|
||||
Action string // "allow" | "defer"
|
||||
Reason string // human/log line; empty on allow unless Peak is set
|
||||
Peak bool // decision was made inside the peak window (context)
|
||||
}
|
||||
|
||||
// Action constants.
|
||||
const (
|
||||
ActionAllow = "allow"
|
||||
ActionDefer = "defer"
|
||||
)
|
||||
|
||||
func allow(peak bool) Decision { return Decision{Action: ActionAllow, Peak: peak} }
|
||||
|
||||
func defer_(reason string, peak bool) Decision {
|
||||
return Decision{Action: ActionDefer, Reason: reason, Peak: peak}
|
||||
}
|
||||
|
||||
// Gate is the loop's back-pressure consultant: quota snapshot (polled or
|
||||
// estimated), peak schedule, thresholds. It never blocks startup on redis
|
||||
// or the provider being down; it degrades to local estimates and allows.
|
||||
type Gate struct {
|
||||
cfg config.QuotaConfig
|
||||
sched Schedule
|
||||
shared *SharedState
|
||||
keys *config.KeyResolver
|
||||
http *http.Client
|
||||
now func() time.Time
|
||||
|
||||
mu sync.Mutex
|
||||
lastSnap *QuotaSnapshot // last successful poll (cache across failures)
|
||||
local5h float64 // process-local estimate fallbacks (redis down)
|
||||
localWeek float64
|
||||
win5hID string
|
||||
winWeekID string
|
||||
}
|
||||
|
||||
// NewGate builds the gate from the [quota] config. shared may be nil
|
||||
// (local-only estimates); keys may be nil when usage_url is unset.
|
||||
func NewGate(cfg config.QuotaConfig, keys *config.KeyResolver, shared *SharedState) (*Gate, error) {
|
||||
sched, err := NewSchedule(cfg.PeakStart, cfg.PeakEnd, cfg.Timezone, cfg.PeakWeekdaysOnly)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("quota schedule: %w", err)
|
||||
}
|
||||
return &Gate{
|
||||
cfg: cfg,
|
||||
sched: sched,
|
||||
shared: shared,
|
||||
keys: keys,
|
||||
http: &http.Client{Timeout: 15 * time.Second},
|
||||
now: time.Now,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// SetClock overrides the gate's clock (tests inject a fake).
|
||||
func (g *Gate) SetClock(now func() time.Time) { g.now = now }
|
||||
|
||||
// InPeak exposes the schedule verdict at the gate's clock.
|
||||
func (g *Gate) InPeak() bool { return g.sched.InPeak(g.now()) }
|
||||
|
||||
// Schedule returns the parsed peak window (status output).
|
||||
func (g *Gate) Schedule() Schedule { return g.sched }
|
||||
|
||||
// Snapshot returns the freshest quota state: a live poll when the provider
|
||||
// endpoint is configured and the cache is stale, else the redis-published
|
||||
// snapshot from any instance, else the local cache, else the estimate
|
||||
// synthesized from this instance's recorded consumption. The returned
|
||||
// snapshot is never nil when the gate is enabled.
|
||||
func (g *Gate) Snapshot(ctx context.Context) *QuotaSnapshot {
|
||||
now := g.now()
|
||||
if g.cfg.UsageURL != "" {
|
||||
if s := g.cachedPoll(); s != nil && now.Sub(s.FetchedAt) < time.Duration(g.cfg.PollIntervalSecs)*time.Second {
|
||||
return s
|
||||
}
|
||||
if s, err := g.poll(ctx, now); err == nil {
|
||||
return s
|
||||
}
|
||||
}
|
||||
if g.shared != nil {
|
||||
if s, err := g.shared.LoadSnapshot(g.cfg.Account); err == nil && s != nil &&
|
||||
now.Sub(s.FetchedAt) < 2*time.Duration(g.cfg.PollIntervalSecs)*time.Second {
|
||||
return s
|
||||
}
|
||||
}
|
||||
if s := g.cachedPoll(); s != nil {
|
||||
return s
|
||||
}
|
||||
return g.estimateSnapshot(now)
|
||||
}
|
||||
|
||||
func (g *Gate) cachedPoll() *QuotaSnapshot {
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
return g.lastSnap
|
||||
}
|
||||
|
||||
func (g *Gate) poll(ctx context.Context, now time.Time) (*QuotaSnapshot, error) {
|
||||
key, err := g.keys.Resolve(ctx, g.cfg.KeyRef)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("quota key: %w", err)
|
||||
}
|
||||
body, err := FetchUsage(ctx, g.http, g.cfg.UsageURL, key)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
snap, err := ParseUsage(g.cfg.Account, now, body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
g.mu.Lock()
|
||||
g.lastSnap = snap
|
||||
g.mu.Unlock()
|
||||
_ = g.shared.PublishSnapshot(snap, time.Duration(g.cfg.PollIntervalSecs)*3*time.Second)
|
||||
return snap, nil
|
||||
}
|
||||
|
||||
// window ids: the 5h bucket advances in fixed 5h steps from the epoch; the
|
||||
// weekly bucket is the Monday 00:00 of the schedule's TZ (z.ai weekly
|
||||
// credits reset 7 days after activation — an approximation until the usage
|
||||
// endpoint ships reset_at, noted in the REPORT).
|
||||
func (g *Gate) windowIDs(now time.Time) (w5h, week string) {
|
||||
five := 5 * time.Hour
|
||||
w5h = fmt.Sprintf("%d", now.Unix()/int64(five.Seconds()))
|
||||
mon := now.In(g.sched.Loc)
|
||||
for mon.Weekday() != time.Monday {
|
||||
mon = mon.AddDate(0, 0, -1)
|
||||
}
|
||||
y, m, d := mon.Date()
|
||||
week = fmt.Sprintf("%04d%02d%02d", y, int(m), d)
|
||||
return w5h, week
|
||||
}
|
||||
|
||||
// RecordTurn records one completed turn's estimated credits into the shared
|
||||
// state (redis when configured; process-local maps always, so a down redis
|
||||
// never loses this instance's own accounting).
|
||||
func (g *Gate) RecordTurn(in EstimateTurnInput) float64 {
|
||||
credits := EstimateCredits(in.Model, in.PromptTokens, in.CachedTokens, in.CompletionTokens, in.Peak)
|
||||
now := g.now()
|
||||
w5h, week := g.windowIDs(now)
|
||||
|
||||
g.mu.Lock()
|
||||
if g.win5hID != w5h { // 5h window rolled: reset the local counter
|
||||
g.local5h, g.win5hID = 0, w5h
|
||||
}
|
||||
if g.winWeekID != week {
|
||||
g.localWeek, g.winWeekID = 0, week
|
||||
}
|
||||
g.local5h += credits
|
||||
g.localWeek += credits
|
||||
g.mu.Unlock()
|
||||
|
||||
if g.shared != nil {
|
||||
k5 := fmt.Sprintf("mopac:quota:%s:est:5h:%s", g.cfg.Account, w5h)
|
||||
kw := fmt.Sprintf("mopac:quota:%s:est:weekly:%s", g.cfg.Account, week)
|
||||
if _, err := g.shared.incrByFloat(k5, credits); err == nil {
|
||||
_, _ = g.shared.command("EXPIRE", k5, "21600") // 5h + 1h slack
|
||||
}
|
||||
_, _ = g.shared.incrByFloat(kw, credits)
|
||||
}
|
||||
return credits
|
||||
}
|
||||
|
||||
// estimateSnapshot synthesizes quota state from recorded consumption against
|
||||
// the configured plan limits (the no-endpoint mode).
|
||||
func (g *Gate) estimateSnapshot(now time.Time) *QuotaSnapshot {
|
||||
w5h, week := g.windowIDs(now)
|
||||
var used5, usedW float64
|
||||
if g.shared != nil {
|
||||
if v, err := g.shared.get(fmt.Sprintf("mopac:quota:%s:est:5h:%s", g.cfg.Account, w5h)); err == nil && v != "" {
|
||||
fmt.Sscanf(v, "%g", &used5)
|
||||
}
|
||||
if v, err := g.shared.get(fmt.Sprintf("mopac:quota:%s:est:weekly:%s", g.cfg.Account, week)); err == nil && v != "" {
|
||||
fmt.Sscanf(v, "%g", &usedW)
|
||||
}
|
||||
}
|
||||
if used5 == 0 || usedW == 0 {
|
||||
g.mu.Lock()
|
||||
if g.win5hID == w5h {
|
||||
used5 = max(used5, g.local5h)
|
||||
}
|
||||
if g.winWeekID == week {
|
||||
usedW = max(usedW, g.localWeek)
|
||||
}
|
||||
g.mu.Unlock()
|
||||
}
|
||||
src := "estimate"
|
||||
if g.shared != nil {
|
||||
src = "estimate+redis"
|
||||
}
|
||||
return &QuotaSnapshot{
|
||||
Account: g.cfg.Account,
|
||||
Source: src,
|
||||
FetchedAt: now,
|
||||
Buckets: []Bucket{
|
||||
{ID: Bucket5h, Used: used5, Limit: g.cfg.Plan5hCredits,
|
||||
WindowReset: now.Add(5 * time.Hour)},
|
||||
{ID: BucketWeekly, Used: usedW, Limit: g.cfg.PlanWeeklyCredits,
|
||||
WindowReset: nextWeekStart(now, g.sched.Loc)},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func nextWeekStart(now time.Time, loc *time.Location) time.Time {
|
||||
t := now.In(loc)
|
||||
for t.Weekday() != time.Monday {
|
||||
t = t.AddDate(0, 0, 1)
|
||||
}
|
||||
y, m, d := t.Date()
|
||||
return time.Date(y, m, d, 0, 0, 0, 0, loc).UTC()
|
||||
}
|
||||
|
||||
// Decide is the pre-dispatch consultation. Priority: hard quota wall first
|
||||
// (the 19:00 failure mode), then peak-window class restriction, then
|
||||
// soft-quota heavy-class deferral. Unknown quota state never defers.
|
||||
func (g *Gate) Decide(ctx context.Context, class string) Decision {
|
||||
peak := g.InPeak()
|
||||
snap := g.Snapshot(ctx)
|
||||
ratio := snap.MaxUsedPct()
|
||||
|
||||
if ratio >= g.cfg.BlockAtPct {
|
||||
return defer_(fmt.Sprintf("quota: %s bucket at %.0f%% (>= block %.0f%%): all classes deferred until reset (source %s)",
|
||||
worstBucketID(snap), ratio, g.cfg.BlockAtPct, snap.Source), peak)
|
||||
}
|
||||
if peak && !g.peakClass(class) {
|
||||
return defer_(fmt.Sprintf("peak window %s-%s %s: class %q deferred to off-peak (flash-tier classes only: %s)",
|
||||
g.cfg.PeakStart, g.cfg.PeakEnd, g.cfg.Timezone, class, strings.Join(g.cfg.PeakClasses, ",")), peak)
|
||||
}
|
||||
if ratio >= g.cfg.DeferAtPct && !g.peakClass(class) {
|
||||
return defer_(fmt.Sprintf("quota: %s bucket at %.0f%% (>= defer %.0f%%): heavy class %q deferred; LLM-lite continues (source %s)",
|
||||
worstBucketID(snap), ratio, g.cfg.DeferAtPct, class, snap.Source), peak)
|
||||
}
|
||||
return allow(peak)
|
||||
}
|
||||
|
||||
func (g *Gate) peakClass(class string) bool {
|
||||
for _, c := range g.cfg.PeakClasses {
|
||||
if c == class {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func worstBucketID(s *QuotaSnapshot) string {
|
||||
id, max := "", 0.0
|
||||
for _, b := range s.Buckets {
|
||||
if p := b.UsedPct(); p > max {
|
||||
max, id = p, b.ID
|
||||
}
|
||||
}
|
||||
return id
|
||||
}
|
||||
|
||||
// StatusLine renders the one-line quota status for loop startup / `harness
|
||||
// quota status`.
|
||||
func (g *Gate) StatusLine(snap *QuotaSnapshot) string {
|
||||
parts := make([]string, 0, len(snap.Buckets)+2)
|
||||
parts = append(parts, fmt.Sprintf("account=%s source=%s", snap.Account, snap.Source))
|
||||
for _, b := range snap.Buckets {
|
||||
parts = append(parts, fmt.Sprintf("%s=%.0f/%.0f(%.0f%%)", b.ID, b.Used, b.Limit, b.UsedPct()))
|
||||
}
|
||||
if g.InPeak() {
|
||||
parts = append(parts, "PEAK")
|
||||
}
|
||||
return strings.Join(parts, " ")
|
||||
}
|
||||
|
||||
// SortedClasses returns PeakClasses sorted (stable status output).
|
||||
func (g *Gate) SortedClasses() []string {
|
||||
out := append([]string{}, g.cfg.PeakClasses...)
|
||||
sort.Strings(out)
|
||||
return out
|
||||
}
|
||||
Reference in New Issue
Block a user