diff --git a/.env.example b/.env.example index 161ff41..7c3041d 100644 --- a/.env.example +++ b/.env.example @@ -50,6 +50,10 @@ OT_PROXY_AFFINITY_TTL=5m # 渠道并发满载时每个候选渠道最多排队多久（0=立刻跳过）。上游 LLM 请求本来就慢， # 短暂排队远好于把 502 直接甩给 agent；排队仍拿不到槽时返回 503 + Retry-After。 OT_PROXY_CHANNEL_ACQUIRE_WAIT=150ms +# 鉴权缓存 TTL（0=关闭，默认 2s）。命中时省掉每个请求的两条随机读（api_keys/users）， +# 实测 p50 从 535µs 降到 131µs。代价：密钥/用户被停用后最多延迟这个时间才生效 +# （管理端与用户的写操作会立即清空缓存，所以实际窗口通常远小于 TTL）。 +OT_PROXY_AUTH_CACHE_TTL=2s # 渠道健康检查 OT_PROXY_HEALTH_INTERVAL=60s diff --git a/README.md b/README.md index 00ccae5..d44d545 100644 --- a/README.md +++ b/README.md @@ -13,7 +13,7 @@ - `GET /v1/models` — 可用模型列表：只列调用方真能用的模型（叠加用户与密钥白名单），带上下文长度与能力位；Anthropic 客户端（带 `x-api-key`/`anthropic-version`）返回 Anthropic 形状 - **三协议互转**：客户端协议 × 渠道协议不匹配时自动转换（如 Chat 调用 Claude、Messages 调用 OpenAI、Responses 调用 Claude），流式逐事件转换；协议匹配时直通 - 错误按客户端协议返回（OpenAI 格式 / Anthropic 格式） -- **热路径开销**：模型/渠道目录进程内缓存（每请求 SQL 从 10 条降到 2 条，管理端写操作与健康状态变化即时失效）、流式正文按字节快路径提取、请求体定点改写（模型映射不再整包 unmarshal/marshal）、tiktoken 编码器复用、会话亲和表 16 片无锁竞争（键为 64 位摘要、零分配）、访问日志走 64 KiB 缓冲（不再每请求一次 stderr 系统调用）；记账落库按驱动自适应（SQLite 单 worker + batch 64，并启用 `_txlock=immediate` 与 `synchronous=NORMAL`，实测 13k → 30k rows/s 且抖动从 ±50% 收敛到 <1%；PostgreSQL 4 个分片 worker） +- **热路径开销**：鉴权缓存（`OT_PROXY_AUTH_CACHE_TTL`，默认 2s，缓存密钥+用户行，状态与过期仍按当前时间判定，写操作即时失效）——开启后命中请求**零 SQL**，16 并发实测 p50 535µs→131µs、RPS 21k→56–81k；模型/渠道目录进程内缓存（每请求 SQL 从 10 条降到 2 条，管理端写操作与健康状态变化即时失效）、流式正文按字节快路径提取、请求体定点改写（模型映射不再整包 unmarshal/marshal）、tiktoken 编码器复用、会话亲和表 16 片无锁竞争（键为 64 位摘要、零分配）、访问日志走 64 KiB 缓冲（不再每请求一次 stderr 系统调用）；记账落库按驱动自适应（SQLite 单 worker + batch 64，并启用 `_txlock=immediate` 与 `synchronous=NORMAL`，实测 13k → 30k rows/s 且抖动从 ±50% 收敛到 <1%；PostgreSQL 4 个分片 worker） - **渠道体系**：按模型绑定选渠道 + 加权负载均衡；模型已有绑定时只用绑定渠道（都不可用则 503，不会把请求发给不认识它的渠道）；每渠道并发信号量（满载排队 `OT_PROXY_CHANNEL_ACQUIRE_WAIT`，默认 150ms，仍满则 503 + Retry-After 而不是 502）；后台健康检查（探测 `/models`，连续失败进 cooldown、恢复放回）+ 运行期健康（按失败原因分类冷却：网络错误/5xx 连续 3 次进冷却且指数退避，余额/配额类第一次即冷却 30 分钟，成功清零退避计数）；可安全重试的失败自动故障转移（网络错误/429/5xx/超时/流式首个内容事件之前的任何失败）；**会话亲和**（`OT_PROXY_AFFINITY_TTL`，默认 5m）：同一会话粘住上次应答的渠道，命中上游提示词缓存而不是全价重算，冷却/下线自动回退加权随机 - **上游超时**：渠道 `timeout_ms` 生效——非流式为整段请求上限，流式仅约束首包（正文不限时长）；上游请求绑定客户端 context，断连即取消 - **故障转移与流式细节**（借鉴 magpie 的 holdWriter）：流式请求在**首个内容事件**到达前由网关缓冲——上游的错误事件、空流、断连、协议形态不符（要流给了整包，didNotStream）都在放行前换渠道消化，客户端无感知；`OT_PROXY_FIRST_TOKEN_WAIT` 限定放行前等内容的最长时间（0=不限）；429 不一定是限流——错误体带余额/配额特征时按配额类长冷却；上游 `Retry-After` 会延长该渠道冷却（夹在 10 分钟内）；重试受 `OT_PROXY_FAILOVER_BUDGET` 时间预算约束（首个渠道不受限，避免坏渠道把客户端拖到 N×timeout）；流式空闲超过 `OT_PROXY_STREAM_HEARTBEAT` 补一行 SSE 注释心跳（扣住期间也照发），防止 nginx/Cloudflare 的空闲超时掐掉长思考；响应头已被心跳先行写出时，全部候选失败以 SSE 错误事件收尾 diff --git a/server/internal/api/admin.go b/server/internal/api/admin.go index fd3e797..a7c5080 100644 --- a/server/internal/api/admin.go +++ b/server/internal/api/admin.go @@ -267,6 +267,19 @@ func (h *Handler) InvalidateCatalogOnWrite(c *gin.Context) { case http.MethodPost, http.MethodPut, http.MethodPatch, http.MethodDelete: if h.gw != nil && c.Writer.Status() < 400 { h.gw.InvalidateCatalog() + h.gw.InvalidateAuth() // 密钥白名单/用户状态/余额可能刚被改 + } + } +} + +// InvalidateAuthOnWrite 用户侧写操作（新建/改/删密钥）成功后清空网关鉴权缓存， +// 让被吊销的密钥立刻失效，而不是等 TTL 到期。 +func (h *Handler) InvalidateAuthOnWrite(c *gin.Context) { + c.Next() + switch c.Request.Method { + case http.MethodPost, http.MethodPut, http.MethodPatch, http.MethodDelete: + if h.gw != nil && c.Writer.Status() < 400 { + h.gw.InvalidateAuth() } } } diff --git a/server/internal/api/router.go b/server/internal/api/router.go index 547fd2a..edb43e8 100644 --- a/server/internal/api/router.go +++ b/server/internal/api/router.go @@ -94,7 +94,7 @@ func NewRouter(a *app.App, gw *proxy.Gateway) *gin.Engine { webauthn.DELETE("/passkeys/:id", middleware.SessionAuth(a), h.PasskeyDelete) } - user := api.Group("", middleware.SessionAuth(a)) + user := api.Group("", middleware.SessionAuth(a), h.InvalidateAuthOnWrite) { user.GET("/user/profile", h.UserProfile) user.GET("/user/balance", h.UserBalance) diff --git a/server/internal/cli/serve.go b/server/internal/cli/serve.go index 4b1e689..eb26ae9 100644 --- a/server/internal/cli/serve.go +++ b/server/internal/cli/serve.go @@ -41,6 +41,7 @@ func serve() error { FirstTokenWait: cfg.Proxy.FirstTokenWait, AffinityTTL: cfg.Proxy.AffinityTTL, ChannelAcquireWait: cfg.Proxy.ChannelAcquireWait, + AuthCacheTTL: cfg.Proxy.AuthCacheTTL, }) // 渠道健康状态变化后让目录缓存立刻失效，而不是等 TTL a.Health.OnChange = gw.InvalidateCatalog diff --git a/server/internal/config/config.go b/server/internal/config/config.go index 05b3572..3e5eb5e 100644 --- a/server/internal/config/config.go +++ b/server/internal/config/config.go @@ -84,6 +84,7 @@ type ProxyConfig struct { FirstTokenWait time.Duration // 流式放行前等内容的最长时间，0=不限；超时未出内容换下一个渠道 AffinityTTL time.Duration // 会话亲和 TTL，0=关闭；同一会话粘住上次应答的渠道以命中上游缓存 ChannelAcquireWait time.Duration // 渠道并发满载时每个候选的排队时长，0=立刻跳过 + AuthCacheTTL time.Duration // 鉴权缓存 TTL，0=关闭（每次请求都查库） HealthInterval time.Duration // 渠道健康检查周期 HealthFailThreshold int // 连续失败 N 次进 cooldown LogRaw bool // 记录管理员原始请求体+响应到 usage_logs（调试用，默认关） @@ -154,6 +155,7 @@ func Load() (*Config, error) { v.SetDefault("proxy.first_token_wait", "0s") v.SetDefault("proxy.affinity_ttl", "5m") v.SetDefault("proxy.channel_acquire_wait", "150ms") + v.SetDefault("proxy.auth_cache_ttl", "2s") v.SetDefault("proxy.health_interval", "60s") v.SetDefault("proxy.health_fail_threshold", 2) v.SetDefault("proxy.log_raw", false) @@ -206,6 +208,7 @@ func Load() (*Config, error) { FirstTokenWait: v.GetDuration("proxy.first_token_wait"), AffinityTTL: v.GetDuration("proxy.affinity_ttl"), ChannelAcquireWait: v.GetDuration("proxy.channel_acquire_wait"), + AuthCacheTTL: v.GetDuration("proxy.auth_cache_ttl"), HealthInterval: v.GetDuration("proxy.health_interval"), HealthFailThreshold: v.GetInt("proxy.health_fail_threshold"), LogRaw: v.GetBool("proxy.log_raw"), diff --git a/server/internal/proxy/authcache.go b/server/internal/proxy/authcache.go new file mode 100644 index 0000000..b14e286 --- /dev/null +++ b/server/internal/proxy/authcache.go @@ -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() +} diff --git a/server/internal/proxy/authcache_test.go b/server/internal/proxy/authcache_test.go new file mode 100644 index 0000000..b1b8360 --- /dev/null +++ b/server/internal/proxy/authcache_test.go @@ -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 diff --git a/server/internal/proxy/gateway.go b/server/internal/proxy/gateway.go index 8499c9f..821f5ec 100644 --- a/server/internal/proxy/gateway.go +++ b/server/internal/proxy/gateway.go @@ -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() diff --git a/server/internal/proxy/sqlbudget_test.go b/server/internal/proxy/sqlbudget_test.go index 305605f..87e4ab0 100644 --- a/server/internal/proxy/sqlbudget_test.go +++ b/server/internal/proxy/sqlbudget_test.go @@ -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)) + } +} diff --git a/server/internal/store/db.go b/server/internal/store/db.go index a9655c9..f090841 100644 --- a/server/internal/store/db.go +++ b/server/internal/store/db.go @@ -26,7 +26,8 @@ import ( // v6：修正 v5 复合索引列序——AutoMigrate 按名跳过已存在的索引，旧库上的 // idx_user_created 是 user_id 单列，光改 struct tag 不会重建，需删掉重迁一次。 // v7：usage_logs 加 reasoning_tokens / cache_write_1h_tokens / served_model 三列。 -const schemaVersion = 7 +// v8：去掉 usage_logs 上没人查询的 key_id / trace_id 索引（每行两次多余的 B-tree 插入）。 +const schemaVersion = 8 // Open 打开数据库连接并自动迁移。 // 开发默认 SQLite（自动追加 WAL 与 busy_timeout，见 withSQLitePragmas），生产可切 postgres。 @@ -64,6 +65,7 @@ func Open(driver, dsn string) (*gorm.DB, error) { return nil, err } reshapeUsageLogIndexes(db) + dropUnusedUsageLogIndexes(db) if driver != "postgres" { setSQLiteVersion(db, schemaVersion) } @@ -132,6 +134,16 @@ func withSQLitePragmas(dsn string) string { // configurePool 连接池上限。此前全仓无任何配置：PG 连接数无上限， // SQLite 多连接并发写叠加无 WAL 加剧锁竞争。 +// dropUnusedUsageLogIndexes 删除历史库里已经没用的索引：AutoMigrate 不会自动删索引， +// 去掉 struct tag 只是不再新建，旧库上它们还会继续拖慢每次插入。 +func dropUnusedUsageLogIndexes(db *gorm.DB) { + for _, name := range []string{"idx_usage_logs_key_id", "idx_usage_logs_trace_id"} { + if err := db.Exec("DROP INDEX IF EXISTS " + name).Error; err != nil { + log.Printf("store: drop index %s: %v", name, err) + } + } +} + func configurePool(db *gorm.DB, driver string) { sqlDB, err := db.DB() if err != nil { diff --git a/server/internal/store/models.go b/server/internal/store/models.go index f59c93a..e987592 100644 --- a/server/internal/store/models.go +++ b/server/internal/store/models.go @@ -183,11 +183,13 @@ type ChannelModelBinding struct { // UsageLog 请求级用量明细（PLANNING §6.5） type UsageLog struct { - ID uint64 `gorm:"primaryKey;autoIncrement" json:"id"` - RequestID string `gorm:"size:128" json:"request_id"` // 上游 request id - TraceID string `gorm:"size:64;index" json:"trace_id"` + ID uint64 `gorm:"primaryKey;autoIncrement" json:"id"` + RequestID string `gorm:"size:128" json:"request_id"` // 上游 request id + // TraceID/KeyID 不建索引：没有查询按它们过滤（幂等检查走 balance_logs.ref_id）， + // 每个索引都是每行一次 B-tree 插入，在记账这条最热的写路径上纯属写放大。 + TraceID string `gorm:"size:64" json:"trace_id"` UserID uint64 `gorm:"index:idx_user_created,priority:1;not null" json:"user_id"` - KeyID uint64 `gorm:"index" json:"key_id"` + KeyID uint64 `json:"key_id"` ChannelID uint64 `gorm:"index:idx_channel_created,priority:1" json:"channel_id"` ModelID uint64 `json:"model_id"` ModelName string `gorm:"size:128" json:"model_name"`