Sakurasan
8e5fe16611
profile（16 并发）显示请求路径的 CPU 有 ~70% 花在鉴权上：每个代理请求都要用 密钥哈希查 api_keys、再用 user_id 查 users，两条随机读。 - 新增鉴权缓存（OT_PROXY_AUTH_CACHE_TTL，默认 2s，0=关闭）：缓存密钥与用户行， 但密钥状态、用户状态、密钥过期仍按当前时间逐请求判定；管理端写操作与用户的 密钥写操作经中间件立即清空缓存。余额随用户行一起缓存，只影响「≤0 就 402」 这道软闸，最长滞后一个 TTL（请求照样计费）。 - last_used_at 节流改由缓存内的 touch 表记账（原先按 DB 行时间戳判断，行被缓存后 时间戳不再变化，会导致每请求都写一次库）。touch 用 sync.Map：早先与缓存共用一把 排他锁，16 并发下互斥饥饿把吞吐从 31k 拖到 3.2k（p99 7ms），换 sync.Map 后消失。 - 去掉 usage_logs 上从未被查询过滤的 key_id / trace_id 索引（没有查询按它们过滤， 每个索引都是每行一次 B-tree 插入）；schema v8 迁移里显式 DROP 旧库上的索引。 实测（16 并发、mock 上游、SQLite）： - 鉴权缓存开：p50 131–177µs、p99 0.6–1.2ms、RPS 56k–81k - 鉴权缓存关：p50 535–631µs、p99 1.6–2.8ms、RPS 21k–27k （4 并发下同样受益：30k→41k） 测试：新增缓存命中零 SQL、状态/过期重判、失效、TTL 到期、缓存关闭、并发七项； go vet、go test ./...、-race（proxy/usage/store/api）全绿。
110 lines
3.6 KiB
Go
110 lines
3.6 KiB
Go
package proxy
|
||
|
||
import (
|
||
"context"
|
||
"io"
|
||
"net/http"
|
||
"net/http/httptest"
|
||
"sync"
|
||
"testing"
|
||
"time"
|
||
|
||
"github.com/openteam/server/internal/channel"
|
||
"gorm.io/gorm"
|
||
gormlogger "gorm.io/gorm/logger"
|
||
)
|
||
|
||
// sqlTracer 记录一次请求路径上打过的 SQL。
|
||
type sqlTracer struct {
|
||
gormlogger.Interface
|
||
mu sync.Mutex
|
||
list []string
|
||
}
|
||
|
||
func (s *sqlTracer) Trace(_ context.Context, _ time.Time, fc func() (string, int64), _ error) {
|
||
sql, _ := fc()
|
||
s.mu.Lock()
|
||
s.list = append(s.list, sql)
|
||
s.mu.Unlock()
|
||
}
|
||
|
||
func (s *sqlTracer) take() []string {
|
||
s.mu.Lock()
|
||
defer s.mu.Unlock()
|
||
out := s.list
|
||
s.list = nil
|
||
return out
|
||
}
|
||
|
||
// 每请求 SQL 预算：热路径上的模型/渠道目录必须走进程内缓存。
|
||
// 原本一次非流式请求要打 10 条（含同一模型查三次），缓存后只剩鉴权两条。
|
||
// 这个上限是防回归的闸门：往请求路径上新增未缓存的查询会立刻让它变红。
|
||
func TestProxyRequestSQLBudget(t *testing.T) {
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
w.Header().Set("Content-Type", "application/json")
|
||
_, _ = io.WriteString(w, `{"id":"x","model":"m1","choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5}}`)
|
||
}))
|
||
defer srv.Close()
|
||
|
||
h := newHarness(t, srv.URL, 0, true, nil)
|
||
tr := &sqlTracer{Interface: gormlogger.Discard}
|
||
traced := h.db.Session(&gorm.Session{Logger: tr})
|
||
h.gw.db = traced
|
||
h.gw.ch = channel.NewService(traced, h.gw.enc)
|
||
|
||
body := `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`
|
||
// 第一次：冷缓存 + last_used_at 首次写入，不计入预算
|
||
if w := h.run(t, nil, body); w.Code != http.StatusOK {
|
||
t.Fatalf("code=%d body=%s", w.Code, w.Body.String())
|
||
}
|
||
tr.take()
|
||
|
||
// 第二次：模型/渠道目录与系统配置都应命中缓存
|
||
if w := h.run(t, nil, body); w.Code != http.StatusOK {
|
||
t.Fatalf("code=%d body=%s", w.Code, w.Body.String())
|
||
}
|
||
sqls := tr.take()
|
||
const budget = 3 // api_keys + users（+ 偶发 last_used_at）
|
||
if len(sqls) > budget {
|
||
t.Fatalf("热路径 SQL 条数 = %d，超过预算 %d：\n%s", len(sqls), budget, joinSQL(sqls))
|
||
}
|
||
t.Logf("第二条请求 SQL = %d 条：\n%s", len(sqls), joinSQL(sqls))
|
||
}
|
||
|
||
func joinSQL(sqls []string) string {
|
||
out := ""
|
||
for i, q := range sqls {
|
||
out += " " + string(rune('1'+i)) + ". " + q + "\n"
|
||
}
|
||
return out
|
||
}
|
||
|
||
// 开启鉴权缓存后，命中期间请求路径上一条 SQL 都不该有（缓存的是密钥+用户行，
|
||
// 状态/过期仍在内存里按当前时间判定）。
|
||
func TestProxyRequestSQLBudgetWithAuthCache(t *testing.T) {
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
w.Header().Set("Content-Type", "application/json")
|
||
_, _ = io.WriteString(w, `{"id":"x","model":"m1","choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5}}`)
|
||
}))
|
||
defer srv.Close()
|
||
|
||
h := newHarness(t, srv.URL, 0, true, nil)
|
||
h.gw.auth = newAuthCache(time.Minute)
|
||
tr := &sqlTracer{Interface: gormlogger.Discard}
|
||
traced := h.db.Session(&gorm.Session{Logger: tr})
|
||
h.gw.db = traced
|
||
h.gw.ch = channel.NewService(traced, h.gw.enc)
|
||
|
||
body := `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`
|
||
if w := h.run(t, nil, body); w.Code != http.StatusOK {
|
||
t.Fatalf("预热请求失败：%d", w.Code)
|
||
}
|
||
tr.take()
|
||
if w := h.run(t, nil, body); w.Code != http.StatusOK {
|
||
t.Fatalf("code=%d", w.Code)
|
||
}
|
||
if sqls := tr.take(); len(sqls) != 0 {
|
||
t.Fatalf("鉴权缓存命中时不应有 SQL，实际 %d 条：\n%s", len(sqls), joinSQL(sqls))
|
||
}
|
||
}
|