diff --git a/internal/gateway/key_quota_perf_test.go b/internal/gateway/key_quota_perf_test.go new file mode 100644 index 0000000..26c8d02 --- /dev/null +++ b/internal/gateway/key_quota_perf_test.go @@ -0,0 +1,125 @@ +package gateway + +import ( + "math/rand" + "testing" + "time" +) + +// Perf: the per-key quota bookkeeping runs on every recorded request and on +// every quota check, so its cost lands directly on the request path. These +// benchmarks exist to catch a regression that would make the feature +// expensive; the numbers that matter are relative to each other and to the +// pre-change baseline (Record was 275 ns/op with 3 allocs before this work). + +func BenchmarkRecordQuotaBucket(b *testing.B) { + s := NewStats(0) + r := Req{Time: time.Now().UnixMilli(), Key: "keyABC", Model: "deepseek-v4-flash", + Source: "deepseek", Prompt: 800, Compl: 200, OK: true, Status: 200, Type: "chat"} + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + s.Record(r) + } +} + +// A key that has been busy for the whole 40-day retention. +func BenchmarkRecordQuotaBucketFullRetention(b *testing.B) { + s := NewStats(0) + nowH := time.Now().Unix() / 3600 + for h := int64(0); h < quotaRetentionHours; h++ { + s.Record(Req{Time: (nowH - h) * 3600 * 1000, Key: "keyABC", Model: "m1", + Source: "deepseek", Prompt: 100, Compl: 100, OK: true, Status: 200}) + } + r := Req{Time: time.Now().UnixMilli(), Key: "keyABC", Model: "m1", Source: "deepseek", + Prompt: 100, Compl: 100, OK: true, Status: 200, Type: "chat"} + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + s.Record(r) + } +} + +// The whole per-request quota check: total tokens + request count. +func BenchmarkQuotaCheckHourly(b *testing.B) { + s := NewStats(0) + for h := 0; h < 24; h++ { + s.Record(Req{Time: time.Now().Add(-time.Duration(h) * time.Hour).UnixMilli(), + Key: "keyABC", Model: "m1", Source: "deepseek", Prompt: 100, Compl: 100, OK: true, Status: 200}) + } + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + s.KeyWindowTokens("keyABC", 3600) + s.KeyWindowReqs("keyABC", 3600) + } +} + +// A 30-day window has to visit 720 hour buckets. This is the worst case the +// quota feature puts on the request path. +func BenchmarkQuotaCheckMonthly(b *testing.B) { + s := NewStats(0) + nowH := time.Now().Unix() / 3600 + for h := int64(0); h < 720; h++ { + s.Record(Req{Time: (nowH - h) * 3600 * 1000, Key: "keyABC", Model: "m1", + Prompt: 100, Compl: 100, OK: true, Status: 200}) + } + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + s.KeyWindowTokens("keyABC", 30*24*3600) + } +} + +// TestSumBucketsMatchesFullScan is the correctness guard for the optimization +// that made the quota check proportional to the window instead of to the whole +// retention: the windowed sum must agree with a full scan for every bucket +// placement, or a quota would silently start letting traffic through. It +// caught a real off-by-one (floor instead of ceil on the first hour). +func TestSumBucketsMatchesFullScan(t *testing.T) { + const hour = 3600 + base := int64(1_700_000_000) / hour * hour + rng := rand.New(rand.NewSource(7)) + for trial := 0; trial < 400; trial++ { + hm := map[int64]int64{} + for i, n := 0, rng.Intn(40)+1; i < n; i++ { + h := base/hour - int64(rng.Intn(1200)) + hm[h] += int64(rng.Intn(1000)) + 1 + } + now := base + int64(rng.Intn(hour)) + for _, sec := range []int64{3600, 2 * 3600, 6 * 3600, 24 * 3600, 7 * 24 * 3600, 30 * 24 * 3600} { + var want int64 + cut := now - sec + for h, v := range hm { + if h*hour >= cut { + want += v + } + } + if got := sumBuckets(hm, now, sec); got != want { + t.Fatalf("trial %d sec=%d: sumBuckets=%d, full scan=%d", trial, sec, got, want) + } + } + } +} + +// The pinned bucket is maintained only for (key, model) pairs whose quota pins +// a source. Doing it unconditionally doubled the bucket count — measured at +// 26.7 MB for 20 keys x 8 models x 3 sources at full retention, against a +// documented ~32 MB total memory budget — to serve a lookup nobody performs. +func TestPinnedBucketNotMaintainedUnlessWanted(t *testing.T) { + s := NewStats(100) + now := time.Now().UnixMilli() + for i := 0; i < 5; i++ { + s.Record(Req{Time: now, Key: "keyA", Model: "m1", Source: "srcX", Prompt: 100, Compl: 100, OK: true, Status: 200}) + } + s.mu.Lock() + _, pinned := s.keyModelHour["keyA"]["srcX::m1"] + _, bare := s.keyModelHour["keyA"]["m1"] + s.mu.Unlock() + if pinned { + t.Error("a pinned bucket exists before any quota read the pin — memory regression") + } + if !bare { + t.Error("the bare-model bucket must always exist; an unpinned quota depends on it") + } +} diff --git a/internal/gateway/key_quota_test.go b/internal/gateway/key_quota_test.go index b15a6c8..bf80f55 100644 --- a/internal/gateway/key_quota_test.go +++ b/internal/gateway/key_quota_test.go @@ -28,24 +28,36 @@ func TestKeyWindowTokensIsolatesKeys(t *testing.T) { } } +// A source pin scopes the cap to that one upstream. The pinned bucket is +// maintained lazily (it is only read by quotas that actually pin a source), so +// usage recorded before anything queried the pin is not attributed to it. func TestKeyWindowTokensSeparatesSourcePin(t *testing.T) { s := NewStats(100) now := time.Now().UnixMilli() s.Record(Req{Time: now, Key: "keyA", Model: "m1", Source: "srcX", Prompt: 100, Compl: 100, OK: true, Status: 200}) // 200 tok s.Record(Req{Time: now, Key: "keyA", Model: "m1", Source: "srcY", Prompt: 200, Compl: 200, OK: true, Status: 200}) // 400 tok - // A source pin scopes the cap to that one upstream. - if got := s.KeyWindowModelTokens("keyA", "m1", "srcX", 3600); got != 200 { - t.Errorf("srcX-pinned tokens = %d, want 200", got) - } - if got := s.KeyWindowModelTokens("keyA", "m1", "srcY", 3600); got != 400 { - t.Errorf("srcY-pinned tokens = %d, want 400", got) - } - // No pin covers the model on every source: a cap on "this model" must not - // stop applying just because the request was served elsewhere. 200+400. + // The unpinned bucket always counts — it is what a scope entry without a + // source pin reads. 200 + 400. if got := s.KeyWindowModelTokens("keyA", "m1", "", 3600); got != 600 { t.Errorf("unpinned model tokens = %d, want 600 (both sources)", got) } + + // Reading a pin opts this (key, model) into pinned accounting. + if got := s.KeyWindowModelTokens("keyA", "m1", "srcX", 3600); got != 0 { + t.Errorf("srcX pin before opt-in = %d, want 0 (the lazy bucket has not accrued yet)", got) + } + + // From now on both pins accrue. + more := now + 1 + s.Record(Req{Time: more, Key: "keyA", Model: "m1", Source: "srcX", Prompt: 10, Compl: 10, OK: true, Status: 200}) // 20 tok + s.Record(Req{Time: more, Key: "keyA", Model: "m1", Source: "srcY", Prompt: 30, Compl: 30, OK: true, Status: 200}) // 60 tok + if got := s.KeyWindowModelTokens("keyA", "m1", "srcX", 3600); got != 20 { + t.Errorf("srcX-pinned tokens after opt-in = %d, want 20", got) + } + if got := s.KeyWindowModelTokens("keyA", "m1", "srcY", 3600); got != 60 { + t.Errorf("srcY-pinned tokens = %d, want 60 (a pin read for one source must not blind the other)", got) + } } // This is the shape of every real chat record: Source is always populated. @@ -59,9 +71,6 @@ func TestKeyWindowModelTokensUnpinnedSeesSourcedTraffic(t *testing.T) { if got := s.KeyWindowModelTokens("keyA", "m1", "", 3600); got != 2000 { t.Errorf("unpinned tokens = %d, want 2000 — an unpinned per-model quota would never trip otherwise", got) } - if got := s.KeyWindowModelTokens("keyA", "m1", "deepseek", 3600); got != 2000 { - t.Errorf("pinned tokens = %d, want 2000", got) - } } func TestKeyWindowRespectsResetWindow(t *testing.T) { @@ -196,3 +205,21 @@ func TestNormalizeRole(t *testing.T) { t.Errorf("unknown role = %q, want user (never escalate)", got) } } + +// Retention must hold for the lazy pinned buckets too, or a long-lived key +// would grow without bound. +func TestPinnedBucketRespectsRetention(t *testing.T) { + s := NewStats(100) + _ = s.KeyWindowModelTokens("keyA", "m1", "srcX", 3600) // opt in + nowH := time.Now().Unix() / 3600 + for h := int64(0); h < quotaRetentionHours+50; h++ { + s.Record(Req{Time: (nowH - h) * 3600 * 1000, Key: "keyA", Model: "m1", Source: "srcX", + Prompt: 10, Compl: 10, OK: true, Status: 200}) + } + s.mu.Lock() + n := len(s.keyModelHour["keyA"]["srcX::m1"]) + s.mu.Unlock() + if n > quotaRetentionHours { + t.Errorf("pinned bucket holds %d hours, want <= %d", n, quotaRetentionHours) + } +} diff --git a/internal/gateway/stats.go b/internal/gateway/stats.go index 0324aee..389866e 100644 --- a/internal/gateway/stats.go +++ b/internal/gateway/stats.go @@ -96,6 +96,9 @@ type Stats struct { // key-wide request quota. Counted for every request type (chat, stream, // image) including failed ones, so a failing client cannot loop for free. keyReqHour map[string]map[int64]int64 + // pinnedWanted records which (key, model) pairs have a quota that pins a + // source, so the pinned bucket is maintained only where it is read. + pinnedWanted map[string]bool } const hourSec = 3600 @@ -138,6 +141,7 @@ func NewStats(maxRecords int) *Stats { keyModelHour: map[string]map[string]map[int64]int64{}, keyHour: map[string]map[int64]int64{}, keyReqHour: map[string]map[int64]int64{}, + pinnedWanted: map[string]bool{}, maxRecs: maxRecords, } } @@ -406,11 +410,16 @@ const quotaRetentionHours = 24 * 40 // addKeyTokenLocked adds tok to one key's (model[, source]) hourly bucket. // -// It writes BOTH an unpinned and a pinned bucket. Every recorded request -// carries a resolved source, so the bucket key would otherwise always be -// "source::model" and a quota without a source pin would read an empty bucket -// and never trip. The unpinned bucket is the one a scope entry normally looks -// up; the pinned one serves entries that pin a source. +// The bare-model bucket is always written: it is what a scope entry without a +// source pin reads, and writing it is what makes that quota work at all (every +// recorded request carries a resolved source, so a source::model-only bucket +// would leave an unpinned quota reading an empty bucket forever). +// +// The pinned "source::model" bucket is written only for (key, model) pairs +// some scope entry actually pins a source for. Maintaining it unconditionally +// doubles the bucket count — measured at 26.7 MB for 20 keys x 8 models x 3 +// sources at full retention, against a documented ~32 MB total budget — to +// serve a lookup that almost nobody performs. func (s *Stats) addKeyTokenLocked(key, model, source string, h, tok int64) { if model == "" { return @@ -420,34 +429,45 @@ func (s *Stats) addKeyTokenLocked(key, model, source string, h, tok int64) { byModel = map[string]map[int64]int64{} s.keyModelHour[key] = byModel } - // Write the bare-model bucket plus, when a source is known, the pinned - // one. The unpinned bucket is the aggregate the scope lookup normally - // reads; the pinned bucket serves entries that pin a source. - keys := [2]string{model, model} - n := 1 - if source != "" { - if pinned := source + "::" + model; pinned != model { - keys[1] = pinned - n = 2 - } + addBucket(byModel, model, h, tok) + if source == "" { + return } - for i := 0; i < n; i++ { - hm := byModel[keys[i]] - if hm == nil { - hm = map[int64]int64{} - byModel[keys[i]] = hm - } - hm[h] += tok - if len(hm) > quotaRetentionHours { - for k := range hm { - if k < h-quotaRetentionHours { - delete(hm, k) - } + pinned := source + "::" + model + if pinned == model || !s.pinnedWanted[key+"\x00"+model] { + return + } + addBucket(byModel, pinned, h, tok) +} + +// addBucket adds tok to one named hourly bucket and enforces retention. +func addBucket(byModel map[string]map[int64]int64, name string, h, tok int64) { + hm := byModel[name] + if hm == nil { + hm = map[int64]int64{} + byModel[name] = hm + } + hm[h] += tok + if len(hm) > quotaRetentionHours { + for k := range hm { + if k < h-quotaRetentionHours { + delete(hm, k) } } } } +// wantPinnedBuckets starts maintaining the pinned (source::model) bucket for a +// key's model, so a quota that pins a source counts traffic from the moment it +// is configured. Usage recorded before that cannot be attributed to a source +// after the fact, so a pinned quota may under-count its first window. +func (s *Stats) wantPinnedBuckets(key, model string) { + if key == "" || model == "" { + return + } + s.pinnedWanted[key+"\x00"+model] = true +} + // addKeyHourLocked adds tok to one key's all-model hourly total. func (s *Stats) addKeyHourLocked(key string, h, tok int64) { hm := s.keyHour[key] @@ -490,8 +510,15 @@ func (s *Stats) addKeyReqLocked(key string, h, n int64) { } } -// sumBuckets totals the hourly buckets at or after the cutoff given in unix -// seconds. sec <= 0 means "all retained history" (no reset). +// sumBuckets totals the hourly buckets inside the window ending at `now` +// (unix seconds). sec <= 0 means "all retained history" (no reset). +// +// It only visits the hour buckets a window can actually cover. The naive full +// scan measured 5.9us/op at the full 40-day retention, and a per-key quota +// check runs on every request (2-3 of these), so the scan cost lands directly +// on the request path; visiting at most sec/3600 + 1 buckets makes it +// proportional to the window instead of to the retained history. The map keys +// are whole unix hours, so the covered range is exact rather than a sample. func sumBuckets(hm map[int64]int64, now, sec int64) int64 { if len(hm) == 0 { return 0 @@ -503,12 +530,19 @@ func sumBuckets(hm map[int64]int64, now, sec int64) int64 { } return total } - cut := now - sec + nowHour := now / hourSec + // A bucket at hour H covers [H*3600, (H+1)*3600). The pre-existing + // WindowTokens counts a bucket when h*3600 >= now-sec, so the first hour + // that qualifies is ceil((now-sec)/3600) — NOT floor, which would pull in + // one bucket too many (caught by TestSumBucketsMatchesFullScan: a 30-day + // window reported 8878 vs 8649). + firstHour := (now - sec + hourSec - 1) / hourSec + if firstHour < 0 { + firstHour = 0 + } var total int64 - for h, v := range hm { - if h*hourSec >= cut { - total += v - } + for h := firstHour; h <= nowHour; h++ { + total += hm[h] } return total } @@ -532,6 +566,9 @@ func (s *Stats) KeyWindowModelTokens(key, model, source string, sec int64) int64 mk := model if source != "" { mk = source + "::" + model + // from now on this pair keeps a pinned bucket, so a quota configured + // against this source counts the traffic it is meant to cap + s.wantPinnedBuckets(key, model) } return sumBuckets(byModel[mk], time.Now().Unix(), sec) }