test(keys): per-key AUTO 链的并发与压力验证

AutoChainFor 在请求热路径上无锁读 keyAutoChains,而 SaveKeyAuto 与
rebuildRegistry 会整表替换它。顺序调用测不到任何交错,所以补并发判据。

- 8 读者 + 2 写者并发改链,断言读到的链必属于 {自身链, 全局链} 这个合法
  集合,不断言固定值(写者本来就在合法地来回切换,钉死值会在正确代码上
  失败)。
- 源增删期间读者保持可用链。churn 源用 "churn" 前缀,避免与链依赖的 s1
  重名——删掉链引用的源会合法地清空该链,拿它当断言等于测虚构。
- 500 个 key 的查表不退化。

变异验证:把 rebuildKeyAutoChains 改成原地清空重填(破坏原子替换),
-race 报 DATA RACE,判据有效。

真实压力(20 key 网关 + 并发 AUTO + 后台持续改链):
- 40 并发 25s:9616 请求 / 382 QPS / 零 panic / 零 DATA RACE
- 出现 111 次 503,根因是 mock 上游被打爆后 "no free slot within 2s",
  属预期降级而非缺陷:并发降到 6 时 1395/1395 全 200。
- 每次压测都断言"未改链的用户始终命中自己该走的模型",未改链的 19 个
  key 全程零错分。
This commit is contained in:
JianFeeeee
2026-10-03 08:01:37 +08:00
parent 22decf2bd3
commit c66e1b5955

View File

@ -0,0 +1,189 @@
package core
import (
"fmt"
"sync"
"sync/atomic"
"testing"
"time"
"llmsproxy/internal/config"
)
// AutoChainFor sits on the request hot path and reads keyAutoChains with no
// lock, while SaveKeyAuto / rebuildRegistry replace that map from the admin
// path. A "call it twice" test never interleaves those, so it would pass even
// if the map were mutated in place. These tests run both sides at once under
// -race and assert what readers actually observe.
//
// The property: a reader sees a chain that is either the key's own or the
// global one — never a torn or empty one. If rebuildKeyAutoChains built the
// map in place instead of swapping a fresh pointer, readers would sometimes
// see a partially-filled map and silently fall back to the global chain, i.e.
// a user quietly losing their configured chain under load.
// concurrentChainConfig builds N keys, half with their own chain pointing at a
// distinct model so a fallback to the global chain is distinguishable.
func concurrentChainConfig(t *testing.T, n int) *config.Config {
t.Helper()
cfg := perKeyTestConfig(t)
cfg.Keys = nil
cfg.GatewayKeys = []string{"sk-gw-admin"}
cfg.Auto = []config.ModelScope{{Model: "gpt-4o", Source: "s1", Tier: 1}}
for i := 0; i < n; i++ {
k := config.GWKey{
Key: fmt.Sprintf("sk-gw-k%03d", i),
Role: "user",
Name: fmt.Sprintf("user%d", i),
}
if i%2 == 0 {
k.Auto = []config.ModelScope{{Model: "gpt-4o-mini", Source: "s1", Tier: 1}}
}
cfg.Keys = append(cfg.Keys, k)
}
return cfg
}
func TestAutoChainForConcurrentWithSaveKeyAuto(t *testing.T) {
cfg := concurrentChainConfig(t, 40)
c := newTestCore(t, cfg)
var stop atomic.Bool
var wg sync.WaitGroup
var bad atomic.Int64
for r := 0; r < 8; r++ {
wg.Add(1)
go func() {
defer wg.Done()
for !stop.Load() {
for i := 0; i < 40; i++ {
key := fmt.Sprintf("sk-gw-k%03d", i)
ch, ok := c.AutoChainFor(key)
if !ok || ch == nil || len(ch.Tiers) == 0 {
bad.Add(1)
continue
}
// Legal values are this key's own chain model (either
// side of the writer's flip) or the global chain's. The
// assertion is about the legal set, not a fixed value:
// the writers legitimately change the chain underneath.
got := ch.Tiers[0].Slots[0].Model
if got != "gpt-4o" && got != "gpt-4o-mini" {
bad.Add(1)
}
}
}
}()
}
for w := 0; w < 2; w++ {
wg.Add(1)
go func() {
defer wg.Done()
model := "gpt-4o-mini"
for i := 0; i < 25 && !stop.Load(); i++ {
for j := 0; j < 40; j += 2 {
if err := c.SaveKeyAuto(fmt.Sprintf("sk-gw-k%03d", j), []config.ModelScope{
{Model: model, Source: "s1", Tier: 1},
}); err != nil {
bad.Add(1)
return
}
}
if model == "gpt-4o-mini" {
model = "gpt-4o"
} else {
model = "gpt-4o-mini"
}
}
}()
}
time.Sleep(400 * time.Millisecond)
stop.Store(true)
wg.Wait()
if n := bad.Load(); n > 0 {
t.Fatalf("%d reads saw a chain outside {own, global} — the cache is "+
"being swapped in place instead of atomically", n)
}
}
// The registry rebuild path (AddSource / RemoveSource) also replaces the
// cache, so readers must keep a usable chain throughout.
//
// Churned sources use a "churn" prefix on purpose: deleting a source a chain
// references legitimately empties that chain (BuildChain drops unresolvable
// slots), and asserting otherwise would be testing a fiction rather than the
// code.
func TestAutoChainForConcurrentWithSourceChanges(t *testing.T) {
cfg := concurrentChainConfig(t, 20)
c := newTestCore(t, cfg)
var stop atomic.Bool
var wg sync.WaitGroup
var bad atomic.Int64
for r := 0; r < 6; r++ {
wg.Add(1)
go func() {
defer wg.Done()
for !stop.Load() {
for i := 0; i < 20; i++ {
ch, ok := c.AutoChainFor(fmt.Sprintf("sk-gw-k%03d", i))
if !ok || ch == nil || len(ch.Tiers) == 0 {
bad.Add(1)
continue
}
if got := ch.Tiers[0].Slots[0].Model; got != "gpt-4o" && got != "gpt-4o-mini" {
bad.Add(1)
}
}
}
}()
}
wg.Add(1)
go func() {
defer wg.Done()
for n := 0; n < 30 && !stop.Load(); n++ {
name := fmt.Sprintf("churn%d", n)
_ = c.AddSource(config.Source{
Name: name, BaseURL: "http://127.0.0.1:9/v1", Adapter: "openai",
Models: []config.Model{{ID: "m-" + name, Priority: 10}}},
)
_ = c.RemoveSource(name)
}
}()
time.Sleep(500 * time.Millisecond)
stop.Store(true)
wg.Wait()
if n := bad.Load(); n > 0 {
t.Fatalf("%d reads lost the chain while unrelated sources were churned", n)
}
}
// Many distinct keys: the lookup must stay a map hit and must not degrade into
// returning the global chain as the key count grows.
func TestAutoChainForManyKeys(t *testing.T) {
cfg := concurrentChainConfig(t, 500)
c := newTestCore(t, cfg)
for i := 0; i < 500; i++ {
key := fmt.Sprintf("sk-gw-k%03d", i)
ch, _ := c.AutoChainFor(key)
if ch == nil || len(ch.Tiers) == 0 {
t.Fatalf("key %s lost its chain with 500 keys configured", key)
}
want := "gpt-4o"
if i%2 == 0 {
want = "gpt-4o-mini"
}
if got := ch.Tiers[0].Slots[0].Model; got != want {
t.Fatalf("key %s got %q, want %q", key, got, want)
}
}
}