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,48 @@
|
||||
package quota
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// FetchUsage GETs the provider usage endpoint with bearer auth. The key is
|
||||
// resolved per call from the [quota] key_ref and never logged; error paths
|
||||
// carry only status codes and redacted/truncated bodies (same discipline as
|
||||
// the keyproxy hop).
|
||||
func FetchUsage(ctx context.Context, client *http.Client, usageURL, key string) ([]byte, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, usageURL, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("usage request: %w", err)
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
req.Header.Set("Accept", "application/json")
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("usage poll: %w", redactQuery(err))
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return nil, fmt.Errorf("usage poll: HTTP %d: %s", resp.StatusCode, truncateBody(string(body)))
|
||||
}
|
||||
return body, nil
|
||||
}
|
||||
|
||||
func redactQuery(err error) error {
|
||||
msg := err.Error()
|
||||
if i := strings.Index(msg, "?"); i >= 0 {
|
||||
msg = msg[:i] + "?..."
|
||||
}
|
||||
return fmt.Errorf("%s", msg)
|
||||
}
|
||||
|
||||
func truncateBody(s string) string {
|
||||
s = strings.TrimSpace(s)
|
||||
if len(s) > 200 {
|
||||
s = s[:200] + "..."
|
||||
}
|
||||
return s
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -0,0 +1,74 @@
|
||||
package quota
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Multipliers are the z.ai credit formula coefficients (per 10k tokens,
|
||||
// docs.z.ai/devpack/overview). Family selects by concrete model name; the
|
||||
// flash tier (GLM-5.3-Flash, routed from glm-4.7*) is the cheap one.
|
||||
type Multipliers struct {
|
||||
Input float64
|
||||
CachedInput float64
|
||||
Output float64
|
||||
}
|
||||
|
||||
// Flagship multipliers: GLM-5.3 (and everything auto-routed to it).
|
||||
func FlagshipMultipliers() Multipliers { return Multipliers{Input: 6.9, CachedInput: 1.7, Output: 24} }
|
||||
|
||||
// Flash multipliers: GLM-5.3-Flash (and glm-4.7* routing).
|
||||
func FlashMultipliers() Multipliers { return Multipliers{Input: 2.3, CachedInput: 0.56, Output: 8} }
|
||||
|
||||
// MultipliersFor maps a concrete proxy model name to its credit multiplier
|
||||
// family. Unknown models read as flagship (conservative: overestimate cost
|
||||
// rather than silently burn quota).
|
||||
func MultipliersFor(model string) Multipliers {
|
||||
m := strings.ToLower(model)
|
||||
switch {
|
||||
case strings.Contains(m, "flash"):
|
||||
return FlashMultipliers()
|
||||
default:
|
||||
return FlagshipMultipliers()
|
||||
}
|
||||
}
|
||||
|
||||
// EstimateCredits computes the z.ai credits one turn consumed from its token
|
||||
// usage: (input*in + cached*cache + output*out) / 10000, halved when the
|
||||
// turn ran off-peak (z.ai charges 50% outside peak hours).
|
||||
func EstimateCredits(model string, promptTokens, cachedTokens, completionTokens int, peak bool) float64 {
|
||||
mult := MultipliersFor(model)
|
||||
credits := (float64(promptTokens)*mult.Input +
|
||||
float64(cachedTokens)*mult.CachedInput +
|
||||
float64(completionTokens)*mult.Output) / 10000
|
||||
if !peak {
|
||||
credits /= 2
|
||||
}
|
||||
// Round to 6 decimals: keeps redis INCRBYFLOAT values readable and the
|
||||
// JSONL compact; sub-microcredit noise is meaningless.
|
||||
return math.Round(credits*1e6) / 1e6
|
||||
}
|
||||
|
||||
// EstimateTurnInput is the per-turn usage record the loop hands the gate.
|
||||
type EstimateTurnInput struct {
|
||||
Model string
|
||||
PromptTokens int
|
||||
CachedTokens int
|
||||
CompletionTokens int
|
||||
Peak bool // turn ran inside the peak window
|
||||
}
|
||||
|
||||
// Describe renders a human summary of one turn's credit cost (logs, REPORTs).
|
||||
func (e EstimateTurnInput) Describe() string {
|
||||
return fmt.Sprintf("%s: %d/%d/%d tokens (in/cached/out) %s = %.4f credits",
|
||||
e.Model, e.PromptTokens, e.CachedTokens, e.CompletionTokens,
|
||||
peakTag(e.Peak), EstimateCredits(e.Model, e.PromptTokens, e.CachedTokens, e.CompletionTokens, e.Peak))
|
||||
}
|
||||
|
||||
func peakTag(peak bool) string {
|
||||
if peak {
|
||||
return "peak"
|
||||
}
|
||||
return "off-peak"
|
||||
}
|
||||
@@ -0,0 +1,143 @@
|
||||
package quota
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"syscall"
|
||||
)
|
||||
|
||||
// ResourceStats is one read-only sample of host load: the numbers the
|
||||
// resource gate compares against its thresholds. Zero fields mean "could
|
||||
// not read" and never trigger a busy verdict on their own.
|
||||
type ResourceStats struct {
|
||||
LoadAvg1m float64
|
||||
MemAvailable float64 // MB
|
||||
DiskFree float64 // MB, at path
|
||||
IODelayPct float64 // /proc/pressure/io "some avg60" percent; 0 if PSI absent
|
||||
HasPSI bool
|
||||
}
|
||||
|
||||
// ReadResourceStats samples /proc/loadavg, /proc/meminfo, /proc/pressure/io
|
||||
// and statfs(path) — all read-only, no host packages. procRoot/sysRoot are
|
||||
// seams (tests point them at fixture trees; production passes /proc, /sys).
|
||||
func ReadResourceStats(procRoot, sysRoot, path string) ResourceStats {
|
||||
var st ResourceStats
|
||||
if la, ok := readLoadAvg(filepath.Join(procRoot, "loadavg")); ok {
|
||||
st.LoadAvg1m = la
|
||||
}
|
||||
if mb, ok := readMemAvailable(filepath.Join(procRoot, "meminfo")); ok {
|
||||
st.MemAvailable = mb
|
||||
}
|
||||
// PSI lives under /proc/pressure (io); older kernels expose none — the
|
||||
// gate then skips the IO check rather than failing.
|
||||
if p, ok := readIODelay(filepath.Join(procRoot, "pressure", "io")); ok {
|
||||
st.IODelayPct, st.HasPSI = p, true
|
||||
}
|
||||
_ = sysRoot // reserved: /sys/class/... sources when PSI is absent
|
||||
if df, ok := readDiskFree(path); ok {
|
||||
st.DiskFree = df
|
||||
}
|
||||
return st
|
||||
}
|
||||
|
||||
func readLoadAvg(path string) (float64, bool) {
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
fields := strings.Fields(string(data))
|
||||
if len(fields) < 1 {
|
||||
return 0, false
|
||||
}
|
||||
f, err := strconv.ParseFloat(fields[0], 64)
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
return f, true
|
||||
}
|
||||
|
||||
func readMemAvailable(path string) (float64, bool) {
|
||||
f, err := os.Open(path)
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
defer f.Close()
|
||||
sc := bufio.NewScanner(f)
|
||||
for sc.Scan() {
|
||||
line := sc.Text()
|
||||
if strings.HasPrefix(line, "MemAvailable:") {
|
||||
fields := strings.Fields(line)
|
||||
if len(fields) < 2 {
|
||||
return 0, false
|
||||
}
|
||||
kb, err := strconv.ParseFloat(fields[1], 64)
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
return kb / 1024, true
|
||||
}
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
|
||||
// readIODelay parses /proc/pressure/io, e.g.
|
||||
// "some avg10=0.00 avg60=0.12 avg300=0.05 total=123456789" — avg60 is the
|
||||
// steady-state signal (a build spiking IO shows here within a minute).
|
||||
func readIODelay(path string) (float64, bool) {
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
for _, line := range strings.Split(string(data), "\n") {
|
||||
if !strings.HasPrefix(line, "some ") {
|
||||
continue
|
||||
}
|
||||
for _, field := range strings.Fields(line)[1:] {
|
||||
if strings.HasPrefix(field, "avg60=") {
|
||||
if v, err := strconv.ParseFloat(strings.TrimPrefix(field, "avg60="), 64); err == nil {
|
||||
return v, true
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
|
||||
func readDiskFree(path string) (float64, bool) {
|
||||
var fs syscall.Statfs_t
|
||||
if err := syscall.Statfs(path, &fs); err != nil {
|
||||
return 0, false
|
||||
}
|
||||
return float64(fs.Bavail) * float64(fs.Bsize) / (1024 * 1024), true
|
||||
}
|
||||
|
||||
// BusyCheck compares a sample against the resource thresholds. Violations
|
||||
// are collected (all reported, not just the first) — the loop defers while
|
||||
// any threshold trips, with every reason surfaced in the defer log line.
|
||||
type BusyCheck struct {
|
||||
MaxLoadAvg float64
|
||||
MinMemAvailableMB float64
|
||||
MinDiskFreeMB float64
|
||||
MaxIODelayPct float64
|
||||
}
|
||||
|
||||
func (c BusyCheck) Evaluate(st ResourceStats) []string {
|
||||
var reasons []string
|
||||
if c.MaxLoadAvg > 0 && st.LoadAvg1m > c.MaxLoadAvg {
|
||||
reasons = append(reasons, fmt.Sprintf("load %.2f > %.2f", st.LoadAvg1m, c.MaxLoadAvg))
|
||||
}
|
||||
if c.MinMemAvailableMB > 0 && st.MemAvailable > 0 && st.MemAvailable < c.MinMemAvailableMB {
|
||||
reasons = append(reasons, fmt.Sprintf("mem available %.0fMB < %.0fMB", st.MemAvailable, c.MinMemAvailableMB))
|
||||
}
|
||||
if c.MinDiskFreeMB > 0 && st.DiskFree > 0 && st.DiskFree < c.MinDiskFreeMB {
|
||||
reasons = append(reasons, fmt.Sprintf("disk free %.0fMB < %.0fMB", st.DiskFree, c.MinDiskFreeMB))
|
||||
}
|
||||
if c.MaxIODelayPct > 0 && st.HasPSI && st.IODelayPct > c.MaxIODelayPct {
|
||||
reasons = append(reasons, fmt.Sprintf("io delay %.1f%% > %.1f%%", st.IODelayPct, c.MaxIODelayPct))
|
||||
}
|
||||
return reasons
|
||||
}
|
||||
@@ -0,0 +1,276 @@
|
||||
package quota
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"fmt"
|
||||
"net"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestResourceStatsFromFixture(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
proc := filepath.Join(root, "proc")
|
||||
os.MkdirAll(filepath.Join(proc, "pressure"), 0o755)
|
||||
os.WriteFile(filepath.Join(proc, "loadavg"), []byte("7.32 5.10 2.20 3/900 12345\n"), 0o644)
|
||||
os.WriteFile(filepath.Join(proc, "meminfo"), []byte("MemTotal: 16000000 kB\nMemAvailable: 3145728 kB\n"), 0o644)
|
||||
os.WriteFile(filepath.Join(proc, "pressure", "io"), []byte("some avg10=1.00 avg60=2.50 avg300=0.10 total=123456789\nfull avg10=0.00 avg60=0.00 avg300=0.00 total=0\n"), 0o644)
|
||||
work := filepath.Join(root, "work")
|
||||
os.MkdirAll(work, 0o755)
|
||||
|
||||
st := ReadResourceStats(proc, filepath.Join(root, "sys"), work)
|
||||
if st.LoadAvg1m != 7.32 {
|
||||
t.Errorf("load = %v, want 7.32", st.LoadAvg1m)
|
||||
}
|
||||
if st.MemAvailable != 3072 {
|
||||
t.Errorf("mem = %v MB, want 3072", st.MemAvailable)
|
||||
}
|
||||
if !st.HasPSI || st.IODelayPct != 2.5 {
|
||||
t.Errorf("io = %v (psi=%v), want 2.5", st.IODelayPct, st.HasPSI)
|
||||
}
|
||||
if st.DiskFree <= 0 {
|
||||
t.Errorf("disk free = %v, want > 0 (statfs on tempdir)", st.DiskFree)
|
||||
}
|
||||
|
||||
check := BusyCheck{MaxLoadAvg: 6, MinMemAvailableMB: 4096, MinDiskFreeMB: 1, MaxIODelayPct: 90}
|
||||
reasons := check.Evaluate(st)
|
||||
joined := strings.Join(reasons, ";")
|
||||
for _, want := range []string{"load 7.32 > 6", "mem available 3072MB < 4096MB"} {
|
||||
if !strings.Contains(joined, want) {
|
||||
t.Errorf("reasons %q missing %q", joined, want)
|
||||
}
|
||||
}
|
||||
// IO under threshold: must NOT appear.
|
||||
if strings.Contains(joined, "io delay") {
|
||||
t.Errorf("io below threshold must not trip: %q", joined)
|
||||
}
|
||||
|
||||
// Missing PSI file: the IO check is skipped, not an error.
|
||||
os.Remove(filepath.Join(proc, "pressure", "io"))
|
||||
st2 := ReadResourceStats(proc, filepath.Join(root, "sys"), work)
|
||||
if st2.HasPSI {
|
||||
t.Error("missing PSI must read HasPSI=false")
|
||||
}
|
||||
if got := (BusyCheck{MaxIODelayPct: 1}).Evaluate(st2); len(got) != 0 {
|
||||
t.Errorf("no PSI must never trip IO: %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestResourceStatsNeverBusyOnUnreadable mirrors the deploy invariant: read
|
||||
// failures defer nothing (the gate is advisory, /proc layout changes must
|
||||
// not stop work).
|
||||
func TestResourceStatsNeverBusyOnUnreadable(t *testing.T) {
|
||||
root := t.TempDir() // empty: no loadavg, no meminfo
|
||||
st := ReadResourceStats(filepath.Join(root, "proc"), filepath.Join(root, "sys"), root)
|
||||
if got := (BusyCheck{MaxLoadAvg: 1, MinMemAvailableMB: 1, MinDiskFreeMB: 1, MaxIODelayPct: 1}).Evaluate(st); len(got) > 1 {
|
||||
t.Errorf("unreadable sources must not trip load/mem (disk may read): %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// fakeRedis is a minimal RESP2 server: enough command surface for the
|
||||
// SharedState client (SET/GET/INCRBYFLOAT/EXPIRE/SELECT), backed by a map.
|
||||
type fakeRedis struct {
|
||||
mu sync.Mutex
|
||||
data map[string]string
|
||||
ttl map[string]int64
|
||||
srv net.Listener
|
||||
}
|
||||
|
||||
func newFakeRedis(t *testing.T) *fakeRedis {
|
||||
t.Helper()
|
||||
f := &fakeRedis{data: map[string]string{}, ttl: map[string]int64{}}
|
||||
srv, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
f.srv = srv
|
||||
go f.serve()
|
||||
t.Cleanup(func() { srv.Close() })
|
||||
return f
|
||||
}
|
||||
|
||||
func (f *fakeRedis) url(t *testing.T) string {
|
||||
t.Helper()
|
||||
return "redis://" + f.srv.Addr().String() + "/0"
|
||||
}
|
||||
|
||||
func (f *fakeRedis) serve() {
|
||||
for {
|
||||
conn, err := f.srv.Accept()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
go f.handle(conn)
|
||||
}
|
||||
}
|
||||
|
||||
func (f *fakeRedis) handle(conn net.Conn) {
|
||||
defer conn.Close()
|
||||
r := bufio.NewReader(conn)
|
||||
w := bufio.NewWriter(conn)
|
||||
for {
|
||||
args, err := readCommand(r)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if len(args) == 0 {
|
||||
continue
|
||||
}
|
||||
f.mu.Lock()
|
||||
reply := f.exec(args)
|
||||
f.mu.Unlock()
|
||||
w.WriteString(reply + "\r\n")
|
||||
if w.Flush() != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func readCommand(r *bufio.Reader) ([]string, error) {
|
||||
line, err := r.ReadString('\n')
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
line = strings.TrimRight(line, "\r\n")
|
||||
if !strings.HasPrefix(line, "*") {
|
||||
return strings.Fields(line), nil
|
||||
}
|
||||
n, _ := strconv.Atoi(line[1:])
|
||||
args := make([]string, 0, n)
|
||||
for i := 0; i < n; i++ {
|
||||
hl, err := r.ReadString('\n')
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
hl = strings.TrimRight(hl, "\r\n")
|
||||
ln, _ := strconv.Atoi(strings.TrimPrefix(hl, "$"))
|
||||
buf := make([]byte, ln+2)
|
||||
if _, err := readFull(r, buf); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
args = append(args, string(buf[:ln]))
|
||||
}
|
||||
return args, nil
|
||||
}
|
||||
|
||||
func readFull(r *bufio.Reader, buf []byte) (int, error) {
|
||||
total := 0
|
||||
for total < len(buf) {
|
||||
n, err := r.Read(buf[total:])
|
||||
total += n
|
||||
if err != nil {
|
||||
return total, err
|
||||
}
|
||||
}
|
||||
return total, nil
|
||||
}
|
||||
|
||||
func (f *fakeRedis) exec(args []string) string {
|
||||
cmd := strings.ToUpper(args[0])
|
||||
switch cmd {
|
||||
case "PING":
|
||||
return "+PONG"
|
||||
case "SELECT":
|
||||
return "+OK"
|
||||
case "SET":
|
||||
f.data[args[1]] = args[2]
|
||||
if len(args) >= 5 && strings.ToUpper(args[3]) == "EX" {
|
||||
f.ttl[args[1]], _ = strconv.ParseInt(args[4], 10, 64)
|
||||
}
|
||||
return "+OK"
|
||||
case "GET":
|
||||
if v, ok := f.data[args[1]]; ok {
|
||||
return fmt.Sprintf("$%d\r\n%s", len(v), v)
|
||||
}
|
||||
return "$-1"
|
||||
case "EXPIRE":
|
||||
if _, ok := f.data[args[1]]; ok {
|
||||
f.ttl[args[1]], _ = strconv.ParseInt(args[2], 10, 64)
|
||||
return ":1"
|
||||
}
|
||||
return ":0"
|
||||
case "INCRBYFLOAT":
|
||||
cur := 0.0
|
||||
if v, ok := f.data[args[1]]; ok {
|
||||
cur, _ = strconv.ParseFloat(v, 64)
|
||||
}
|
||||
delta, err := strconv.ParseFloat(args[2], 64)
|
||||
if err != nil {
|
||||
return "-ERR bad delta"
|
||||
}
|
||||
cur += delta
|
||||
f.data[args[1]] = strconv.FormatFloat(cur, 'f', 6, 64)
|
||||
return "$" + strconv.Itoa(len(f.data[args[1]])) + "\r\n" + f.data[args[1]]
|
||||
default:
|
||||
return "-ERR unknown command '" + cmd + "'"
|
||||
}
|
||||
}
|
||||
|
||||
func TestSharedStateRoundTrip(t *testing.T) {
|
||||
fake := newFakeRedis(t)
|
||||
st, err := NewSharedState(fake.url(t))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer st.Close()
|
||||
|
||||
if err := st.setEx("k", "v1", time.Second); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if v, err := st.get("k"); err != nil || v != "v1" {
|
||||
t.Fatalf("get = %q err=%v", v, err)
|
||||
}
|
||||
if v, err := st.incrByFloat("ctr", 1.5); err != nil || v != 1.5 {
|
||||
t.Fatalf("incr = %v err=%v", v, err)
|
||||
}
|
||||
if v, err := st.incrByFloat("ctr", 2.25); err != nil || v != 3.75 {
|
||||
t.Fatalf("incr2 = %v err=%v (want 3.75)", v, err)
|
||||
}
|
||||
if v, err := st.get("missing"); err != nil || v != "" {
|
||||
t.Fatalf("missing key = %q err=%v (want empty, no error)", v, err)
|
||||
}
|
||||
|
||||
// snapshot publish/load through the typed helpers
|
||||
snap := &QuotaSnapshot{Account: "zai-9", Source: "usage_url", FetchedAt: time.Now(),
|
||||
Buckets: []Bucket{{ID: Bucket5h, Used: 3.5, Limit: 28000}}}
|
||||
if err := st.PublishSnapshot(snap, time.Minute); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, err := st.LoadSnapshot("zai-9")
|
||||
if err != nil || got == nil {
|
||||
t.Fatalf("load = %v err=%v", got, err)
|
||||
}
|
||||
if b, _ := got.Bucket(Bucket5h); b.Used != 3.5 {
|
||||
t.Errorf("bucket round trip: %+v", got.Buckets)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSharedStateDownNeverFatal(t *testing.T) {
|
||||
// Nothing listens on this port: every op must error softly (nil state
|
||||
// ops are no-ops; a dead server returns errors, never panics).
|
||||
dead, err := NewSharedState("redis://127.0.0.1:1/0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := dead.setEx("k", "v", time.Second); err == nil {
|
||||
t.Error("dead redis SET must error")
|
||||
}
|
||||
if _, err := dead.incrByFloat("k", 1); err == nil {
|
||||
t.Error("dead redis INCR must error")
|
||||
}
|
||||
var nilState *SharedState
|
||||
if err := nilState.setEx("k", "v", time.Second); err != nil {
|
||||
t.Errorf("nil state must no-op: %v", err)
|
||||
}
|
||||
if st, err := NewSharedState(""); err != nil || st != nil {
|
||||
t.Errorf("empty url must return nil state, got %v %v", st, err)
|
||||
}
|
||||
if _, err := NewSharedState("redis://"); err == nil {
|
||||
t.Error("hostless url must error")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,77 @@
|
||||
package quota
|
||||
|
||||
import (
|
||||
_ "time/tzdata" // embedded zone database: TZ-aware windows work in scratch containers
|
||||
|
||||
"time"
|
||||
|
||||
"ukrrs.com/mopac/harness/internal/config"
|
||||
)
|
||||
|
||||
// Schedule is the TZ-aware peak window. z.ai peak hours are documented as
|
||||
// Monday-Friday 14:00-18:00 Singapore (UTC+8), which is 01:00-05:00
|
||||
// America/Chicago in winter (CST) / 00:00-04:00 during US DST — the default
|
||||
// matches Charles's "0100 to 0500 CST" and is configurable to the minute.
|
||||
// The window may wrap midnight (start > end): then it covers evenings of
|
||||
// the start day plus early mornings of the following day.
|
||||
type Schedule struct {
|
||||
Start, End time.Duration // minutes-since-midnight in Loc
|
||||
Loc *time.Location
|
||||
WeekdaysOnly bool
|
||||
}
|
||||
|
||||
// NewSchedule parses the window out of the [quota] config fields.
|
||||
func NewSchedule(peakStart, peakEnd, tzName string, weekdaysOnly bool) (Schedule, error) {
|
||||
start, err := config.ParseHHMM(peakStart)
|
||||
if err != nil {
|
||||
return Schedule{}, err
|
||||
}
|
||||
end, err := config.ParseHHMM(peakEnd)
|
||||
if err != nil {
|
||||
return Schedule{}, err
|
||||
}
|
||||
loc, err := time.LoadLocation(tzName)
|
||||
if err != nil {
|
||||
return Schedule{}, err
|
||||
}
|
||||
return Schedule{Start: start, End: end, Loc: loc, WeekdaysOnly: weekdaysOnly}, nil
|
||||
}
|
||||
|
||||
// SameDay reports whether t is inside the window without crossing midnight
|
||||
// (start <= t < end, single-day window).
|
||||
func (s Schedule) inWindow(mins time.Duration) bool {
|
||||
if s.Start <= s.End {
|
||||
return mins >= s.Start && mins < s.End
|
||||
}
|
||||
// Wrapped window: [start, 24h) of the start day, [0, end) of the next.
|
||||
return mins >= s.Start || mins < s.End
|
||||
}
|
||||
|
||||
// InPeak reports whether t falls inside the peak window, evaluated in the
|
||||
// schedule's timezone. For wrapped windows the weekday check applies to the
|
||||
// day the window STARTED on (a Sunday 22:00-02:00 window is off all week
|
||||
// when weekdays-only, because it starts on Sunday).
|
||||
func (s Schedule) InPeak(t time.Time) bool {
|
||||
local := t.In(s.Loc)
|
||||
mins := time.Duration(local.Hour())*time.Hour + time.Duration(local.Minute())*time.Minute
|
||||
if s.Start <= s.End {
|
||||
if !s.inWindow(mins) {
|
||||
return false
|
||||
}
|
||||
return !s.WeekdaysOnly || isWeekday(local)
|
||||
}
|
||||
// Evening half: today carries the window.
|
||||
if mins >= s.Start {
|
||||
return !s.WeekdaysOnly || isWeekday(local)
|
||||
}
|
||||
// Morning half: the window started yesterday.
|
||||
if mins < s.End {
|
||||
return !s.WeekdaysOnly || isWeekday(local.AddDate(0, 0, -1))
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func isWeekday(t time.Time) bool {
|
||||
wd := t.Weekday()
|
||||
return wd >= time.Monday && wd <= time.Friday
|
||||
}
|
||||
@@ -0,0 +1,161 @@
|
||||
// Package quota is the z.ai coding-plan quota gate (Redmine 490) plus the
|
||||
// read-only system resource monitor (Redmine 491): it tracks plan credit
|
||||
// buckets (5-hour + weekly), computes per-turn credit consumption from the
|
||||
// documented z.ai formula, applies TZ-aware peak-hour scheduling and
|
||||
// back-pressure thresholds, and shares quota state across harness instances
|
||||
// through an optional redis container. The gate NEVER hard-fails the loop:
|
||||
// unknown state is permissive, exhaustion defers work with a logged reason.
|
||||
//
|
||||
// z.ai credit model (docs.z.ai/devpack/overview, 2026-08-29):
|
||||
//
|
||||
// credits = (input*in_mult + cached_input*cache_mult + output*out_mult) / 10000
|
||||
//
|
||||
// GLM-5.3: 6.9/1.7/24; GLM-5.3-Flash: 2.3/0.56/8. Off-peak hours charge 50%.
|
||||
// Plans: Lite 2,000/10,000, Pro 12,000/60,000, Max 28,000/140,000 credits
|
||||
// (5-hour / weekly). LIVE VERIFICATION of the usage endpoint is OPEN: z.ai
|
||||
// documents the buckets but no public usage REST route exists today
|
||||
// (probed 2026-08-29, see REPORT-20260829-0500-quota); the parser targets
|
||||
// the documented shape and runs against a fake server until z.ai ships it.
|
||||
package quota
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Bucket IDs for the z.ai coding-plan windows.
|
||||
const (
|
||||
Bucket5h = "5h"
|
||||
BucketWeekly = "weekly"
|
||||
)
|
||||
|
||||
// Bucket is one usage window: credits used against the plan limit and when
|
||||
// the window resets. Used/Limit are in z.ai credits, not tokens or dollars.
|
||||
type Bucket struct {
|
||||
ID string `json:"id"`
|
||||
Used float64 `json:"used"`
|
||||
Limit float64 `json:"limit"`
|
||||
WindowReset time.Time `json:"window_reset"`
|
||||
}
|
||||
|
||||
// UsedPct is the bucket's consumption ratio in percent of the limit.
|
||||
// A zero/negative limit reads as 0 (never blocks).
|
||||
func (b Bucket) UsedPct() float64 {
|
||||
if b.Limit <= 0 {
|
||||
return 0
|
||||
}
|
||||
return b.Used / b.Limit * 100
|
||||
}
|
||||
|
||||
// QuotaSnapshot is one account's quota state at a point in time, either
|
||||
// polled from the provider (Source "usage_url") or synthesized from locally
|
||||
// estimated consumption against the configured plan limits (Source
|
||||
// "estimate").
|
||||
type QuotaSnapshot struct {
|
||||
Account string `json:"account"`
|
||||
Source string `json:"source"`
|
||||
FetchedAt time.Time `json:"fetched_at"`
|
||||
Buckets []Bucket `json:"buckets"`
|
||||
}
|
||||
|
||||
// Bucket returns the bucket with the given id, if present.
|
||||
func (s *QuotaSnapshot) Bucket(id string) (Bucket, bool) {
|
||||
for _, b := range s.Buckets {
|
||||
if b.ID == id {
|
||||
return b, true
|
||||
}
|
||||
}
|
||||
return Bucket{}, false
|
||||
}
|
||||
|
||||
// MaxUsedPct is the highest consumption ratio across all buckets — the
|
||||
// number back-pressure thresholds compare against.
|
||||
func (s *QuotaSnapshot) MaxUsedPct() float64 {
|
||||
if s == nil {
|
||||
return 0
|
||||
}
|
||||
max := 0.0
|
||||
for _, b := range s.Buckets {
|
||||
if p := b.UsedPct(); p > max {
|
||||
max = p
|
||||
}
|
||||
}
|
||||
return max
|
||||
}
|
||||
|
||||
// providerUsage is the z.ai usage response shape this parser targets. z.ai
|
||||
// has not published the route yet; field aliases (used_credits|credits_used,
|
||||
// total_credits|limit_credits) keep the parser tolerant of either naming.
|
||||
type providerUsage struct {
|
||||
Usage struct {
|
||||
FiveHour providerBucket `json:"five_hour"`
|
||||
Weekly providerBucket `json:"weekly"`
|
||||
} `json:"usage"`
|
||||
}
|
||||
|
||||
type providerBucket struct {
|
||||
Used any `json:"used_credits"`
|
||||
UsedAlias any `json:"credits_used"`
|
||||
Limit any `json:"total_credits"`
|
||||
LimitAlias any `json:"limit_credits"`
|
||||
ResetAt string `json:"reset_at"`
|
||||
ResetAlias string `json:"reset_time"`
|
||||
}
|
||||
|
||||
// ParseUsage decodes a provider usage response into a QuotaSnapshot for the
|
||||
// account. It is strict about structure (unknown shapes error so a silent
|
||||
// HTML error page never reads as "quota fine") and lenient about field
|
||||
// naming within the documented buckets.
|
||||
func ParseUsage(account string, fetchedAt time.Time, body []byte) (*QuotaSnapshot, error) {
|
||||
var p providerUsage
|
||||
if err := json.Unmarshal(body, &p); err != nil {
|
||||
return nil, fmt.Errorf("usage decode: %w", err)
|
||||
}
|
||||
num := func(v any) (float64, bool) {
|
||||
switch n := v.(type) {
|
||||
case float64:
|
||||
return n, true
|
||||
case string:
|
||||
var f float64
|
||||
if _, err := fmt.Sscanf(n, "%g", &f); err == nil {
|
||||
return f, true
|
||||
}
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
snap := &QuotaSnapshot{Account: account, Source: "usage_url", FetchedAt: fetchedAt}
|
||||
for id, pb := range map[string]providerBucket{Bucket5h: p.Usage.FiveHour, BucketWeekly: p.Usage.Weekly} {
|
||||
used, ok := num(pb.Used)
|
||||
if !ok {
|
||||
used, ok = num(pb.UsedAlias)
|
||||
}
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("usage decode: bucket %q has no used_credits", id)
|
||||
}
|
||||
limit, ok := num(pb.Limit)
|
||||
if !ok {
|
||||
limit, ok = num(pb.LimitAlias)
|
||||
}
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("usage decode: bucket %q has no total_credits", id)
|
||||
}
|
||||
reset := pb.ResetAt
|
||||
if reset == "" {
|
||||
reset = pb.ResetAlias
|
||||
}
|
||||
rt := time.Time{}
|
||||
if reset != "" {
|
||||
var err error
|
||||
rt, err = time.Parse(time.RFC3339, reset)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("usage decode: bucket %q reset_at: %w", id, err)
|
||||
}
|
||||
}
|
||||
snap.Buckets = append(snap.Buckets, Bucket{ID: id, Used: used, Limit: limit, WindowReset: rt})
|
||||
}
|
||||
if len(snap.Buckets) != 2 {
|
||||
return nil, fmt.Errorf("usage decode: want five_hour and weekly buckets, got %d", len(snap.Buckets))
|
||||
}
|
||||
return snap, nil
|
||||
}
|
||||
@@ -0,0 +1,406 @@
|
||||
package quota
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"ukrrs.com/mopac/harness/internal/config"
|
||||
)
|
||||
|
||||
func TestParseUsage(t *testing.T) {
|
||||
ok := `{"usage":{"five_hour":{"used_credits":15400,"total_credits":28000,"reset_at":"2026-08-29T22:00:00Z"},
|
||||
"weekly":{"used_credits":130200,"total_credits":140000,"reset_at":"2026-09-02T00:00:00Z"}}}`
|
||||
alias := `{"usage":{"five_hour":{"credits_used":100,"limit_credits":200,"reset_time":"2026-08-29T22:00:00Z"},
|
||||
"weekly":{"credits_used":"150","limit_credits":"200","reset_time":"2026-09-02T00:00:00Z"}}}`
|
||||
bad := []struct {
|
||||
name, body string
|
||||
}{
|
||||
{"html error page", `<html><body>502 Bad Gateway</body></html>`},
|
||||
{"missing buckets", `{"usage":{}}`},
|
||||
{"missing used", `{"usage":{"five_hour":{"total_credits":28000},"weekly":{"used_credits":1,"total_credits":2}}}`},
|
||||
{"bad reset", `{"usage":{"five_hour":{"used_credits":1,"total_credits":2,"reset_at":"tomorrow"},"weekly":{"used_credits":1,"total_credits":2}}}`},
|
||||
}
|
||||
cases := []struct {
|
||||
name string
|
||||
body string
|
||||
want5h float64
|
||||
want5p float64
|
||||
wantWkd float64
|
||||
wantErr bool
|
||||
}{
|
||||
{"canonical shape", ok, 15400, 55.0, 93.0, false},
|
||||
{"field aliases + string numbers", alias, 100, 50, 75, false},
|
||||
}
|
||||
for _, c := range cases {
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
snap, err := ParseUsage("zai-1", time.Now(), []byte(c.body))
|
||||
if err != nil {
|
||||
t.Fatalf("ParseUsage: %v", err)
|
||||
}
|
||||
b5, ok := snap.Bucket(Bucket5h)
|
||||
if !ok || b5.Used != c.want5h || !nearly(b5.UsedPct(), c.want5p) {
|
||||
t.Errorf("5h bucket = %+v, want used=%v pct=%v", b5, c.want5h, c.want5p)
|
||||
}
|
||||
bw, _ := snap.Bucket(BucketWeekly)
|
||||
if !nearly(bw.UsedPct(), c.wantWkd) {
|
||||
t.Errorf("weekly pct = %v, want %v", bw.UsedPct(), c.wantWkd)
|
||||
}
|
||||
if got := snap.MaxUsedPct(); !nearly(got, c.wantWkd) {
|
||||
t.Errorf("MaxUsedPct = %v, want %v", got, c.wantWkd)
|
||||
}
|
||||
})
|
||||
}
|
||||
for _, c := range bad {
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
if _, err := ParseUsage("zai-1", time.Now(), []byte(c.body)); err == nil {
|
||||
t.Fatal("expected error, got nil (a garbage body must never read as 'quota fine')")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func nearly(a, b float64) bool {
|
||||
d := a - b
|
||||
return d < 1e-9 && d > -1e-9
|
||||
}
|
||||
|
||||
func TestParseUsageZeroLimitNeverBlocks(t *testing.T) {
|
||||
snap, err := ParseUsage("a", time.Now(), []byte(`{"usage":{"five_hour":{"used_credits":999,"total_credits":0,"reset_at":""},"weekly":{"used_credits":1,"total_credits":2}}}`))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if b, _ := snap.Bucket(Bucket5h); b.UsedPct() != 0 {
|
||||
t.Errorf("zero-limit bucket must read 0%%, got %v", b.UsedPct())
|
||||
}
|
||||
}
|
||||
|
||||
func TestGateUsesBearerAndNeverLogsKey(t *testing.T) {
|
||||
var gotAuth string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotAuth = r.Header.Get("Authorization")
|
||||
fmt.Fprintf(w, `{"usage":{"five_hour":{"used_credits":1,"total_credits":28000,"reset_at":"2026-08-29T22:00:00Z"},"weekly":{"used_credits":1,"total_credits":140000,"reset_at":"2026-09-02T00:00:00Z"}}}`)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
cfg := gateTestCfg()
|
||||
cfg.UsageURL = srv.URL
|
||||
cfg.KeyRef = "literal:seekrit-zai-key"
|
||||
gate := newGate(t, cfg)
|
||||
snap := gate.Snapshot(context.Background())
|
||||
if snap.Source != "usage_url" {
|
||||
t.Fatalf("source = %s, want usage_url", snap.Source)
|
||||
}
|
||||
if gotAuth != "Bearer seekrit-zai-key" {
|
||||
t.Errorf("auth header = %q", gotAuth)
|
||||
}
|
||||
// Poll failures (5xx) must fall back without panicking and never leak
|
||||
// the key into the returned error strings.
|
||||
srv.Config.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
http.Error(w, "boom", http.StatusInternalServerError)
|
||||
})
|
||||
gate2 := newGate(t, cfg)
|
||||
snap2 := gate2.Snapshot(context.Background())
|
||||
if snap2 == nil {
|
||||
t.Fatal("nil snapshot on poll failure; gate must degrade to estimates")
|
||||
}
|
||||
if strings.Contains(snap2.Source, "seekrit") {
|
||||
t.Errorf("key leaked into snapshot source %q", snap2.Source)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGateSharedSnapshotRoundTrip(t *testing.T) {
|
||||
fake := newFakeRedis(t)
|
||||
st, err := NewSharedState(fake.url(t))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer st.Close()
|
||||
|
||||
snap := &QuotaSnapshot{Account: "zai-1", Source: "usage_url", FetchedAt: time.Now(),
|
||||
Buckets: []Bucket{{ID: Bucket5h, Used: 1, Limit: 28000}}}
|
||||
if err := st.PublishSnapshot(snap, time.Minute); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, err := st.LoadSnapshot("zai-1")
|
||||
if err != nil || got == nil || got.Source != "usage_url" {
|
||||
t.Fatalf("round trip: %v %+v", err, got)
|
||||
}
|
||||
if b, _ := got.Bucket(Bucket5h); b.Used != 1 {
|
||||
t.Errorf("bucket lost: %+v", got.Buckets)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEstimateCredits(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
model string
|
||||
prompt, cached, comp int
|
||||
peak bool
|
||||
want float64
|
||||
}{
|
||||
// flagship: (10000*6.9 + 2000*1.7 + 10000*24)/10000 = 31.24
|
||||
{"flagship peak", "glm-5.3", 10000, 2000, 10000, true, 31.24},
|
||||
// off-peak halves: 15.62
|
||||
{"flagship off-peak", "glm-5.3", 10000, 2000, 10000, false, 15.62},
|
||||
// flash: (10000*2.3 + 2000*0.56 + 10000*8)/10000 = 10.412 -> /2
|
||||
{"flash off-peak", "glm-4.7-flash", 10000, 2000, 10000, false, 5.206},
|
||||
{"flash peak", "GLM-5.3-Flash", 10000, 0, 10000, true, 10.3},
|
||||
}
|
||||
for _, c := range cases {
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
got := EstimateCredits(c.model, c.prompt, c.cached, c.comp, c.peak)
|
||||
if diff := got - c.want; diff > 1e-9 || diff < -1e-9 {
|
||||
t.Errorf("credits = %v, want %v", got, c.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
if MultipliersFor("glm-9.9-omega") != FlagshipMultipliers() {
|
||||
t.Error("unknown model must read as flagship (conservative)")
|
||||
}
|
||||
}
|
||||
|
||||
func TestScheduleEdges(t *testing.T) {
|
||||
cst, err := time.LoadLocation("America/Chicago")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
at := func(day string, hm string) time.Time {
|
||||
tm, err := time.ParseInLocation("2006-01-02 15:04", day+" "+hm, cst)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return tm
|
||||
}
|
||||
// Default window 01:00-05:00 CST, weekdays only.
|
||||
sched, err := NewSchedule("01:00", "05:00", "America/Chicago", true)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cases := []struct {
|
||||
name string
|
||||
t time.Time
|
||||
want bool
|
||||
}{
|
||||
{"just before start (Fri)", at("2026-08-28", "00:59"), false},
|
||||
{"at start (Fri)", at("2026-08-28", "01:00"), true},
|
||||
{"mid window (Fri)", at("2026-08-28", "03:00"), true},
|
||||
{"last minute (Fri)", at("2026-08-28", "04:59"), true},
|
||||
{"at end (Fri)", at("2026-08-28", "05:00"), false},
|
||||
{"evening off-peak (Fri)", at("2026-08-28", "19:00"), false},
|
||||
{"weekend inside window (Sat)", at("2026-08-29", "02:00"), false},
|
||||
{"weekend inside window (Sun)", at("2026-08-30", "02:00"), false},
|
||||
{"Monday early window", at("2026-08-31", "02:00"), true},
|
||||
}
|
||||
for _, c := range cases {
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
if got := sched.InPeak(c.t); got != c.want {
|
||||
t.Errorf("InPeak(%s) = %v, want %v", c.t, got, c.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// Wrapped window 22:00-02:00, weekdays only: evening part carries the
|
||||
// start day's weekday; morning part checks the PREVIOUS day.
|
||||
wrap, err := NewSchedule("22:00", "02:00", "America/Chicago", true)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
wrapCases := []struct {
|
||||
name string
|
||||
t time.Time
|
||||
want bool
|
||||
}{
|
||||
{"Fri 23:00 (started Fri)", at("2026-08-28", "23:00"), true},
|
||||
{"Sat 01:00 (window started Fri)", at("2026-08-29", "01:00"), true},
|
||||
{"Sat 23:00 (started Sat)", at("2026-08-29", "23:00"), false},
|
||||
{"Sun 01:00 (window started Sat)", at("2026-08-30", "01:00"), false},
|
||||
}
|
||||
for _, c := range wrapCases {
|
||||
t.Run("wrap: "+c.name, func(t *testing.T) {
|
||||
if got := wrap.InPeak(c.t); got != c.want {
|
||||
t.Errorf("InPeak(%s) = %v, want %v", c.t, got, c.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// Every-day window ignores weekday flags.
|
||||
daily, _ := NewSchedule("01:00", "05:00", "America/Chicago", false)
|
||||
if !daily.InPeak(at("2026-08-29", "02:00")) { // Saturday
|
||||
t.Error("weekdays_only=false must peak on Saturday too")
|
||||
}
|
||||
|
||||
// UTC window evaluated on the INSTANT: 21:00 UTC is outside, but the
|
||||
// same wall-clock digits in CST (21:00 CDT == 02:00 UTC) are inside.
|
||||
utc, _ := NewSchedule("01:00", "05:00", "UTC", false)
|
||||
utcT := time.Date(2026, 8, 28, 21, 0, 0, 0, time.UTC)
|
||||
if utc.InPeak(utcT) {
|
||||
t.Error("21:00 UTC must be outside the UTC window")
|
||||
}
|
||||
cstWall := time.Date(2026, 8, 28, 21, 0, 0, 0, cst)
|
||||
if !utc.InPeak(cstWall) {
|
||||
t.Error("21:00 CST (== 02:00 UTC) must be inside the UTC window")
|
||||
}
|
||||
|
||||
if _, err := NewSchedule("25:00", "05:00", "UTC", false); err == nil {
|
||||
t.Error("25:00 must fail to parse")
|
||||
}
|
||||
if _, err := NewSchedule("01:00", "05:00", "Mars/Olympus", false); err == nil {
|
||||
t.Error("unknown timezone must fail")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecordTurnFeedsEstimateSnapshot(t *testing.T) {
|
||||
fake := newFakeRedis(t)
|
||||
cfg := gateTestCfg()
|
||||
cfg.RedisURL = fake.url(t)
|
||||
gate := newGate(t, cfg)
|
||||
// Friday 03:00 CST: inside the default peak window.
|
||||
fri := time.Date(2026, 8, 28, 3, 0, 0, 0, mustLoc(t, "America/Chicago"))
|
||||
gate.SetClock(func() time.Time { return fri })
|
||||
|
||||
credits := gate.RecordTurn(EstimateTurnInput{Model: "glm-5.3", PromptTokens: 10000, CachedTokens: 2000, CompletionTokens: 10000, Peak: true})
|
||||
if credits != 31.24 {
|
||||
t.Fatalf("credits = %v, want 31.24", credits)
|
||||
}
|
||||
|
||||
snap := gate.estimateSnapshot(fri)
|
||||
b5, _ := snap.Bucket(Bucket5h)
|
||||
bw, _ := snap.Bucket(BucketWeekly)
|
||||
if b5.Used < 31.23 || b5.Used > 31.25 || bw.Used < 31.23 || bw.Used > 31.25 {
|
||||
t.Errorf("estimate buckets after one turn: 5h=%v weekly=%v", b5.Used, bw.Used)
|
||||
}
|
||||
if b5.Limit != cfg.Plan5hCredits || bw.Limit != cfg.PlanWeeklyCredits {
|
||||
t.Errorf("limits = %v/%v, want plan limits", b5.Limit, bw.Limit)
|
||||
}
|
||||
|
||||
// Local-only mode (no redis) keeps its own accounting.
|
||||
local := newGate(t, gateTestCfg())
|
||||
local.SetClock(func() time.Time { return fri })
|
||||
local.RecordTurn(EstimateTurnInput{Model: "glm-4.7-flash", PromptTokens: 10000, CompletionTokens: 10000, Peak: false})
|
||||
ls := local.estimateSnapshot(fri)
|
||||
lb, _ := ls.Bucket(Bucket5h)
|
||||
if lb.Used < 5.14 || lb.Used > 5.16 {
|
||||
t.Errorf("local estimate = %v, want ~5.15", lb.Used)
|
||||
}
|
||||
}
|
||||
|
||||
func mustLoc(t *testing.T, name string) *time.Location {
|
||||
t.Helper()
|
||||
l, err := time.LoadLocation(name)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return l
|
||||
}
|
||||
|
||||
// TestDecideNineteenWall replays the 2026-08-28 19:00 quota wall: weekly
|
||||
// bucket exhausted at 19:00 CST — the gate must DEFER every class with a
|
||||
// surfaced reason instead of letting the turn die at the provider.
|
||||
func TestDecideNineteenWall(t *testing.T) {
|
||||
gate := newGate(t, gateTestCfg())
|
||||
cst := mustLoc(t, "America/Chicago")
|
||||
wall := time.Date(2026, 8, 28, 19, 0, 0, 0, cst) // Friday evening
|
||||
gate.SetClock(func() time.Time { return wall })
|
||||
gate.mu.Lock()
|
||||
gate.lastSnap = &QuotaSnapshot{Account: "zai-1", Source: "usage_url", FetchedAt: wall,
|
||||
Buckets: []Bucket{
|
||||
{ID: Bucket5h, Used: 12000, Limit: 28000},
|
||||
{ID: BucketWeekly, Used: 135800, Limit: 140000}, // 97%
|
||||
}}
|
||||
gate.mu.Unlock()
|
||||
|
||||
for _, class := range []string{"primary", "code", "study", "read", "review"} {
|
||||
d := gate.Decide(context.Background(), class)
|
||||
if d.Action != ActionDefer {
|
||||
t.Fatalf("class %s at 97%% weekly: action = %s, want defer", class, d.Action)
|
||||
}
|
||||
if !strings.Contains(d.Reason, "weekly") || !strings.Contains(d.Reason, "97%") {
|
||||
t.Errorf("defer reason must surface the bucket + ratio: %q", d.Reason)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDecideLevels(t *testing.T) {
|
||||
cfg := gateTestCfg()
|
||||
gate := newGate(t, cfg)
|
||||
// Wednesday noon CST: off-peak, mid-week.
|
||||
noon := time.Date(2026, 9, 2, 12, 0, 0, 0, mustLoc(t, "America/Chicago"))
|
||||
gate.SetClock(func() time.Time { return noon })
|
||||
snap := func(pct5, pctW float64) {
|
||||
gate.mu.Lock()
|
||||
defer gate.mu.Unlock()
|
||||
gate.lastSnap = &QuotaSnapshot{Account: cfg.Account, Source: "usage_url", FetchedAt: noon,
|
||||
Buckets: []Bucket{
|
||||
{ID: Bucket5h, Used: 28000 * pct5 / 100, Limit: 28000},
|
||||
{ID: BucketWeekly, Used: 140000 * pctW / 100, Limit: 140000},
|
||||
}}
|
||||
}
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
pct5, pctW float64
|
||||
class string
|
||||
want string
|
||||
wantIn string
|
||||
}{
|
||||
{"healthy allows flagship", 10, 10, "primary", ActionAllow, ""},
|
||||
{"healthy allows flash", 10, 10, "study", ActionAllow, ""},
|
||||
{"defer level: heavy defers", 90, 50, "code", ActionDefer, ">= defer"},
|
||||
{"defer level: flash continues", 90, 50, "study", ActionAllow, ""},
|
||||
{"block level: flash defers too", 50, 96, "study", ActionDefer, ">= block"},
|
||||
}
|
||||
for _, c := range cases {
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
snap(c.pct5, c.pctW)
|
||||
d := gate.Decide(context.Background(), c.class)
|
||||
if d.Action != c.want {
|
||||
t.Fatalf("action = %s (%s), want %s", d.Action, d.Reason, c.want)
|
||||
}
|
||||
if c.wantIn != "" && !strings.Contains(d.Reason, c.wantIn) {
|
||||
t.Errorf("reason %q missing %q", d.Reason, c.wantIn)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// Peak window restriction (quota healthy): Friday 03:00 CST.
|
||||
peak := time.Date(2026, 8, 28, 3, 0, 0, 0, mustLoc(t, "America/Chicago"))
|
||||
gate.SetClock(func() time.Time { return peak })
|
||||
snap(10, 10)
|
||||
if d := gate.Decide(context.Background(), "primary"); d.Action != ActionDefer || !strings.Contains(d.Reason, "peak window") {
|
||||
t.Errorf("peak + flagship class: %+v, want defer/peak reason", d)
|
||||
}
|
||||
if d := gate.Decide(context.Background(), "study"); d.Action != ActionAllow {
|
||||
t.Errorf("peak + flash class must allow: %+v", d)
|
||||
}
|
||||
// Block beats peak: even flash defers when the wall is hit in-peak.
|
||||
snap(10, 97)
|
||||
if d := gate.Decide(context.Background(), "study"); d.Action != ActionDefer {
|
||||
t.Errorf("peak + wall: %+v, want defer", d)
|
||||
}
|
||||
}
|
||||
|
||||
func gateTestCfg() config.QuotaConfig {
|
||||
return config.QuotaConfig{
|
||||
Enabled: true, Account: "zai-1",
|
||||
Plan5hCredits: 28000, PlanWeeklyCredits: 140000,
|
||||
PollIntervalSecs: 300, DeferAtPct: 85, BlockAtPct: 95,
|
||||
PeakStart: "01:00", PeakEnd: "05:00", Timezone: "America/Chicago",
|
||||
PeakWeekdaysOnly: true, PeakClasses: []string{"study", "read"},
|
||||
}
|
||||
}
|
||||
|
||||
func newGate(t *testing.T, cfg config.QuotaConfig) *Gate {
|
||||
t.Helper()
|
||||
gate, err := NewGate(cfg, nil, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return gate
|
||||
}
|
||||
|
||||
// (end of quota gate tests)
|
||||
@@ -0,0 +1,233 @@
|
||||
package quota
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// SharedState is the cross-instance quota state bus: one redis container on
|
||||
// the LAN (DECISION 2026-08-29, see REPORT-20260829-0500-quota) so the nine
|
||||
// harness instances across two hosts read/write the same quota snapshot and
|
||||
// credit estimates. It is a deliberately tiny stdlib-only RESP2 client
|
||||
// (SET/GET/INCRBYFLOAT/EXPIRE — the whole surface this package needs), no
|
||||
// host packages and no external Go deps. Every op fails soft: with redis
|
||||
// down the gate degrades to this instance's local estimate and the loop
|
||||
// keeps running.
|
||||
type SharedState struct {
|
||||
addr string
|
||||
db int
|
||||
mu sync.Mutex
|
||||
conn net.Conn
|
||||
rw *bufio.ReadWriter
|
||||
dialTO time.Duration
|
||||
}
|
||||
|
||||
// NewSharedState parses redis://host:port[/db]; empty URL returns nil (the
|
||||
// nil state is valid and fully functional as a no-op).
|
||||
func NewSharedState(redisURL string) (*SharedState, error) {
|
||||
if redisURL == "" {
|
||||
return nil, nil
|
||||
}
|
||||
u, err := url.Parse(redisURL)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("redis url: %w", err)
|
||||
}
|
||||
host := u.Host
|
||||
if host == "" {
|
||||
return nil, fmt.Errorf("redis url: missing host")
|
||||
}
|
||||
if !strings.Contains(host, ":") {
|
||||
host += ":6379"
|
||||
}
|
||||
db := 0
|
||||
if s := strings.TrimPrefix(u.Path, "/"); s != "" {
|
||||
if db, err = strconv.Atoi(s); err != nil {
|
||||
return nil, fmt.Errorf("redis url db: %w", err)
|
||||
}
|
||||
}
|
||||
return &SharedState{addr: host, db: db, dialTO: 3 * time.Second}, nil
|
||||
}
|
||||
|
||||
// Close releases the connection, if any.
|
||||
func (s *SharedState) Close() error {
|
||||
if s == nil {
|
||||
return nil
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.conn != nil {
|
||||
err := s.conn.Close()
|
||||
s.conn, s.rw = nil, nil
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// command runs one RESP command and returns the reply. Transport failures
|
||||
// reset the connection so the next op redials once.
|
||||
func (s *SharedState) command(args ...string) (string, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.conn == nil {
|
||||
conn, err := net.DialTimeout("tcp", s.addr, s.dialTO)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("redis dial %s: %w", s.addr, err)
|
||||
}
|
||||
s.conn = conn
|
||||
s.rw = bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn))
|
||||
if s.db != 0 {
|
||||
if _, err := s.execLocked("SELECT", strconv.Itoa(s.db)); err != nil {
|
||||
s.resetLocked()
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
}
|
||||
reply, err := s.execLocked(args...)
|
||||
if err != nil {
|
||||
s.resetLocked()
|
||||
}
|
||||
return reply, err
|
||||
}
|
||||
|
||||
// execLocked writes one command and reads one reply (mutex held).
|
||||
func (s *SharedState) execLocked(args ...string) (string, error) {
|
||||
var b strings.Builder
|
||||
fmt.Fprintf(&b, "*%d\r\n", len(args))
|
||||
for _, a := range args {
|
||||
fmt.Fprintf(&b, "$%d\r\n%s\r\n", len(a), a)
|
||||
}
|
||||
if _, err := s.rw.WriteString(b.String()); err != nil {
|
||||
return "", fmt.Errorf("redis write: %w", err)
|
||||
}
|
||||
if err := s.rw.Flush(); err != nil {
|
||||
return "", fmt.Errorf("redis flush: %w", err)
|
||||
}
|
||||
return readReply(s.rw.Reader)
|
||||
}
|
||||
|
||||
func (s *SharedState) resetLocked() {
|
||||
if s.conn != nil {
|
||||
_ = s.conn.Close()
|
||||
}
|
||||
s.conn, s.rw = nil, nil
|
||||
}
|
||||
|
||||
// readReply parses one RESP2 reply: +simple / -error / :integer / $bulk.
|
||||
func readReply(r *bufio.Reader) (string, error) {
|
||||
line, err := r.ReadString('\n')
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("redis read: %w", err)
|
||||
}
|
||||
line = strings.TrimRight(line, "\r\n")
|
||||
if line == "" {
|
||||
return "", errors.New("redis: empty reply")
|
||||
}
|
||||
switch line[0] {
|
||||
case '+', ':':
|
||||
return line[1:], nil
|
||||
case '-':
|
||||
return "", fmt.Errorf("redis: %s", line[1:])
|
||||
case '$':
|
||||
n, err := strconv.Atoi(line[1:])
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("redis bulk len: %w", err)
|
||||
}
|
||||
if n < 0 {
|
||||
return "", nil // nil bulk: missing key
|
||||
}
|
||||
buf := make([]byte, n+2)
|
||||
if _, err := ioReadFull(r, buf); err != nil {
|
||||
return "", fmt.Errorf("redis bulk read: %w", err)
|
||||
}
|
||||
return string(buf[:n]), nil
|
||||
default:
|
||||
return "", fmt.Errorf("redis: unexpected reply %q", line)
|
||||
}
|
||||
}
|
||||
|
||||
func ioReadFull(r *bufio.Reader, buf []byte) (int, error) {
|
||||
total := 0
|
||||
for total < len(buf) {
|
||||
n, err := r.Read(buf[total:])
|
||||
total += n
|
||||
if err != nil {
|
||||
return total, err
|
||||
}
|
||||
}
|
||||
return total, nil
|
||||
}
|
||||
|
||||
// Key layout (one account's plan is one shared universe):
|
||||
//
|
||||
// mopac:quota:<account>:snapshot — latest polled QuotaSnapshot (JSON)
|
||||
// mopac:quota:<account>:est:5h — estimated credits in the rolling 5h window
|
||||
// mopac:quota:<account>:est:weekly — estimated credits in the plan week
|
||||
//
|
||||
// Estimate keys carry the window id in the VALUE-side bookkeeping done by
|
||||
// the caller (Gate): 5h keys expire after 5h+slack; weekly keys are keyed
|
||||
// by ISO week via the caller and need no expiry.
|
||||
|
||||
func (s *SharedState) get(key string) (string, error) {
|
||||
if s == nil {
|
||||
return "", nil
|
||||
}
|
||||
return s.command("GET", key)
|
||||
}
|
||||
|
||||
func (s *SharedState) setEx(key, val string, ttl time.Duration) error {
|
||||
if s == nil {
|
||||
return nil
|
||||
}
|
||||
_, err := s.command("SET", key, val, "EX", strconv.Itoa(int(ttl.Seconds())))
|
||||
return err
|
||||
}
|
||||
|
||||
// incrByFloat adds delta to key (creating at delta) and returns the new value.
|
||||
func (s *SharedState) incrByFloat(key string, delta float64) (float64, error) {
|
||||
if s == nil {
|
||||
return 0, nil
|
||||
}
|
||||
reply, err := s.command("INCRBYFLOAT", key, strconv.FormatFloat(delta, 'f', 6, 64))
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return strconv.ParseFloat(strings.TrimSpace(reply), 64)
|
||||
}
|
||||
|
||||
// PublishSnapshot stores the polled snapshot for all instances to read.
|
||||
func (s *SharedState) PublishSnapshot(snap *QuotaSnapshot, ttl time.Duration) error {
|
||||
if s == nil || snap == nil {
|
||||
return nil
|
||||
}
|
||||
body, err := json.Marshal(snap)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return s.setEx(snapshotKey(snap.Account), string(body), ttl)
|
||||
}
|
||||
|
||||
// LoadSnapshot returns the last published snapshot, or nil when absent.
|
||||
func (s *SharedState) LoadSnapshot(account string) (*QuotaSnapshot, error) {
|
||||
body, err := s.get(snapshotKey(account))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if body == "" {
|
||||
return nil, nil
|
||||
}
|
||||
var snap QuotaSnapshot
|
||||
if err := json.Unmarshal([]byte(body), &snap); err != nil {
|
||||
return nil, fmt.Errorf("shared snapshot decode: %w", err)
|
||||
}
|
||||
return &snap, nil
|
||||
}
|
||||
|
||||
func snapshotKey(account string) string { return "mopac:quota:" + account + ":snapshot" }
|
||||
Reference in New Issue
Block a user