From d072a03c9a5c73140ae08315a15234e53421a40a Mon Sep 17 00:00:00 2001 From: JianFeeeee Date: Sun, 27 Sep 2026 18:34:13 +0800 Subject: [PATCH] =?UTF-8?q?perf(gateway):=20=E9=85=8D=E9=A2=9D=E6=A1=B6?= =?UTF-8?q?=E6=89=AB=E6=8F=8F=E6=94=B9=E4=B8=BA=E7=AA=97=E5=8F=A3=E5=8C=96?= =?UTF-8?q?=20+=20pinned=20=E6=A1=B6=E6=83=B0=E6=80=A7=E5=88=9B=E5=BB=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 审查本特性线的性能时发现两个问题，均有实测数据。 ## 1. 窗口查询是全扫，代价落在每个请求上 sumBuckets 原来遍历整个 map（最多 960 个小时桶），实测 5.9us/op。 配额检查在每个请求上跑 2-3 次（key 总 token、key 请求数、scope token）， 于是单请求多付约 18us。 注意这**不是本改动引入的成本**：main 上既有的 WindowTokens 同样是 5907ns/op（全扫）。是本改动让它在请求路径上被调用得更多。 改为只遍历窗口可能覆盖的桶（键是整点小时，范围是精确的，不是采样）： - 24h 窗口 200ns -> 55ns - 1h 窗口 80ns -> 49ns - 30d 窗口 5.9us -> 3.9us（720 次查找，只有配 month 配额时才走到） 等价性由 TestSumBucketsMatchesFullScan 保证（400 组随机桶位置 x 6 种 窗口，对全扫逐项比对）。★ 第一次写错成 floor，判据立刻抓到： 30 天窗口报 8878 而全扫是 8649 —— 正确是 ceil。 ## 2. pinned 桶无条件创建，内存最坏 26.7MB 每条记录写两个桶：裸 model 与 "source::model"。但 pinned 桶只有 「配额里显式 pin 了 source」时才会被查。 实测最坏情况（20 key x 8 model x 3 source x 40 天全 retention）： HeapAlloc 26.67MB —— 而 README 宣传「16 源生产实例 ~32-35MB」， 等于吃掉 80% 内存预算。 改为惰性：只有 KeyWindowModelTokens 带 source 查询过某个 (key, model) 之后，才开始维护它的 pinned 桶。 20key x 8model x 3src 26.67MB -> 8.87MB (-67%) 5key x 6model（真实） 3.36MB -> 2.26MB (-33%) 5key x 12model 5.82MB -> 3.59MB (-38%) 代价：配 pinned 配额之前发生的用量无法事后按源拆分（记录里虽然有 Source，但桶只存了裸 model），所以 pinned 配额的首个窗口可能少算。 已在代码注释与判据中写明。 ## 其余实测 Record main 基线 275ns/429B/3allocs -> 283ns/429B/3allocs （+8ns，分配数不变；3 allocs 来自 ring buffer） 配额检查全路径 115ns / 0 allocs（每请求新增） 纯读路径 6.5ns / 0 allocs ## 100 并发调度/拒绝压测（真实进程 + 可报并发峰值的假上游） 容量 100（4+96），0.15s/请求，100 并发 ok=100 fail=0 上游峰值 42 容量 100，0.15s/请求，200 并发 ok=200 fail=0 上游峰值 97 容量 8，3s/请求，100 并发 ok=8 fail=92 上游峰值 8 容量 8，0.6s/请求，100 并发 ok=32 fail=68 上游峰值 8 容量 8，0.6s/请求，40 并发 ok=32 fail=8 上游峰值 8 上游峰值恒定不超过 max_concurrent，容量拒绝返回 503 + busyWait 有界 等待（约 2.6s）。main 基线在同条件下 ok=32 fail=68、上游峰值 8、 延迟分布相同 —— 配额改动没有触碰调度/拒绝路径。 配额拒绝单独验证（低并发避开容量拒绝）：req_quota=50 用尽后 100 并发全部 429 rate_limit_exceeded + Retry-After: 1661， **延迟仅 19-21ms**、上游 total 未增加 —— 配额在入口廉价拒绝， 不占用任何上游槽位，与容量不足的昂贵等待形成明确分工。 ## 判据 key_quota_perf_test.go：4 个基准 + 2 个判据（sumBuckets 等价性、 pinned 桶惰性）。3/3 变异全被抓（无条件建 pinned 桶、firstHour 用 floor、keyHour 不再写）。 (cherry picked from commit be11a06a46c6730245a801d833de1f34d4dc240a) --- internal/gateway/key_quota_perf_test.go | 125 ++++++++++++++++++++++++ internal/gateway/key_quota_test.go | 51 +++++++--- internal/gateway/stats.go | 105 +++++++++++++------- 3 files changed, 235 insertions(+), 46 deletions(-) create mode 100644 internal/gateway/key_quota_perf_test.go 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) }