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,110 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/openteam/server/internal/store"
|
||||
)
|
||||
|
||||
// 鉴权缓存：每个代理请求都要用密钥哈希查 api_keys、再用 user_id 查 users，
|
||||
// 这两条随机读在 profile 里占请求路径 CPU 的 ~70%。
|
||||
//
|
||||
// 缓存的是「密钥 + 用户」这两行，但**状态与过期每次请求都重新判定**：
|
||||
// 密钥停用、用户停用、密钥过期都按缓存值判断，最长延迟就是 TTL。
|
||||
// 因此 TTL 要短（默认 2s），且管理端/用户的写操作会立即清空缓存。
|
||||
//
|
||||
// 余额也随用户行一起缓存：它只用于「余额 ≤ 0 就 402」这道软闸，
|
||||
// 最长 2s 的滞后不会造成实际损失（请求照样计费，余额允许为负）。
|
||||
const (
|
||||
authCacheMax = 8192
|
||||
authCacheDefault = 2 * time.Second
|
||||
)
|
||||
|
||||
type authEntry struct {
|
||||
key store.APIKey
|
||||
user store.User
|
||||
until time.Time
|
||||
}
|
||||
|
||||
type authCache struct {
|
||||
mu sync.RWMutex
|
||||
m map[string]authEntry
|
||||
ttl time.Duration
|
||||
// touch 每个密钥最近一次写 last_used_at 的时间。行被缓存后它自己是不会变的，
|
||||
// 节流必须按这里的记录判断，否则每请求都会写一次库（反而比不缓存更糟）。
|
||||
//
|
||||
// 用 sync.Map 而不是和 m 共用一把锁：这个表是「每请求读、每分钟写一次」的
|
||||
// 读多写少结构。早先共用一个 Mutex 时，16 并发下每个请求都要抢排他锁，
|
||||
// 直接把吞吐从 31k 拖到 3.2k（Goroutine 互斥饥饿，p99 7ms）。
|
||||
touch sync.Map // keyID(uint64) → time.Time
|
||||
}
|
||||
|
||||
func newAuthCache(ttl time.Duration) *authCache {
|
||||
if ttl < 0 {
|
||||
ttl = 0
|
||||
}
|
||||
return &authCache{m: map[string]authEntry{}, ttl: ttl}
|
||||
}
|
||||
|
||||
// on 缓存是否启用。
|
||||
func (a *authCache) on() bool { return a != nil && a.ttl > 0 }
|
||||
|
||||
func (a *authCache) get(hash string) (store.APIKey, store.User, bool) {
|
||||
if !a.on() {
|
||||
return store.APIKey{}, store.User{}, false
|
||||
}
|
||||
a.mu.RLock()
|
||||
e, ok := a.m[hash]
|
||||
a.mu.RUnlock()
|
||||
if !ok || time.Now().After(e.until) {
|
||||
return store.APIKey{}, store.User{}, false
|
||||
}
|
||||
return e.key, e.user, true
|
||||
}
|
||||
|
||||
func (a *authCache) put(hash string, k store.APIKey, u store.User) {
|
||||
if !a.on() {
|
||||
return
|
||||
}
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
if len(a.m) >= authCacheMax {
|
||||
now := time.Now()
|
||||
for h, e := range a.m {
|
||||
if now.After(e.until) {
|
||||
delete(a.m, h)
|
||||
}
|
||||
}
|
||||
if len(a.m) >= authCacheMax {
|
||||
return // 全是活的：放弃本次缓存，下次请求再查库
|
||||
}
|
||||
}
|
||||
a.m[hash] = authEntry{key: k, user: u, until: time.Now().Add(a.ttl)}
|
||||
}
|
||||
|
||||
// shouldTouch 该密钥现在是否应该写一次 last_used_at（每 every 最多一次）。
|
||||
// 快路径是 sync.Map 的无锁读。
|
||||
func (a *authCache) shouldTouch(keyID uint64, every time.Duration) bool {
|
||||
if a == nil {
|
||||
return true
|
||||
}
|
||||
now := time.Now()
|
||||
if v, ok := a.touch.Load(keyID); ok {
|
||||
if t, ok := v.(time.Time); ok && now.Sub(t) < every {
|
||||
return false
|
||||
}
|
||||
}
|
||||
a.touch.Store(keyID, now)
|
||||
return true
|
||||
}
|
||||
|
||||
// clear 清空缓存（密钥/用户/渠道等发生写操作后调用）。
|
||||
func (a *authCache) clear() {
|
||||
if a == nil {
|
||||
return
|
||||
}
|
||||
a.mu.Lock()
|
||||
a.m = map[string]authEntry{}
|
||||
a.mu.Unlock()
|
||||
}
|
||||
@@ -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
|
||||
@@ -54,9 +54,11 @@ type Gateway struct {
|
||||
affinityTTL time.Duration
|
||||
// 渠道并发满载时每个候选最多排队多久（0=立刻跳过）。
|
||||
channelAcquireWait time.Duration
|
||||
aff *affinityTable
|
||||
tpMu sync.Mutex
|
||||
tps map[time.Duration]*http.Transport
|
||||
// 鉴权缓存（密钥 + 用户行），0=关闭。
|
||||
auth *authCache
|
||||
aff *affinityTable
|
||||
tpMu sync.Mutex
|
||||
tps map[time.Duration]*http.Transport
|
||||
|
||||
policyMu sync.Mutex
|
||||
policy modelPolicy
|
||||
@@ -102,6 +104,9 @@ func (g *Gateway) modelByName(name string) (store.Model, bool) {
|
||||
return e.m, e.ok
|
||||
}
|
||||
|
||||
// InvalidateAuth 清空鉴权缓存（密钥/用户/密钥白名单等写操作后调用）。
|
||||
func (g *Gateway) InvalidateAuth() { g.auth.clear() }
|
||||
|
||||
// InvalidateCatalog 清空模型与渠道目录缓存（管理端改动、健康状态变化后调用）。
|
||||
func (g *Gateway) InvalidateCatalog() {
|
||||
g.modelMu.Lock()
|
||||
@@ -221,6 +226,7 @@ type Options struct {
|
||||
FirstTokenWait time.Duration // 流式放行前等内容的最长时间（0=不限）
|
||||
AffinityTTL time.Duration // 会话亲和 TTL（0=关闭）
|
||||
ChannelAcquireWait time.Duration // 渠道满载排队时长（0=立刻跳过）
|
||||
AuthCacheTTL time.Duration // 鉴权缓存 TTL（0=关闭）
|
||||
}
|
||||
|
||||
func NewGateway(db *gorm.DB, enc *crypto.Encryptor, rec *usage.Recorder, lim *ratelimit.Limiter, o Options) *Gateway {
|
||||
@@ -239,6 +245,7 @@ func NewGateway(db *gorm.DB, enc *crypto.Encryptor, rec *usage.Recorder, lim *ra
|
||||
firstTokenWait: o.FirstTokenWait,
|
||||
affinityTTL: o.AffinityTTL,
|
||||
channelAcquireWait: o.ChannelAcquireWait,
|
||||
auth: newAuthCache(o.AuthCacheTTL),
|
||||
aff: newAffinityTable(),
|
||||
modelCat: map[string]modelEntry{},
|
||||
tps: map[time.Duration]*http.Transport{},
|
||||
@@ -317,14 +324,28 @@ func (g *Gateway) Auth(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
hash := apikey.Hash(key)
|
||||
var k store.APIKey
|
||||
if err := g.db.Where("key_hash = ? AND status = ?", hash, store.KeyStatusActive).First(&k).Error; err != nil {
|
||||
// 命中缓存时省掉两条随机读；状态/过期仍在下面按当前时间判定
|
||||
k, u, ok := g.auth.get(hash)
|
||||
if !ok {
|
||||
if err := g.db.Where("key_hash = ? AND status = ?", hash, store.KeyStatusActive).First(&k).Error; err != nil {
|
||||
apiError(c, http.StatusUnauthorized, "invalid_api_key", "Invalid API key")
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
if err := g.db.First(&u, k.UserID).Error; err != nil {
|
||||
apiError(c, http.StatusForbidden, "user_disabled", "User account is disabled")
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
g.auth.put(hash, k, u)
|
||||
}
|
||||
// 缓存值可能是 TTL 内的旧状态：每次请求都按当前时间重判
|
||||
if k.Status != store.KeyStatusActive {
|
||||
apiError(c, http.StatusUnauthorized, "invalid_api_key", "Invalid API key")
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
var u store.User
|
||||
if err := g.db.First(&u, k.UserID).Error; err != nil || u.Status != store.UserStatusActive {
|
||||
if u.Status != store.UserStatusActive {
|
||||
apiError(c, http.StatusForbidden, "user_disabled", "User account is disabled")
|
||||
c.Abort()
|
||||
return
|
||||
@@ -359,8 +380,9 @@ func (g *Gateway) Auth(c *gin.Context) {
|
||||
c.Set(CtxKeyID, k.ID)
|
||||
c.Set(CtxKeyModel, k.AllowedModels)
|
||||
c.Set(CtxTrace, newTraceID())
|
||||
// last_used_at 按分钟节流：纯展示字段，不值得每请求一次 DB 写（高 QPS 下是可观开销）
|
||||
if k.LastUsedAt == nil || time.Since(*k.LastUsedAt) >= time.Minute {
|
||||
// last_used_at 按分钟节流：纯展示字段，不值得每请求一次 DB 写（高 QPS 下是可观开销）。
|
||||
// 密钥行可能来自缓存（时间戳不会自己变），所以节流按缓存里的 touch 记录判断。
|
||||
if g.auth.shouldTouch(k.ID, time.Minute) {
|
||||
g.db.Model(&store.APIKey{}).Where("id = ?", k.ID).Update("last_used_at", time.Now())
|
||||
}
|
||||
c.Next()
|
||||
|
||||
@@ -78,3 +78,32 @@ func joinSQL(sqls []string) string {
|
||||
}
|
||||
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))
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user