mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-10-03 23:54:06 +00:00
perf(gateway): 配额桶扫描改为窗口化 + pinned 桶惰性创建
审查本特性线的性能时发现两个问题，均有实测数据。
## 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 be11a06a46)
This commit is contained in:
JianFeeeee
125
internal/gateway/key_quota_perf_test.go
Normal file
125
internal/gateway/key_quota_perf_test.go
Normal file
@ -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")
|
||||
}
|
||||
}
|
||||
@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@ -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)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user