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:
1 parent
117e03673b
commit
8e5fe16611
12 files changed
+391
-16
No files matched your search
@@ -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
|
||||
|
||||
@@ -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 错误事件收尾
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"),
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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"`
|
||||
|
||||
Reference in new issue
Block a user