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,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")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user