perf: 鉴权缓存让命中请求零 SQL，去掉两个无用索引
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）全绿。
This commit is contained in:
Sakurasan committed
1 parent
117e03673b
commit
8e5fe16611
12 files changed
+391
-16
No files matched your search
@@ -0,0 +1,179 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/openteam/server/internal/channel"
|
||||
"github.com/openteam/server/internal/store"
|
||||
"gorm.io/gorm"
|
||||
gormlogger "gorm.io/gorm/logger"
|
||||
)
|
||||
|
||||
// 第二条请求应完全命中鉴权缓存：不再有 api_keys / users 查询。
|
||||
func TestAuthCacheSkipsQueries(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":1,"completion_tokens":1}}`)
|
||||
}))
|
||||
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"}]}`
|
||||
for i := 0; i < 2; i++ {
|
||||
if w := h.run(t, nil, body); w.Code != http.StatusOK {
|
||||
t.Fatalf("第 %d 次 code=%d body=%s", i+1, w.Code, w.Body.String())
|
||||
}
|
||||
tr.take()
|
||||
}
|
||||
// 第三次：缓存已热，不应再有 SQL
|
||||
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("缓存命中后不应再查库，实际 %d 条：%v", len(sqls), sqls)
|
||||
}
|
||||
}
|
||||
|
||||
// 缓存值可能是 TTL 内的旧状态：密钥/用户被停用后，缓存命中也要被拒。
|
||||
func TestAuthCacheRechecksStatus(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = io.WriteString(w, `{"id":"x","choices":[{"message":{"content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
h := newHarness(t, srv.URL, 0, true, nil)
|
||||
h.gw.auth = newAuthCache(time.Minute)
|
||||
body := `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`
|
||||
if w := h.run(t, nil, body); w.Code != http.StatusOK {
|
||||
t.Fatalf("首选请求应成功：%d", w.Code)
|
||||
}
|
||||
// 直接改缓存里的状态（模拟"库里已停用、缓存尚未过期"）
|
||||
h.gw.auth.mu.Lock()
|
||||
for k, e := range h.gw.auth.m {
|
||||
e.user.Status = store.UserStatusDisabled
|
||||
h.gw.auth.m[k] = e
|
||||
}
|
||||
h.gw.auth.mu.Unlock()
|
||||
if w := h.run(t, nil, body); w.Code != http.StatusForbidden {
|
||||
t.Fatalf("缓存里的停用用户也应被拒，实际 %d", w.Code)
|
||||
}
|
||||
}
|
||||
|
||||
// 密钥过期按当前时间判定，不受缓存影响。
|
||||
func TestAuthCacheRespectsKeyExpiry(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = io.WriteString(w, `{"id":"x","choices":[{"message":{"content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
h := newHarness(t, srv.URL, 0, true, nil)
|
||||
h.gw.auth = newAuthCache(time.Minute)
|
||||
body := `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`
|
||||
if w := h.run(t, nil, body); w.Code != http.StatusOK {
|
||||
t.Fatalf("首选请求应成功：%d", w.Code)
|
||||
}
|
||||
past := time.Now().Add(-time.Minute)
|
||||
h.gw.auth.mu.Lock()
|
||||
for k, e := range h.gw.auth.m {
|
||||
e.key.ExpiresAt = &past
|
||||
h.gw.auth.m[k] = e
|
||||
}
|
||||
h.gw.auth.mu.Unlock()
|
||||
if w := h.run(t, nil, body); w.Code != http.StatusUnauthorized {
|
||||
t.Fatalf("过期密钥应 401，实际 %d", w.Code)
|
||||
}
|
||||
}
|
||||
|
||||
// InvalidateAuth 之后必须重新查库（吊销立刻生效）。
|
||||
func TestAuthCacheInvalidate(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = io.WriteString(w, `{"id":"x","choices":[{"message":{"content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
h := newHarness(t, srv.URL, 0, true, nil)
|
||||
h.gw.auth = newAuthCache(time.Minute)
|
||||
body := `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`
|
||||
if w := h.run(t, nil, body); w.Code != http.StatusOK {
|
||||
t.Fatalf("首选请求应成功：%d", w.Code)
|
||||
}
|
||||
// 吊销密钥后清缓存：下一次请求应 401
|
||||
h.db.Model(&store.APIKey{}).Where("id > 0").Update("status", store.KeyStatusRevoked)
|
||||
h.gw.InvalidateAuth()
|
||||
if w := h.run(t, nil, body); w.Code != http.StatusUnauthorized {
|
||||
t.Fatalf("吊销并清缓存后应 401，实际 %d", w.Code)
|
||||
}
|
||||
}
|
||||
|
||||
// TTL 到期后重新查库：直接改库（绕过失效）也能在一个 TTL 内生效。
|
||||
func TestAuthCacheTTLExpiry(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = io.WriteString(w, `{"id":"x","choices":[{"message":{"content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
h := newHarness(t, srv.URL, 0, true, nil)
|
||||
h.gw.auth = newAuthCache(50 * time.Millisecond)
|
||||
body := `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`
|
||||
if w := h.run(t, nil, body); w.Code != http.StatusOK {
|
||||
t.Fatalf("首选请求应成功：%d", w.Code)
|
||||
}
|
||||
h.db.Model(&store.APIKey{}).Where("id > 0").Update("status", store.KeyStatusRevoked)
|
||||
time.Sleep(70 * time.Millisecond)
|
||||
if w := h.run(t, nil, body); w.Code != http.StatusUnauthorized {
|
||||
t.Fatalf("TTL 到期后应重新查库并 401，实际 %d", w.Code)
|
||||
}
|
||||
}
|
||||
|
||||
// 关闭缓存（TTL=0）时行为不变：每次都查库。
|
||||
func TestAuthCacheDisabled(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = io.WriteString(w, `{"id":"x","choices":[{"message":{"content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
h := newHarness(t, srv.URL, 0, true, nil)
|
||||
h.gw.auth = newAuthCache(0)
|
||||
body := `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`
|
||||
if w := h.run(t, nil, body); w.Code != http.StatusOK {
|
||||
t.Fatalf("code=%d", w.Code)
|
||||
}
|
||||
if h.gw.auth.on() {
|
||||
t.Fatal("TTL=0 时缓存应关闭")
|
||||
}
|
||||
}
|
||||
|
||||
// 并发访问（-race 验证）
|
||||
func TestAuthCacheConcurrent(t *testing.T) {
|
||||
c := newAuthCache(time.Minute)
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 8; i++ {
|
||||
wg.Add(1)
|
||||
go func(n int) {
|
||||
defer wg.Done()
|
||||
for j := 0; j < 200; j++ {
|
||||
h := "hash"
|
||||
c.put(h, store.APIKey{ID: uint64(n)}, store.User{ID: uint64(n)})
|
||||
c.get(h)
|
||||
if j%50 == 0 {
|
||||
c.clear()
|
||||
}
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
var _ = context.Background
|
||||
Reference in new issue
Block a user