Sakurasan
8c12207478
按性能分析逐项落地，实测对比（16 并发、mock 上游零延迟、SQLite）： - 每请求 SQL 10 → 2：模型/渠道目录（模型行、绑定、渠道行）进程内缓存 10s， 运行期状态（冷却、健康）仍每次实时判定；管理端写操作经中间件、 健康状态变化经回调立即失效。新增 TestProxyRequestSQLBudget 钉住预算。 - 端到端 p50 1.3–1.6ms → 0.66–0.70ms，RPS 3.5–8.9k → 16–18k。 - 流式正文累积：字符串 += 改 strings.Builder（原为 O(n²)）， 逐行 JSON 提取改字节快路径，并加廉价的子串预过滤、拿到输出 usage 后停累积。 500 块：949µs/7.9MB 分配 → 83µs/136KB。 - 请求体定点改写：模型映射从整包 unmarshal/marshal 改为顶层字段替换 （100KB：585µs → 29µs），stream_options 注入同理；requestText 改为 只在「上游没报用量」时按需提取，亲和键改用原始体前缀。 - 渠道满载背压：新增 Acquire 有界排队（OT_PROXY_CHANNEL_ACQUIRE_WAIT， 默认 150ms）；排队仍满时返回 503 + Retry-After 而不是 502。 实测渠道上限 4、16 并发：无排队 81 次 503；排队 150ms 后 400/400 成功。 - 记账按批聚合：同批同用户一次扣款、流水批量插入、日表按 (user,model,日期) 聚合，替代每行 5 条语句；幂等仍靠 ref_id 唯一索引兜底。 - gorm 开启 PrepareStmt（SQL 解析此前占 CPU 15%）；tiktoken 编码器按模型复用。 - 管理端用量列表不再返回原始请求/响应体（可能几十 MB）， 改为 has_raw 标记 + GET /admin/usage/:id 按需拉取，前端弹窗打开时取。 验证：go vet、go test ./...、go test -race（proxy/channel/usage/tokenizer）、 vue-tsc、eslint、vite build 全绿；本地起 mock 上游实测列表/详情接口与 管理页弹窗按需加载。
443 lines
14 KiB
Go
443 lines
14 KiB
Go
package channel
|
|
|
|
import (
|
|
"context"
|
|
"os"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/openteam/server/internal/pkg/crypto"
|
|
"github.com/openteam/server/internal/store"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
func newTestDB(t *testing.T) *gorm.DB {
|
|
t.Helper()
|
|
dsn := "/tmp/channel_test_" + time.Now().Format("150405.000000000") + ".db"
|
|
db, err := store.Open("sqlite", dsn)
|
|
if err != nil {
|
|
t.Fatalf("open db: %v", err)
|
|
}
|
|
t.Cleanup(func() {
|
|
sqlDB, _ := db.DB()
|
|
sqlDB.Close()
|
|
os.Remove(dsn)
|
|
})
|
|
return db
|
|
}
|
|
|
|
func seedChannel(t *testing.T, db *gorm.DB, name string, weight, priority, maxConc int) uint64 {
|
|
t.Helper()
|
|
enc := crypto.NewEncryptor("test-master-key-1234567890")
|
|
key, _ := enc.Encrypt("upstream-key")
|
|
ch := store.Channel{
|
|
Name: name, Provider: store.ChannelProviderOpenAI, BaseURL: "http://localhost:9000",
|
|
APIKeyEnc: key, Weight: weight, Priority: priority, MaxConcurrency: maxConc,
|
|
HealthStatus: store.ChannelHealthHealthy, Enabled: true,
|
|
}
|
|
if err := db.Create(&ch).Error; err != nil {
|
|
t.Fatalf("create channel: %v", err)
|
|
}
|
|
return ch.ID
|
|
}
|
|
|
|
func TestCandidatesFiltersUnhealthy(t *testing.T) {
|
|
db := newTestDB(t)
|
|
seedChannel(t, db, "a", 1, 0, 4)
|
|
seedChannel(t, db, "b", 5, 1, 4)
|
|
// 禁用渠道不入候选
|
|
bad := store.Channel{Name: "bad", Provider: store.ChannelProviderOpenAI, BaseURL: "http://x", APIKeyEnc: "x", HealthStatus: store.ChannelHealthCooldown, Enabled: true}
|
|
db.Create(&bad)
|
|
|
|
s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890"))
|
|
cands := s.Candidates("")
|
|
if len(cands) != 2 {
|
|
t.Fatalf("candidates = %d, want 2", len(cands))
|
|
}
|
|
if cands[0].Channel.Name != "a" {
|
|
t.Fatalf("first by priority should be a, got %s", cands[0].Channel.Name)
|
|
}
|
|
}
|
|
|
|
func TestPickWeighted(t *testing.T) {
|
|
db := newTestDB(t)
|
|
seedChannel(t, db, "only", 1, 0, 4)
|
|
s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890"))
|
|
cands := s.Candidates("")
|
|
chosen := s.Pick(cands)
|
|
if chosen == nil || chosen.Name != "only" {
|
|
t.Fatalf("pick single: %+v", chosen)
|
|
}
|
|
// 空候选返回 nil
|
|
if p := s.Pick(nil); p != nil {
|
|
t.Fatal("expected nil for empty candidates")
|
|
}
|
|
}
|
|
|
|
func TestTryAcquire(t *testing.T) {
|
|
db := newTestDB(t)
|
|
id := seedChannel(t, db, "c", 1, 0, 1)
|
|
var ch store.Channel
|
|
db.First(&ch, id)
|
|
s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890"))
|
|
release, ok := s.TryAcquire(&ch)
|
|
if !ok {
|
|
t.Fatal("first acquire should succeed")
|
|
}
|
|
if _, ok := s.TryAcquire(&ch); ok {
|
|
t.Fatal("second acquire should fail (capacity 1)")
|
|
}
|
|
release()
|
|
if _, ok := s.TryAcquire(&ch); !ok {
|
|
t.Fatal("acquire after release should succeed")
|
|
}
|
|
}
|
|
|
|
func TestHealthCooldownAndRecovery(t *testing.T) {
|
|
db := newTestDB(t)
|
|
id := seedChannel(t, db, "h", 1, 0, 4)
|
|
var ch store.Channel
|
|
db.First(&ch, id)
|
|
|
|
m := NewHealthMonitor(db, crypto.NewEncryptor("test-master-key-1234567890"), HealthConfig{
|
|
Interval: time.Minute, FailThreshold: 2, Timeout: time.Second,
|
|
})
|
|
|
|
// 两次失败 → cooldown
|
|
m.record(&ch, false)
|
|
if ch.HealthStatus != store.ChannelHealthHealthy {
|
|
t.Fatal("one failure should stay healthy")
|
|
}
|
|
m.record(&ch, false)
|
|
if ch.HealthStatus != store.ChannelHealthCooldown {
|
|
t.Fatalf("two failures should cooldown, got %s", ch.HealthStatus)
|
|
}
|
|
// 恢复
|
|
m.record(&ch, true)
|
|
if ch.HealthStatus != store.ChannelHealthHealthy {
|
|
t.Fatalf("success should recover, got %s", ch.HealthStatus)
|
|
}
|
|
}
|
|
|
|
func TestHealthStartStops(t *testing.T) {
|
|
db := newTestDB(t)
|
|
seedChannel(t, db, "probe", 1, 0, 4)
|
|
m := NewHealthMonitor(db, crypto.NewEncryptor("test-master-key-1234567890"), HealthConfig{
|
|
Interval: time.Hour, FailThreshold: 2, Timeout: time.Second,
|
|
})
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
m.Start(ctx)
|
|
cancel() // 不应 panic
|
|
}
|
|
|
|
// 模型被渠道绑定后只允许走绑定渠道；绑定渠道不可用时返回空候选（上层 503），
|
|
// 不能退化到"不认识这个模型"的全局渠道。
|
|
func TestBoundModelDoesNotFallBackToGlobal(t *testing.T) {
|
|
db := newTestDB(t)
|
|
bound := seedChannel(t, db, "bound", 1, 0, 4)
|
|
seedChannel(t, db, "other", 1, 1, 4)
|
|
m := store.Model{Name: "gpt-m", Enabled: true}
|
|
if err := db.Create(&m).Error; err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := db.Create(&store.ChannelModelBinding{ChannelID: bound, ModelID: m.ID, UpstreamModel: "gpt-m"}).Error; err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890"))
|
|
cands := s.Candidates("gpt-m")
|
|
if len(cands) != 1 || cands[0].Channel.Name != "bound" {
|
|
t.Fatalf("绑定模型应只命中 bound，实际 %+v", cands)
|
|
}
|
|
if cands[0].UpstreamModel != "gpt-m" {
|
|
t.Fatalf("映射名丢失：%s", cands[0].UpstreamModel)
|
|
}
|
|
|
|
db.Model(&store.Channel{}).Where("id = ?", bound).Update("enabled", false)
|
|
s.InvalidateCatalog() // 管理端停用渠道会失效缓存；这里直接写库，手动失效
|
|
if got := s.Candidates("gpt-m"); len(got) != 0 {
|
|
t.Fatalf("绑定渠道停用后应为空候选，实际 %d", len(got))
|
|
}
|
|
// 完全未绑定的模型仍然走全局
|
|
if got := s.Candidates("unmanaged"); len(got) != 1 || got[0].Channel.Name != "other" {
|
|
t.Fatalf("未纳管模型应退回全局渠道，实际 %+v", got)
|
|
}
|
|
}
|
|
|
|
// 运行期冷却：连续失败才踢，成功即清零，冷却到期自动恢复。
|
|
func TestRuntimeCooldownThresholdAndExpiry(t *testing.T) {
|
|
db := newTestDB(t)
|
|
id := seedChannel(t, db, "r", 1, 0, 4)
|
|
var ch store.Channel
|
|
db.First(&ch, id)
|
|
|
|
s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890"))
|
|
s.cooldownFor = 60 * time.Millisecond
|
|
// 默认阈值 3：前两次失败不该踢，第三次才进冷却
|
|
|
|
s.ReportFailure(&ch, FailNet, 0)
|
|
s.ReportFailure(&ch, FailNet, 0)
|
|
if s.InCooldown(id) {
|
|
t.Fatal("未达阈值（3 次连续）不应进冷却")
|
|
}
|
|
s.ReportFailure(&ch, FailNet, 0)
|
|
if !s.InCooldown(id) {
|
|
t.Fatal("连续达阈值应进冷却")
|
|
}
|
|
if got := s.Candidates(""); len(got) != 0 {
|
|
t.Fatalf("冷却渠道不应进候选，实际 %d", len(got))
|
|
}
|
|
if got := s.Candidates("whatever"); len(got) != 0 {
|
|
t.Fatalf("全局候选也应过滤冷却渠道，实际 %d", len(got))
|
|
}
|
|
|
|
time.Sleep(80 * time.Millisecond)
|
|
if s.InCooldown(id) {
|
|
t.Fatal("冷却到期应恢复可用")
|
|
}
|
|
if got := s.Candidates(""); len(got) != 1 {
|
|
t.Fatalf("恢复后应重新进候选，实际 %d", len(got))
|
|
}
|
|
|
|
s.ReportFailure(&ch, FailNet, 0)
|
|
s.ReportSuccess(&ch)
|
|
s.ReportFailure(&ch, FailNet, 0)
|
|
if s.InCooldown(id) {
|
|
t.Fatal("非连续失败不应触发冷却")
|
|
}
|
|
}
|
|
|
|
// 管理后台可见性：RuntimeState 报累计失败与冷却剩余，ClearRuntime 立即清零。
|
|
func TestRuntimeStateAndClear(t *testing.T) {
|
|
db := newTestDB(t)
|
|
id := seedChannel(t, db, "s", 1, 0, 4)
|
|
var ch store.Channel
|
|
db.First(&ch, id)
|
|
|
|
s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890"))
|
|
if fails, rem, _ := s.RuntimeState(id); fails != 0 || rem != 0 {
|
|
t.Fatalf("初始状态 = %d/%s，期望 0/0", fails, rem)
|
|
}
|
|
s.ReportFailure(&ch, FailNet, 0)
|
|
s.ReportFailure(&ch, FailNet, 0)
|
|
if fails, rem, _ := s.RuntimeState(id); fails != 2 || rem != 0 {
|
|
t.Fatalf("两次失败后 = %d/%s，期望 2/0", fails, rem)
|
|
}
|
|
s.ReportFailure(&ch, FailNet, 0)
|
|
fails, rem, _ := s.RuntimeState(id)
|
|
if fails != 0 || rem <= 0 || rem > s.cooldownFor {
|
|
t.Fatalf("冷却中 = %d/%s，期望 0/(0,%s]", fails, rem, s.cooldownFor)
|
|
}
|
|
if got := len(s.Candidates("")); got != 0 {
|
|
t.Fatalf("冷却渠道不应进候选，实际 %d", got)
|
|
}
|
|
|
|
s.ClearRuntime(id)
|
|
if fails, rem, _ := s.RuntimeState(id); fails != 0 || rem != 0 {
|
|
t.Fatalf("清除后 = %d/%s，期望 0/0", fails, rem)
|
|
}
|
|
if s.InCooldown(id) {
|
|
t.Fatal("清除后不应仍在冷却")
|
|
}
|
|
if got := len(s.Candidates("")); got != 1 {
|
|
t.Fatalf("恢复后应重新进候选，实际 %d", got)
|
|
}
|
|
}
|
|
|
|
// 上游 429 带的 Retry-After 可以拉长冷却，但要有上限，不能被钉死几小时。
|
|
func TestReportFailureCooldownOverrideAndClamp(t *testing.T) {
|
|
db := newTestDB(t)
|
|
id := seedChannel(t, db, "ra", 1, 0, 4)
|
|
var ch store.Channel
|
|
db.First(&ch, id)
|
|
|
|
s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890"))
|
|
s.cooldownFor = 30 * time.Second
|
|
|
|
for i := 0; i < 3; i++ {
|
|
s.ReportFailure(&ch, FailRate, 2*time.Minute)
|
|
}
|
|
if _, rem, _ := s.RuntimeState(id); rem < 115*time.Second || rem > 2*time.Minute {
|
|
t.Fatalf("应按 Retry-After 冷却 2min，实际剩余 %s", rem)
|
|
}
|
|
|
|
s.ClearRuntime(id)
|
|
for i := 0; i < 3; i++ {
|
|
s.ReportFailure(&ch, FailRate, time.Hour)
|
|
}
|
|
if _, rem, _ := s.RuntimeState(id); rem > maxCooldown {
|
|
t.Fatalf("冷却应夹到上限 %s，实际 %s", maxCooldown, rem)
|
|
}
|
|
|
|
// 比默认值短的 Retry-After 不缩短冷却（否则 1s 就把坏渠道放回轮转）
|
|
s.ClearRuntime(id)
|
|
for i := 0; i < 3; i++ {
|
|
s.ReportFailure(&ch, FailRate, 5*time.Second)
|
|
}
|
|
if _, rem, _ := s.RuntimeState(id); rem < 25*time.Second {
|
|
t.Fatalf("短 Retry-After 不应缩短默认冷却，实际 %s", rem)
|
|
}
|
|
}
|
|
|
|
// Pick 只在最高优先级层内加权随机；低优先级渠道即使权重再大也不该被首选打中。
|
|
func TestPickStaysInTopPriorityTier(t *testing.T) {
|
|
db := newTestDB(t)
|
|
seedChannel(t, db, "low-pri", 100000, 5, 4) // 低优先级 + 巨大权重
|
|
seedChannel(t, db, "top-pri", 1, 0, 4)
|
|
|
|
s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890"))
|
|
cands := s.Candidates("")
|
|
if len(cands) != 2 {
|
|
t.Fatalf("candidates = %d, want 2", len(cands))
|
|
}
|
|
for i := 0; i < 50; i++ {
|
|
ch := s.Pick(cands)
|
|
if ch.Name != "top-pri" {
|
|
t.Fatalf("pick %d: 应始终选最高优先级渠道，实际 %s", i, ch.Name)
|
|
}
|
|
}
|
|
}
|
|
|
|
// 绑定级权重与渠道级权重相乘：默认绑定权重 1 时等价纯渠道权重，显式加权可单独放大。
|
|
func TestPickBindingWeightMultiplier(t *testing.T) {
|
|
db := newTestDB(t)
|
|
seedChannel(t, db, "plain", 1, 0, 4)
|
|
seedChannel(t, db, "boosted", 1, 0, 4)
|
|
|
|
var m store.Model
|
|
m = store.Model{Name: "mm", Enabled: true}
|
|
db.Create(&m)
|
|
var chs []store.Channel
|
|
db.Order("id").Find(&chs)
|
|
db.Create(&store.ChannelModelBinding{ChannelID: chs[0].ID, ModelID: m.ID, UpstreamModel: "mm"}) // 默认 weight=1
|
|
db.Create(&store.ChannelModelBinding{ChannelID: chs[1].ID, ModelID: m.ID, UpstreamModel: "mm", Weight: 100}) // 显式放大
|
|
|
|
s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890"))
|
|
cands := s.Candidates("mm")
|
|
if len(cands) != 2 {
|
|
t.Fatalf("candidates = %d, want 2", len(cands))
|
|
}
|
|
boosted := 0
|
|
for i := 0; i < 200; i++ {
|
|
if ch := s.Pick(cands); ch.Name == "boosted" {
|
|
boosted++
|
|
}
|
|
}
|
|
if boosted < 150 || boosted > 200 {
|
|
t.Fatalf("权重 100:1 下 boosted 命中 %d/200，应接近 200", boosted)
|
|
}
|
|
}
|
|
|
|
// MaxConcurrency 改配置后在旧槽排空时热生效，无需重启。
|
|
func TestTryAcquireCapacityChange(t *testing.T) {
|
|
db := newTestDB(t)
|
|
id := seedChannel(t, db, "c", 1, 0, 1)
|
|
var ch store.Channel
|
|
db.First(&ch, id)
|
|
s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890"))
|
|
|
|
release, ok := s.TryAcquire(&ch)
|
|
if !ok {
|
|
t.Fatal("容量 1 首次获取应成功")
|
|
}
|
|
// 容量 1 → 2：旧槽被占用，本次应沿用旧容量（拿不到）
|
|
ch.MaxConcurrency = 2
|
|
if _, ok := s.TryAcquire(&ch); ok {
|
|
t.Fatal("旧槽占用期间不应扩容")
|
|
}
|
|
release()
|
|
// 排空后再取：新容量生效，连取 2 个都成功
|
|
r1, ok1 := s.TryAcquire(&ch)
|
|
r2, ok2 := s.TryAcquire(&ch)
|
|
if !ok1 || !ok2 {
|
|
t.Fatal("排空后容量 2 应允许两个并发")
|
|
}
|
|
r1()
|
|
r2()
|
|
}
|
|
|
|
// 失败分类冷却：配额类第一次失败即长冷却；普通类指数退避；成功清零退避计数。
|
|
func TestFailKindCooldownPolicy(t *testing.T) {
|
|
db := newTestDB(t)
|
|
id := seedChannel(t, db, "kind", 1, 0, 4)
|
|
var ch store.Channel
|
|
db.First(&ch, id)
|
|
|
|
s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890"))
|
|
s.cooldownFor = 30 * time.Second
|
|
|
|
// 配额耗尽：不攒阈值，第一次失败立即冷却 quotaCooldown
|
|
s.ReportFailure(&ch, FailQuota, 0)
|
|
if !s.InCooldown(id) {
|
|
t.Fatal("配额类失败应立即进冷却")
|
|
}
|
|
if _, rem, reason := s.RuntimeState(id); reason != FailQuota.String() || rem <= 0 {
|
|
t.Fatalf("冷却原因应为 quota_exhausted 且剩余>0，实际 %s/%s", reason, rem)
|
|
}
|
|
s.ClearRuntime(id)
|
|
|
|
// 指数退避：连续进冷却时长翻倍，封顶 maxCooldown
|
|
s.ReportFailure(&ch, FailServer, 0)
|
|
s.ReportFailure(&ch, FailServer, 0)
|
|
s.ReportFailure(&ch, FailServer, 0) // 第 1 次冷却：30s
|
|
s.mu.Lock() // 只看 trips，把 until 抹掉以便观察下一次退避
|
|
s.runt[id].until = time.Time{}
|
|
trips1 := s.runt[id].trips
|
|
s.mu.Unlock()
|
|
if trips1 != 1 {
|
|
t.Fatalf("第 1 次冷却后 trips=1，实际 %d", trips1)
|
|
}
|
|
s.ReportFailure(&ch, FailServer, 0)
|
|
s.ReportFailure(&ch, FailServer, 0)
|
|
s.ReportFailure(&ch, FailServer, 0) // 第 2 次冷却：60s
|
|
s.mu.Lock()
|
|
d := s.runt[id].until
|
|
s.mu.Unlock()
|
|
if time.Until(d) <= 45*time.Second || time.Until(d) > 61*time.Second {
|
|
t.Fatalf("第 2 次冷却应为 60s 左右，实际 %s", time.Until(d))
|
|
}
|
|
|
|
// Retry-After 比退避值长时覆盖（夹上限）
|
|
s.ClearRuntime(id)
|
|
s.ReportFailure(&ch, FailRate, 2*time.Minute)
|
|
s.ReportFailure(&ch, FailRate, 2*time.Minute)
|
|
s.ReportFailure(&ch, FailRate, 2*time.Minute)
|
|
if _, rem, _ := s.RuntimeState(id); rem <= 0 {
|
|
t.Fatal("Retry-After 覆盖后应有冷却剩余")
|
|
}
|
|
|
|
// 一次成功清零退避计数
|
|
s.ClearRuntime(id)
|
|
s.ReportFailure(&ch, FailNet, 0)
|
|
s.ReportFailure(&ch, FailNet, 0)
|
|
s.ReportFailure(&ch, FailNet, 0) // trips=1
|
|
s.ReportSuccess(&ch)
|
|
s.mu.Lock()
|
|
rt := s.runt[id]
|
|
fails, trips := 0, 0
|
|
if rt != nil {
|
|
fails, trips = rt.fails, rt.trips
|
|
}
|
|
s.mu.Unlock()
|
|
if fails != 0 || trips != 0 {
|
|
t.Fatalf("成功后失败计数与退避计数都应清零，实际 fails=%d trips=%d", fails, trips)
|
|
}
|
|
}
|
|
|
|
// 短 Retry-After 不缩短默认冷却（1s 不能把坏渠道立刻放回轮转）。
|
|
func TestShortRetryAfterDoesNotShortenCooldown(t *testing.T) {
|
|
db := newTestDB(t)
|
|
id := seedChannel(t, db, "short-ra", 1, 0, 4)
|
|
var ch store.Channel
|
|
db.First(&ch, id)
|
|
|
|
s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890"))
|
|
s.cooldownFor = 30 * time.Second
|
|
s.ReportFailure(&ch, FailRate, time.Second)
|
|
s.ReportFailure(&ch, FailRate, time.Second)
|
|
s.ReportFailure(&ch, FailRate, time.Second)
|
|
if _, rem, _ := s.RuntimeState(id); rem < 25*time.Second {
|
|
t.Fatalf("短 Retry-After 不应缩短冷却，实际剩余 %s", rem)
|
|
}
|
|
}
|