perf: 目录缓存与流式正文快路径,请求热路径开销减半

按性能分析逐项落地,实测对比(16 并发、mock 上游零延迟、SQLite):

- 每请求 SQL 10 → 2:模型/渠道目录(模型行、绑定、渠道行)进程内缓存 10s,
  运行期状态(冷却、健康)仍每次实时判定;管理端写操作经中间件、
  健康状态变化经回调立即失效。新增 TestProxyRequestSQLBudget 钉住预算。
- 端到端 p50 1.3–1.6ms → 0.66–0.70ms,RPS 3.5–8.9k → 16–18k。
- 流式正文累积:字符串 += 改 strings.Builder(原为 O(n²)),
  逐行 JSON 提取改字节快路径,并加廉价的子串预过滤、拿到输出 usage 后停累积。
  500 块:949µs/7.9MB 分配 → 83µs/136KB。
- 请求体定点改写:模型映射从整包 unmarshal/marshal 改为顶层字段替换
  (100KB:585µs → 29µs),stream_options 注入同理;requestText 改为
  只在「上游没报用量」时按需提取,亲和键改用原始体前缀。
- 渠道满载背压:新增 Acquire 有界排队(OT_PROXY_CHANNEL_ACQUIRE_WAIT,
  默认 150ms);排队仍满时返回 503 + Retry-After 而不是 502。
  实测渠道上限 4、16 并发:无排队 81 次 503;排队 150ms 后 400/400 成功。
- 记账按批聚合:同批同用户一次扣款、流水批量插入、日表按 (user,model,日期)
  聚合,替代每行 5 条语句;幂等仍靠 ref_id 唯一索引兜底。
- gorm 开启 PrepareStmt(SQL 解析此前占 CPU 15%);tiktoken 编码器按模型复用。
- 管理端用量列表不再返回原始请求/响应体(可能几十 MB),
  改为 has_raw 标记 + GET /admin/usage/:id 按需拉取,前端弹窗打开时取。

验证:go vet、go test ./...、go test -race(proxy/channel/usage/tokenizer)、
vue-tsc、eslint、vite build 全绿;本地起 mock 上游实测列表/详情接口与
管理页弹窗按需加载。
This commit is contained in:
Sakurasan committed 2026-10-30 14:14:21 +08:00
1 parent d34d9bd2e6
commit 8c12207478
23 files changed
+1050 -170

No files matched your search

+3
View File
@@ -47,6 +47,9 @@ OT_PROXY_FIRST_TOKEN_WAIT=0s
# 会话亲和 TTL(0=关闭,默认 5m):同一会话粘住上次应答的渠道,命中上游的提示词
# 缓存(整段上下文每轮重发),而不是被加权随机打散后全价重算。
OT_PROXY_AFFINITY_TTL=5m
# 渠道并发满载时每个候选渠道最多排队多久(0=立刻跳过)。上游 LLM 请求本来就慢,
# 短暂排队远好于把 502 直接甩给 agent;排队仍拿不到槽时返回 503 + Retry-After。
OT_PROXY_CHANNEL_ACQUIRE_WAIT=150ms
# 渠道健康检查
OT_PROXY_HEALTH_INTERVAL=60s
+2 -1
View File
@@ -13,7 +13,8 @@
- `GET /v1/models` — 可用模型列表:只列调用方真能用的模型(叠加用户与密钥白名单),带上下文长度与能力位;Anthropic 客户端(带 `x-api-key`/`anthropic-version`)返回 Anthropic 形状
- **三协议互转**:客户端协议 × 渠道协议不匹配时自动转换(如 Chat 调用 Claude、Messages 调用 OpenAI、Responses 调用 Claude),流式逐事件转换;协议匹配时直通
- 错误按客户端协议返回(OpenAI 格式 / Anthropic 格式)
- **渠道体系**:按模型绑定选渠道 + 加权负载均衡;模型已有绑定时只用绑定渠道(都不可用则 503,不会把请求发给不认识它的渠道);每渠道并发信号量(满载溢出);后台健康检查(探测 `/models`,连续失败进 cooldown、恢复放回)+ 运行期健康(按失败原因分类冷却:网络错误/5xx 连续 3 次进冷却且指数退避,余额/配额类第一次即冷却 30 分钟,成功清零退避计数);可安全重试的失败自动故障转移(网络错误/429/5xx/超时/流式首个内容事件之前的任何失败);**会话亲和**(`OT_PROXY_AFFINITY_TTL`,默认 5m):同一会话粘住上次应答的渠道,命中上游提示词缓存而不是全价重算,冷却/下线自动回退加权随机
- **热路径开销**:模型/渠道目录进程内缓存(每请求 SQL 从 10 条降到 2 条,管理端写操作与健康状态变化即时失效)、流式正文按字节快路径提取、请求体定点改写(模型映射不再整包 unmarshal/marshal)、tiktoken 编码器复用
- **渠道体系**:按模型绑定选渠道 + 加权负载均衡;模型已有绑定时只用绑定渠道(都不可用则 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 错误事件收尾
- **渠道观测**:`GET /api/v1/admin/channels/stats?window=5m|1h|6h|24h|7d` 按渠道聚合请求量、成功率、错误码分布与 P50/P95 延迟;渠道页直接展示近况与运行期冷却剩余时间,可一键恢复
+13
View File
@@ -257,3 +257,16 @@ func (h *Handler) AdminPutConfig(c *gin.Context) {
}
resp.OK(c, gin.H{"ok": true})
}
// InvalidateCatalogOnWrite 管理端写操作成功后清空网关的模型/渠道目录缓存。
// 缓存 TTL 只有 10s,但「改完渠道立刻生效」是管理员的合理预期;
// 用中间件统一处理,避免每个写 handler 都要记得调一次。
func (h *Handler) InvalidateCatalogOnWrite(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.InvalidateCatalog()
}
}
}
+37 -1
View File
@@ -2,6 +2,7 @@ package api
import (
"math"
"net/http"
"sort"
"strconv"
"time"
@@ -261,9 +262,44 @@ func (h *Handler) AdminUsage(c *gin.Context) {
"reasoning_tokens": l.ReasoningTokens, "cache_write_1h_tokens": l.CacheWrite1hTokens,
"served_model": l.ServedModel, "cost": l.Cost,
"latency_ms": l.LatencyMS, "status": l.Status, "error_code": l.ErrorCode,
"raw_request": l.RawRequest, "raw_response": l.RawResponse,
// 原始体(可能是整段 SSE,百 KB~MB 级)不进列表:15 行一页就能到几十 MB。
// 只给一个标记,需要时走 /admin/usage/:id 单独取。
"has_raw": l.RawRequest != "" || l.RawResponse != "",
"created_at": l.CreatedAt,
})
}
resp.OK(c, gin.H{"items": out, "total": total, "page": page, "page_size": size})
}
// AdminUsageDetail 单条用量的完整明细(含原始请求/响应体)。
// 列表接口刻意不返回原始体:它们可能是整段 SSE,把列表响应撑到几十 MB 没有意义。
func (h *Handler) AdminUsageDetail(c *gin.Context) {
id, err := strconv.ParseUint(c.Param("id"), 10, 64)
if err != nil {
resp.Fail(c, http.StatusBadRequest, "invalid usage id")
return
}
var l store.UsageLog
if err := h.a.DB.First(&l, id).Error; err != nil {
resp.Fail(c, http.StatusNotFound, "usage log not found")
return
}
var u store.User
h.a.DB.Select("id, username").First(&u, l.UserID)
var ch store.Channel
if l.ChannelID > 0 {
h.a.DB.Select("id, name").First(&ch, l.ChannelID)
}
resp.OK(c, gin.H{
"id": l.ID, "user": u.Username, "user_id": l.UserID, "model": l.ModelName, "protocol": l.Protocol,
"channel_id": l.ChannelID, "channel": ch.Name,
"input_tokens": l.InputTokens, "output_tokens": l.OutputTokens,
"cache_read_tokens": l.CacheReadTokens, "cache_creation_tokens": l.CacheCreationTokens,
"reasoning_tokens": l.ReasoningTokens, "cache_write_1h_tokens": l.CacheWrite1hTokens,
"served_model": l.ServedModel, "cost": l.Cost,
"latency_ms": l.LatencyMS, "status": l.Status, "error_code": l.ErrorCode,
"raw_request": l.RawRequest, "raw_response": l.RawResponse,
"has_raw": l.RawRequest != "" || l.RawResponse != "",
"created_at": l.CreatedAt,
})
}
+2 -1
View File
@@ -108,7 +108,7 @@ func NewRouter(a *app.App, gw *proxy.Gateway) *gin.Engine {
user.DELETE("/keys/:id", h.DeleteKey)
}
admin := api.Group("/admin", middleware.SessionAuth(a), middleware.AdminOnly)
admin := api.Group("/admin", middleware.SessionAuth(a), middleware.AdminOnly, h.InvalidateCatalogOnWrite)
{
// 用户
admin.GET("/users", h.AdminUsers)
@@ -138,6 +138,7 @@ func NewRouter(a *app.App, gw *proxy.Gateway) *gin.Engine {
// 统计与用量
admin.GET("/stats/overview", h.AdminStatsOverview)
admin.GET("/usage", h.AdminUsage)
admin.GET("/usage/:id", h.AdminUsageDetail)
// 配置
admin.GET("/config", h.AdminConfig)
admin.PUT("/config", h.AdminPutConfig)
+130 -25
View File
@@ -3,6 +3,7 @@
package channel
import (
"context"
"crypto/rand"
"errors"
"log"
@@ -61,6 +62,27 @@ type Service struct {
// 运行期冷却参数;NewService 给默认值,测试可覆盖。
failThreshold int
cooldownFor time.Duration
// 目录缓存:模型→候选渠道的静态部分(渠道行、绑定、映射)缓存 catalogTTL,
// 运行期状态(冷却、健康)每次读取时重新判定。管理端写操作与健康状态变化
// 通过 InvalidateCatalog 立即失效。
catMu sync.Mutex
cat map[string]catEntry
catalogTTL time.Duration
}
// catEntry 一次目录查询的结果(未按运行期状态过滤)。
type catEntry struct {
rows []catRow
at time.Time
// bound:该模型有渠道绑定(只用绑定渠道);否则为全局渠道池
bound bool
}
type catRow struct {
ch store.Channel
upstreamModel string
weight int
}
// runtimeHealth 运行期健康:连续可重试失败(网络错误/超时/429/5xx)累计到阈值进入冷却。
@@ -81,6 +103,8 @@ func NewService(db *gorm.DB, enc *crypto.Encryptor) *Service {
runt: map[uint64]*runtimeHealth{},
failThreshold: 3,
cooldownFor: 30 * time.Second,
cat: map[string]catEntry{},
catalogTTL: 10 * time.Second,
}
}
@@ -95,7 +119,55 @@ type Candidate struct {
// model 非空且该模型有渠道绑定时,只返回绑定渠道(携带 upstream_model 映射);
// 绑定渠道全部不可用则返回空(上层报 no_channel),不会退化到不认识该模型的全局渠道。
// 仅当模型完全没有任何绑定时才退回全局渠道。
//
// 目录查询(模型/绑定/渠道)每请求原本要打 3 条 SQL,这里缓存 catalogTTL;
// 冷却与健康状态每次读取时重新判定,所以坏渠道照样立刻摘除。
func (s *Service) Candidates(model string) []Candidate {
rows, bound := s.catalog(model)
out := make([]Candidate, 0, len(rows))
for i := range rows {
ch := rows[i].ch
if !ch.Enabled || ch.HealthStatus != store.ChannelHealthHealthy {
continue // 缓存行里的健康/启用状态可能落后,这里按读到的值再过滤
}
if s.InCooldown(ch.ID) {
continue
}
cand := Candidate{Channel: &ch}
if bound {
cand.UpstreamModel = rows[i].upstreamModel
cand.Weight = rows[i].weight
}
out = append(out, cand)
}
return out
}
// catalog 取该模型的目录行(缓存命中返回副本,避免调用方改写缓存内容)。
// bound 表示这些行来自「该模型的渠道绑定」而不是全局渠道池。
func (s *Service) catalog(model string) (rows []catRow, bound bool) {
if s.catalogTTL > 0 {
s.catMu.Lock()
if e, ok := s.cat[model]; ok && time.Since(e.at) < s.catalogTTL {
s.catMu.Unlock()
return e.rows, e.bound
}
s.catMu.Unlock()
}
rows, bound = s.loadCatalog(model)
if s.catalogTTL > 0 {
s.catMu.Lock()
if s.cat == nil {
s.cat = map[string]catEntry{}
}
s.cat[model] = catEntry{rows: rows, at: time.Now(), bound: bound}
s.catMu.Unlock()
}
return rows, bound
}
// loadCatalog 查询模型/绑定/渠道,组装目录行(不判运行期状态)。
func (s *Service) loadCatalog(model string) (rows []catRow, bound bool) {
if model != "" {
var modelIDs []uint64
var b []store.ChannelModelBinding
@@ -104,28 +176,21 @@ func (s *Service) Candidates(model string) []Candidate {
s.db.Where("model_id IN ?", modelIDs).Find(&b)
}
if len(b) > 0 {
return s.loadBound(b)
return s.boundRows(b), true
}
}
var chs []store.Channel
s.db.Where("enabled = ? AND health_status = ?", true, store.ChannelHealthHealthy).
Order("priority ASC, weight DESC, id ASC").Find(&chs)
out := make([]Candidate, 0, len(chs))
rows = make([]catRow, 0, len(chs))
for i := range chs {
if s.InCooldown(chs[i].ID) {
continue
}
out = append(out, Candidate{Channel: &chs[i]})
rows = append(rows, catRow{ch: chs[i]})
}
return out
return rows, false
}
// loadBound 按绑定顺序加载渠道候选,过滤健康/启用/运行期冷却,携带 upstream_model 映射。
func (s *Service) loadBound(bindings []store.ChannelModelBinding) []Candidate {
if len(bindings) == 0 {
return nil
}
// channel_id -> 绑定(取该渠道对该模型的映射)
// boundRows 按绑定组装目录行(保持 priority ASC 的故障转移顺序)。
func (s *Service) boundRows(bindings []store.ChannelModelBinding) []catRow {
byChannel := map[uint64]store.ChannelModelBinding{}
ids := make([]uint64, 0, len(bindings))
for _, b := range bindings {
@@ -137,20 +202,19 @@ func (s *Service) loadBound(bindings []store.ChannelModelBinding) []Candidate {
var chs []store.Channel
s.db.Where("id IN ? AND enabled = ? AND health_status = ?", ids, true, store.ChannelHealthHealthy).
Order("priority ASC, weight DESC, id ASC").Find(&chs)
out := make([]Candidate, 0, len(chs))
// 按查询排好序的 chs 输出(而不是绑定返回顺序),保证候选保持 priority ASC 的故障转移顺序
rows := make([]catRow, 0, len(chs))
for i := range chs {
ch := &chs[i]
if s.InCooldown(ch.ID) {
continue
}
out = append(out, Candidate{
Channel: ch,
UpstreamModel: byChannel[ch.ID].UpstreamModel,
Weight: byChannel[ch.ID].Weight,
})
b := byChannel[chs[i].ID]
rows = append(rows, catRow{ch: chs[i], upstreamModel: b.UpstreamModel, weight: b.Weight})
}
return out
return rows
}
// InvalidateCatalog 清空目录缓存:管理端改了渠道/绑定/模型,或渠道健康状态变化时调用。
func (s *Service) InvalidateCatalog() {
s.catMu.Lock()
s.cat = map[string]catEntry{}
s.catMu.Unlock()
}
// AvailableModelIDs 返回对外可见的模型 ID:启用的模型且至少绑定到一个启用且健康的渠道。
@@ -246,6 +310,47 @@ func (s *Service) TryAcquire(ch *store.Channel) (func(), bool) {
}
}
// Acquire 获取渠道并发槽;满载时最多等 wait(0=立刻返回 false,等价 TryAcquire)。
// ctx 结束(客户端断开/请求取消)即放弃等待。
//
// 满载时直接跳过会让超出并发上限的请求立刻失败:实测 16 并发、渠道上限 4 时
// 有 18–25% 的请求直接 502。LLM 请求本来就慢,短暂排队远好于把错误返回给 agent
// (一次 502 可能毁掉一整轮工具链)。
func (s *Service) Acquire(ctx context.Context, ch *store.Channel, wait time.Duration) (func(), bool) {
if ch.MaxConcurrency <= 0 {
return func() {}, true
}
s.mu.Lock()
e, ok := s.sems[ch.ID]
if !ok {
e = &semEntry{sem: make(chan struct{}, ch.MaxConcurrency), cap: ch.MaxConcurrency}
s.sems[ch.ID] = e
} else if e.cap != ch.MaxConcurrency && len(e.sem) == 0 {
e = &semEntry{sem: make(chan struct{}, ch.MaxConcurrency), cap: ch.MaxConcurrency}
s.sems[ch.ID] = e
}
sem := e.sem
s.mu.Unlock()
select {
case sem <- struct{}{}:
return func() { <-sem }, true
default:
}
if wait <= 0 {
return nil, false
}
t := time.NewTimer(wait)
defer t.Stop()
select {
case sem <- struct{}{}:
return func() { <-sem }, true
case <-t.C:
return nil, false
case <-ctx.Done():
return nil, false
}
}
// UpstreamKey 解密渠道上游密钥。
func (s *Service) UpstreamKey(ch *store.Channel) (string, error) {
return s.enc.Decrypt(ch.APIKeyEnc)
+1
View File
@@ -154,6 +154,7 @@ func TestBoundModelDoesNotFallBackToGlobal(t *testing.T) {
}
db.Model(&store.Channel{}).Where("id = ?", bound).Update("enabled", false)
s.InvalidateCatalog() // 管理端停用渠道会失效缓存;这里直接写库,手动失效
if got := s.Candidates("gpt-m"); len(got) != 0 {
t.Fatalf("绑定渠道停用后应为空候选,实际 %d", len(got))
}
+5
View File
@@ -27,6 +27,8 @@ type HealthMonitor struct {
hc *http.Client
mu sync.Mutex
fail map[uint64]int // 渠道连续失败次数
// OnChange 健康状态变化后的回调(目录缓存据此立即失效,不必等 TTL)。
OnChange func()
}
func NewHealthMonitor(db *gorm.DB, enc *crypto.Encryptor, cfg HealthConfig) *HealthMonitor {
@@ -127,4 +129,7 @@ func (h *HealthMonitor) record(ch *store.Channel, ok bool) {
func (h *HealthMonitor) setStatus(ch *store.Channel, status string) {
h.db.Model(&store.Channel{}).Where("id = ?", ch.ID).Update("health_status", status)
ch.HealthStatus = status
if h.OnChange != nil {
h.OnChange()
}
}
+11 -8
View File
@@ -32,15 +32,18 @@ func serve() error {
defer a.Shutdown(context.Background())
gw := proxy.NewGateway(a.DB, a.Enc, a.Usage, a.Limit, proxy.Options{
UserRPS: cfg.RateLimit.UserRPS,
LogRaw: cfg.Proxy.LogRaw,
GWTimeout: cfg.Proxy.Timeout,
MaxBodyMB: cfg.Proxy.MaxBodyMB,
FailoverBudget: cfg.Proxy.FailoverBudget,
StreamHeartbeat: cfg.Proxy.StreamHeartbeat,
FirstTokenWait: cfg.Proxy.FirstTokenWait,
AffinityTTL: cfg.Proxy.AffinityTTL,
UserRPS: cfg.RateLimit.UserRPS,
LogRaw: cfg.Proxy.LogRaw,
GWTimeout: cfg.Proxy.Timeout,
MaxBodyMB: cfg.Proxy.MaxBodyMB,
FailoverBudget: cfg.Proxy.FailoverBudget,
StreamHeartbeat: cfg.Proxy.StreamHeartbeat,
FirstTokenWait: cfg.Proxy.FirstTokenWait,
AffinityTTL: cfg.Proxy.AffinityTTL,
ChannelAcquireWait: cfg.Proxy.ChannelAcquireWait,
})
// 渠道健康状态变化后让目录缓存立刻失效,而不是等 TTL
a.Health.OnChange = gw.InvalidateCatalog
router := api.NewRouter(a, gw)
srv := &http.Server{
+3
View File
@@ -83,6 +83,7 @@ type ProxyConfig struct {
StreamHeartbeat time.Duration // 流式空闲多久发一次 SSE 注释心跳,0=关闭
FirstTokenWait time.Duration // 流式放行前等内容的最长时间,0=不限;超时未出内容换下一个渠道
AffinityTTL time.Duration // 会话亲和 TTL,0=关闭;同一会话粘住上次应答的渠道以命中上游缓存
ChannelAcquireWait time.Duration // 渠道并发满载时每个候选的排队时长,0=立刻跳过
HealthInterval time.Duration // 渠道健康检查周期
HealthFailThreshold int // 连续失败 N 次进 cooldown
LogRaw bool // 记录管理员原始请求体+响应到 usage_logs(调试用,默认关)
@@ -152,6 +153,7 @@ func Load() (*Config, error) {
v.SetDefault("proxy.stream_heartbeat", "15s")
v.SetDefault("proxy.first_token_wait", "0s")
v.SetDefault("proxy.affinity_ttl", "5m")
v.SetDefault("proxy.channel_acquire_wait", "150ms")
v.SetDefault("proxy.health_interval", "60s")
v.SetDefault("proxy.health_fail_threshold", 2)
v.SetDefault("proxy.log_raw", false)
@@ -203,6 +205,7 @@ func Load() (*Config, error) {
StreamHeartbeat: v.GetDuration("proxy.stream_heartbeat"),
FirstTokenWait: v.GetDuration("proxy.first_token_wait"),
AffinityTTL: v.GetDuration("proxy.affinity_ttl"),
ChannelAcquireWait: v.GetDuration("proxy.channel_acquire_wait"),
HealthInterval: v.GetDuration("proxy.health_interval"),
HealthFailThreshold: v.GetInt("proxy.health_fail_threshold"),
LogRaw: v.GetBool("proxy.log_raw"),
+46 -7
View File
@@ -2,19 +2,28 @@
// 用于流式中断时对已生成内容做近似计费;未知模型回退 cl100k_base。
package tokenizer
import "github.com/tiktoken-go/tokenizer"
import (
"sync"
"github.com/tiktoken-go/tokenizer"
)
// 编码器缓存:tokenizer.Get/ForModel 每次调用都新建 codec,其中包含 regexp2
// 正则编译(回溯引擎,编译昂贵)与 map 分配。实测短文本 6.9µs、18KB 文本
// 1.05ms,绝大部分是这套固定开销而非真正的分词。
//
// codec 构建后只读(vocabulary/splitRegexp 不再变),regexp2.Regexp 自身带锁,
// 因此缓存实例可并发复用。键是模型名:调用方模型名有限,缓存自然有界。
var codecs sync.Map // string → tokenizer.Codec
// Count 估算文本 token 数;空文本返回 0。
func Count(text, model string) int {
if text == "" {
return 0
}
enc, err := tokenizer.ForModel(tokenizer.Model(model))
if err != nil {
enc, err = tokenizer.Get(tokenizer.Cl100kBase)
if err != nil {
return 0
}
enc := codecFor(model)
if enc == nil {
return 0
}
toks, _, err := enc.Encode(text)
if err != nil {
@@ -22,3 +31,33 @@ func Count(text, model string) int {
}
return len(toks)
}
// codecFor 取(必要时构建并缓存)该模型对应的 codec;未知模型回退 cl100k_base。
func codecFor(model string) tokenizer.Codec {
if v, ok := codecs.Load(model); ok {
return v.(tokenizer.Codec)
}
enc, err := tokenizer.ForModel(tokenizer.Model(model))
if err != nil || enc == nil {
enc = fallback()
}
if enc == nil {
return nil
}
codecs.Store(model, enc)
return enc
}
// fallback 取 cl100k_base 编码器(单独缓存,避免每次未知模型都重建)。
func fallback() tokenizer.Codec {
const key = "\x00cl100k"
if v, ok := codecs.Load(key); ok {
return v.(tokenizer.Codec)
}
enc, err := tokenizer.Get(tokenizer.Cl100kBase)
if err != nil || enc == nil {
return nil
}
codecs.Store(key, enc)
return enc
}
@@ -0,0 +1,51 @@
package tokenizer
import (
"sync"
"testing"
)
// 同一模型连续计数应复用同一个 codec(缓存命中),且结果稳定。
func TestCountCachesCodec(t *testing.T) {
text := "请帮我分析这段代码的性能问题"
first := Count(text, "gpt-4o-mini")
if first == 0 {
t.Fatal("期望非 0 的 token 数")
}
c1 := codecFor("gpt-4o-mini")
c2 := codecFor("gpt-4o-mini")
if c1 != c2 {
t.Fatal("同一模型应复用同一个 codec 实例")
}
if got := Count(text, "gpt-4o-mini"); got != first {
t.Fatalf("结果不稳定:%d vs %d", got, first)
}
}
// 未知模型回退 cl100k_base,且不重复构建。
func TestCountUnknownModelFallsBack(t *testing.T) {
if got := Count("hello world", "some-unknown-model-x"); got == 0 {
t.Fatal("未知模型应回退到 cl100k_base 并给出非 0 结果")
}
if fallback() != fallback() {
t.Fatal("回退编码器应被缓存")
}
}
// 并发调用不得 panic/竞态(-race 下验证)。
func TestCountConcurrent(t *testing.T) {
var wg sync.WaitGroup
for i := 0; i < 16; i++ {
wg.Add(1)
go func() {
defer wg.Done()
for j := 0; j < 50; j++ {
if Count("并发测试文本,检查 codec 复用是否安全。", "gpt-4o-mini") == 0 {
t.Error("期望非 0")
return
}
}
}()
}
wg.Wait()
}
+88 -38
View File
@@ -52,12 +52,63 @@ type Gateway struct {
firstTokenWait time.Duration
// 会话亲和 TTL(0=关闭):同一会话粘住上次应答的渠道,保住上游提示词缓存。
affinityTTL time.Duration
aff *affinityTable
tpMu sync.Mutex
tps map[time.Duration]*http.Transport
// 渠道并发满载时每个候选最多排队多久(0=立刻跳过)。
channelAcquireWait time.Duration
aff *affinityTable
tpMu sync.Mutex
tps map[time.Duration]*http.Transport
policyMu sync.Mutex
policy modelPolicy
// 模型目录缓存:每请求原本要查「是否存在」与「计价」两次同名查询,
// 这里按名称缓存 catalogTTL;管理端写操作立即失效。
modelMu sync.Mutex
modelCat map[string]modelEntry
modelAt time.Time
}
// catalogTTL 目录缓存有效期。运行期状态(冷却、健康)不受它影响,
// 它们每次读取时重新判定;这里只缓存「配置类」数据(模型、渠道、绑定)。
const catalogTTL = 10 * time.Second
type modelEntry struct {
m store.Model
ok bool
}
// modelByName 取模型(带进程内缓存);第二个返回值表示模型存在且启用。
func (g *Gateway) modelByName(name string) (store.Model, bool) {
g.modelMu.Lock()
if g.modelCat != nil && time.Since(g.modelAt) < catalogTTL {
if e, ok := g.modelCat[name]; ok {
g.modelMu.Unlock()
return e.m, e.ok
}
}
g.modelMu.Unlock()
var m store.Model
err := g.db.Where("name = ? AND enabled = ?", name, true).First(&m).Error
e := modelEntry{m: m, ok: err == nil}
g.modelMu.Lock()
if g.modelCat == nil || time.Since(g.modelAt) >= catalogTTL {
g.modelCat = map[string]modelEntry{}
g.modelAt = time.Now()
}
g.modelCat[name] = e
g.modelMu.Unlock()
return e.m, e.ok
}
// InvalidateCatalog 清空模型与渠道目录缓存(管理端改动、健康状态变化后调用)。
func (g *Gateway) InvalidateCatalog() {
g.modelMu.Lock()
g.modelCat = nil
g.modelAt = time.Time{}
g.modelMu.Unlock()
g.ch.InvalidateCatalog()
}
// modelPolicy 全局模型限制策略(来自系统配置,短时缓存)。
@@ -141,9 +192,7 @@ func (g *Gateway) checkModelExists(c *gin.Context, model string) bool {
apiError(c, http.StatusBadRequest, "invalid_request_error", "model is required")
return false
}
var n int64
g.db.Model(&store.Model{}).Where("name = ? AND enabled = ?", model, true).Count(&n)
if n == 0 {
if _, ok := g.modelByName(model); !ok {
apiError(c, http.StatusNotFound, "model_not_found",
fmt.Sprintf("The model %q does not exist or is not enabled", model))
return false
@@ -163,33 +212,36 @@ func contains(list []string, s string) bool {
// Options 网关配置项。用结构体而不是十个位置参数:里面三个 time.Duration、
// 两个 int,传错的编译期发现不了。
type Options struct {
UserRPS int
LogRaw bool
GWTimeout time.Duration // 上游超时兜底(渠道 timeout_ms 优先)
MaxBodyMB int // /v1 请求体上限(MB)
FailoverBudget time.Duration // 故障转移重试的时间预算
StreamHeartbeat time.Duration // 流式空闲心跳间隔
FirstTokenWait time.Duration // 流式放行前等内容的最长时间(0=不限)
AffinityTTL time.Duration // 会话亲和 TTL(0=关闭)
UserRPS int
LogRaw bool
GWTimeout time.Duration // 上游超时兜底(渠道 timeout_ms 优先)
MaxBodyMB int // /v1 请求体上限(MB)
FailoverBudget time.Duration // 故障转移重试的时间预算
StreamHeartbeat time.Duration // 流式空闲心跳间隔
FirstTokenWait time.Duration // 流式放行前等内容的最长时间(0=不限)
AffinityTTL time.Duration // 会话亲和 TTL(0=关闭)
ChannelAcquireWait time.Duration // 渠道满载排队时长(0=立刻跳过)
}
func NewGateway(db *gorm.DB, enc *crypto.Encryptor, rec *usage.Recorder, lim *ratelimit.Limiter, o Options) *Gateway {
return &Gateway{
db: db,
ch: channel.NewService(db, enc),
rec: rec,
enc: enc,
lim: lim,
userRPS: o.UserRPS,
logRaw: o.LogRaw,
gwTimeout: o.GWTimeout,
maxBodyMB: o.MaxBodyMB,
failoverBudget: o.FailoverBudget,
streamHeartbeat: o.StreamHeartbeat,
firstTokenWait: o.FirstTokenWait,
affinityTTL: o.AffinityTTL,
aff: newAffinityTable(),
tps: map[time.Duration]*http.Transport{},
db: db,
ch: channel.NewService(db, enc),
rec: rec,
enc: enc,
lim: lim,
userRPS: o.UserRPS,
logRaw: o.LogRaw,
gwTimeout: o.GWTimeout,
maxBodyMB: o.MaxBodyMB,
failoverBudget: o.FailoverBudget,
streamHeartbeat: o.StreamHeartbeat,
firstTokenWait: o.FirstTokenWait,
affinityTTL: o.AffinityTTL,
channelAcquireWait: o.ChannelAcquireWait,
aff: newAffinityTable(),
modelCat: map[string]modelEntry{},
tps: map[time.Duration]*http.Transport{},
}
}
@@ -493,17 +545,15 @@ func prepareUpstream(ch *store.Channel, clientProto string, body []byte, upstrea
return plan, nil
}
// rewriteModel 改写请求体中的 model 字段(三种协议 model 都在顶层)。
// rewriteModel 改写请求体顶层的 model 字段(模型映射)。
// 定点替换而不是整包 unmarshal/marshal:后者在 100KB 请求体上实测 ~585µs、
// 470KB 分配,而这里只有一次扫描 + 拼接。
func rewriteModel(body []byte, model string) ([]byte, error) {
var m map[string]any
if err := json.Unmarshal(body, &m); err != nil {
return body, nil
out, ok := rewriteModelField(body, model)
if !ok {
return body, nil // 没有顶层 model(非法或空体):原样转发,交给上游判
}
if cur, _ := m["model"].(string); cur == model {
return body, nil
}
m["model"] = model
return json.Marshal(m)
return out, nil
}
var errNoChannel = errors.New("no available channel")
+2
View File
@@ -34,6 +34,7 @@ func addChannel(t *testing.T, h *harness, name, baseURL string, weight int) {
if err := h.db.Create(&store.ChannelModelBinding{ChannelID: ch.ID, ModelID: m.ID, UpstreamModel: "m1"}).Error; err != nil {
t.Fatal(err)
}
h.gw.InvalidateCatalog() // 管理端写操作会失效目录缓存,测试里直接写库要手动失效
}
// pushWeightTo 把指定名字渠道的权重改大,让加权随机几乎必然从它开始。
@@ -42,6 +43,7 @@ func pushWeightTo(t *testing.T, h *harness, name string, weight int) {
if err := h.db.Model(&store.Channel{}).Where("name = ?", name).Update("weight", weight).Error; err != nil {
t.Fatal(err)
}
h.gw.InvalidateCatalog()
}
// 流式容错:渠道1 在首个内容事件之前回错误事件(200+流式形式的错误),应换渠道2
+155
View File
@@ -0,0 +1,155 @@
package proxy
import (
"bytes"
"encoding/json"
)
// 请求体上的定点改写:只动需要改的那一段,不做整包 unmarshal → map → marshal。
// 整包往返在 agent 场景(上下文几十到几百 KB)实测约 0.5–0.6ms 且分配翻倍,
// 而这里只做一次 O(n) 扫描。
// rewriteModelField 改写请求体顶层 model 字段,其余字节原样保留。
// 嵌套出现的 "model"(比如消息内容里引用的 JSON)不会被误改:只认顶层键。
// 找不到顶层 model 或请求体不是对象时返回 (body, false)。
func rewriteModelField(body []byte, model string) ([]byte, bool) {
span, ok := topLevelValueSpan(body, "model")
if !ok {
return body, false
}
quoted, err := json.Marshal(model)
if err != nil {
return body, false
}
out := make([]byte, 0, len(body)-(span[1]-span[0])+len(quoted))
out = append(out, body[:span[0]]...)
out = append(out, quoted...)
out = append(out, body[span[1]:]...)
return out, true
}
// injectStreamOptions 在请求体顶层追加 stream_options,让流式末块带上 usage。
// 只在没有该键时插入;body 必须是 JSON 对象(以 '}' 结尾,允许尾部空白)。
func injectStreamOptions(body []byte, options string) ([]byte, bool) {
trimmed := bytes.TrimRight(body, " \t\r\n")
if len(trimmed) < 2 || trimmed[0] != '{' || trimmed[len(trimmed)-1] != '}' {
return body, false
}
if bytes.Contains(trimmed, []byte(`"stream_options"`)) {
return body, false
}
end := len(trimmed) - 1
out := make([]byte, 0, len(body)+len(options)+16)
out = append(out, trimmed[:end]...)
// 空对象 {} 不能加逗号
if last := bytes.TrimRight(trimmed[:end], " \t\r\n"); len(last) > 1 {
out = append(out, ',')
}
out = append(out, `"stream_options":`...)
out = append(out, options...)
out = append(out, '}')
return out, true
}
// topLevelValueSpan 找顶层键 key 的值在 body 中的字节区间 [start, end)。
// 只遍历顶层(depth==1)的键,跳过字符串字面量与转义。
func topLevelValueSpan(body []byte, key string) ([2]int, bool) {
var zero [2]int
trimmed := bytes.TrimLeft(body, " \t\r\n")
if len(trimmed) == 0 || trimmed[0] != '{' {
return zero, false
}
// 重新对齐到 trimmed 起点后的偏移
base := len(body) - len(trimmed)
keyPat := []byte(`"` + key + `"`)
i := base + 1
depth := 1
for i < len(body) {
switch body[i] {
case '{', '[':
depth++
i++
case '}', ']':
depth--
if depth == 0 {
return zero, false
}
i++
case '"':
start := i
i = skipJSONString(body, i)
// 顶层字符串:可能是"键",也可能是"值"
if depth == 1 && bytes.Equal(body[start:i], keyPat) {
j := i
for j < len(body) && (body[j] == ' ' || body[j] == '\t' || body[j] == '\r' || body[j] == '\n') {
j++
}
if j < len(body) && body[j] == ':' {
j++
for j < len(body) && (body[j] == ' ' || body[j] == '\t' || body[j] == '\r' || body[j] == '\n') {
j++
}
vs := j
ve, ok := jsonValueEnd(body, vs)
if !ok {
return zero, false
}
return [2]int{vs, ve}, true
}
}
default:
i++
}
}
return zero, false
}
// skipJSONString 返回从 body[i]('"')开始的下一个未转义引号之后的位置。
func skipJSONString(body []byte, i int) int {
i++ // 跳过开引号
for i < len(body) {
switch body[i] {
case '\\':
i += 2
case '"':
return i + 1
default:
i++
}
}
return i
}
// jsonValueEnd 返回从 i 开始的一个 JSON 值的结束位置(不含尾部空白)。
func jsonValueEnd(body []byte, i int) (int, bool) {
if i >= len(body) {
return 0, false
}
switch body[i] {
case '"':
return skipJSONString(body, i), true
case '{', '[':
depth := 0
for j := i; j < len(body); j++ {
switch body[j] {
case '"':
j = skipJSONString(body, j) - 1
case '{', '[':
depth++
case '}', ']':
depth--
if depth == 0 {
return j + 1, true
}
}
}
return 0, false
default:
j := i
for j < len(body) && body[j] != ',' && body[j] != '}' && body[j] != ']' &&
body[j] != ' ' && body[j] != '\t' && body[j] != '\r' && body[j] != '\n' {
j++
}
return j, j > i
}
}
+100
View File
@@ -0,0 +1,100 @@
package proxy
import (
"encoding/json"
"strings"
"testing"
)
func TestRewriteModelFieldTopLevelOnly(t *testing.T) {
in := `{"model":"a","messages":[{"role":"user","content":"{\"model\":\"nested\"}"}],"max_tokens":5}`
out, ok := rewriteModelField([]byte(in), "b")
if !ok {
t.Fatal("应改写顶层 model")
}
var m map[string]any
if err := json.Unmarshal(out, &m); err != nil {
t.Fatalf("输出必须是合法 JSON: %v\n%s", err, out)
}
if m["model"] != "b" {
t.Fatalf("顶层 model 未改写:%v", m["model"])
}
// 嵌套的 model 引用必须原样保留
if !strings.Contains(string(out), `\"model\":\"nested\"`) {
t.Fatalf("嵌套内容被改动了:%s", out)
}
if m["max_tokens"].(float64) != 5 {
t.Fatalf("其它字段丢失:%v", out)
}
}
func TestRewriteModelFieldKeepsByteShape(t *testing.T) {
in := "{\n \"model\" : \"old\",\n \"stream\": true\n}"
out, ok := rewriteModelField([]byte(in), "new")
if !ok {
t.Fatal("应改写")
}
want := "{\n \"model\" : \"new\",\n \"stream\": true\n}"
if string(out) != want {
t.Fatalf("应只替换值本身,实际:%q", out)
}
}
func TestRewriteModelFieldMissing(t *testing.T) {
if _, ok := rewriteModelField([]byte(`{"stream":true}`), "x"); ok {
t.Fatal("没有 model 时不应改写")
}
if _, ok := rewriteModelField([]byte(`not json`), "x"); ok {
t.Fatal("非 JSON 不应改写")
}
}
func TestInjectStreamOptions(t *testing.T) {
cases := []struct{ in, want string }{
{`{"model":"m"}`, `{"model":"m","stream_options":{"include_usage":true}}`},
{"{\n \"model\": \"m\"\n}", "{\n \"model\": \"m\"\n,\"stream_options\":{\"include_usage\":true}}"},
{`{}`, `{"stream_options":{"include_usage":true}}`}, // 空对象也能安全追加
}
for _, c := range cases {
out, ok := injectStreamOptions([]byte(c.in), `{"include_usage":true}`)
if !ok {
t.Fatalf("应注入:%s", c.in)
}
if string(out) != c.want {
t.Fatalf("注入结果 = %q,期望 %q", out, c.want)
}
var m map[string]any
if err := json.Unmarshal(out, &m); err != nil {
t.Fatalf("注入后不是合法 JSON: %v (%s)", err, out)
}
if _, ok := m["stream_options"]; !ok {
t.Fatalf("缺少 stream_options: %s", out)
}
}
// 已有 stream_options 时不重复注入
if _, ok := injectStreamOptions([]byte(`{"stream_options":{"include_usage":true}}`), `{}`); ok {
t.Fatal("已有该键时不应再注入")
}
}
// 快路径必须与 JSON 解析结果一致(含转义、非正文事件、三协议形态)。
func TestFastContentTextMatchesJSON(t *testing.T) {
lines := []string{
`data: {"choices":[{"delta":{"content":"hello"}}]}`,
`data: {"choices":[{"delta":{"content":"中文内容"}}]}`,
`data: {"choices":[{"delta":{"content":"带\"引号\"和\\反斜杠"}}]}`,
`data: {"choices":[{"delta":{"role":"assistant"}}]}`,
`data: {"choices":[{"delta":{},"finish_reason":"stop"}]}`,
`data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"块文本"}}`,
`data: {"type":"response.output_text.delta","delta":"响应增量"}`,
`data: {"type":"message_start","message":{"usage":{}}}`,
`data: [DONE]`,
}
for _, line := range lines {
got := fastContentText([]byte(line))
want := sseContentText([]byte(line))
if got != want {
t.Fatalf("不一致\n行: %s\n快路径: %q\nJSON: %q", line, got, want)
}
}
}
+122 -22
View File
@@ -43,8 +43,10 @@ func parseBody(c *gin.Context) (*bodyReq, []byte, error) {
return nil, nil, err
}
c.Request.Body = io.NopCloser(bytes.NewReader(body))
// 中断请求输入文本(token 估算用,仅存文本不立即分词)
c.Set("est_input_text", requestText(body))
// 只留存原始请求体:输入文本提取(requestText)要整包 JSON 解析,
// 而它只在「上游没报用量」时才用得到(实测 100KB 请求体 ~506µs)。
// 亲和键与估算都改成用到时再提取。
c.Set("req_body", body)
br := &bodyReq{}
_ = json.Unmarshal(body, br) // 解析失败按空处理,直通仍可转发
return br, body, nil
@@ -162,6 +164,47 @@ func sseContentText(line []byte) string {
return ""
}
// fastContentText 从一条 SSE 行里取出增量正文,走字节扫描而不是整行 JSON 解析。
// 逐条 JSON 反序列化到 map 大约 1–2µs、二十多次分配,一份长回答几千条事件就是
// 毫秒级纯开销;这里只在三种固定键位上定位字符串值,遇到转义再回退到 JSON 解析。
//
// 覆盖三协议的增量形态:
// - chat:{"choices":[{"delta":{"content":"…"}}]}
// - messages:{"type":"content_block_delta","delta":{"text":"…"}}
// - responses:{"type":"response.output_text.delta","delta":"…"}
func fastContentText(line []byte) string {
for _, key := range []string{`"content":"`, `"text":"`, `"delta":"`} {
idx := bytes.Index(line, []byte(key))
if idx < 0 {
continue
}
start := idx + len(key) // 指向值的第一个字符(引号已含在 key 里)
end := start
escaped := false
for end < len(line) {
switch line[end] {
case '\\':
escaped = true
end += 2
continue
case '"':
if escaped {
// 有转义:这一段交给 JSON 解析,保证还原正确
var out string
if json.Unmarshal(line[idx+len(key)-1:end+1], &out) == nil {
return out
}
return ""
}
return string(line[start:end])
}
end++
}
return "" // 行被截断:留给下一行
}
return ""
}
// retryAfter 解析上游 Retry-After(delta-秒或 HTTP 日期),用于按对端意愿延长渠道冷却。
// 解析不出或为非正返回 0,交给 channel.Service 用默认时长(并夹到上限)。
func retryAfter(h http.Header) time.Duration {
@@ -236,6 +279,9 @@ func (g *Gateway) doProxy(c *gin.Context, cands []channel.Candidate, clientProto
// 最后真正试过的渠道:全部失败时把错误归到它头上,否则 usage_logs 里
// channel_id=0,渠道维度的统计与排障就看不到这类失败。
var lastCh *store.Channel
// saturated:有候选因并发满载被跳过。一次都没打到上游就全被跳过时,
// 那不是「上游坏了」(502),而是「我方满了」(503 + Retry-After)。
saturated := false
startAll := time.Now()
// 从选中的渠道开始遍历,到末尾后再从头遍历到选中渠道之前
for offset := 0; offset < len(cands); offset++ {
@@ -262,9 +308,10 @@ func (g *Gateway) doProxy(c *gin.Context, cands []channel.Candidate, clientProto
lastStatus, lastBody = http.StatusInternalServerError, []byte("conversion error: "+err.Error())
continue
}
release, ok := g.ch.TryAcquire(ch)
release, ok := g.ch.Acquire(c.Request.Context(), ch, g.channelAcquireWait)
if !ok {
continue // 渠道满载,溢出到下一个
saturated = true // 满载:稍等仍拿不到槽,溢出到下一个
continue
}
// 失败/成功计数由 proxyOne 就地报:只有它读得到上游响应头(429 的 Retry-After)
lastCh = ch
@@ -291,6 +338,19 @@ func (g *Gateway) doProxy(c *gin.Context, cands []channel.Candidate, clientProto
}
lastStatus, lastBody = st, b
}
// 全部候选都没打出去、且都是被并发满载挡住的:503 让客户端稍后重试,
// 而不是 502(那意味着上游故障,客户端通常会直接放弃)。
if saturated && lastCh == nil {
c.Header("Retry-After", "1")
msg := "all upstream channels are at capacity, retry shortly"
if sh != nil && sh.headSent() {
_ = sh.writeDirect(streamErrorEvent(clientProto, http.StatusServiceUnavailable, msg))
} else {
apiError(c, http.StatusServiceUnavailable, "channel_saturated", msg)
}
g.recordError(c, nil, nil, now(), "channel_saturated")
return
}
// 全部候选重试性失败
if sh != nil && sh.headSent() {
_ = sh.writeDirect(streamErrorEvent(clientProto, lastStatus, string(lastBody)))
@@ -313,9 +373,15 @@ func (g *Gateway) affinityKeyFor(c *gin.Context) string {
}
mn, _ := c.Get("model_name")
model, _ := mn.(string)
text, _ := c.Get("est_input_text")
firstText, _ := text.(string)
return affinityKey(kID, model, firstText)
// 亲和键取原始请求体的前缀:每轮重发全量上下文时开头是稳定的(模型、系统提示、
// 最早的几条消息),足够区分会话;不必为了它整包解析 JSON。
var prefix string
if v, ok := c.Get("req_body"); ok {
if raw, ok2 := v.([]byte); ok2 {
prefix = string(raw) // affinityKey 内部只取前 512 字节
}
}
return affinityKey(kID, model, prefix)
}
// streamErrorEvent 把错误构造成客户端协议的 SSE 错误事件(响应头已写出时的唯一交付方式)。
@@ -349,12 +415,9 @@ func (g *Gateway) proxyOne(c *gin.Context, ch *store.Channel, plan *upstreamPlan
upBody := plan.body
// 直通 chat 流式:注入 stream_options.include_usage,保证末块带 usage(OpenAI 行为)
if stream && plan.path == "/chat/completions" && plan.lineConv == nil && !bytes.Contains(upBody, []byte(`"include_usage"`)) {
var m map[string]any
if json.Unmarshal(upBody, &m) == nil {
m["stream_options"] = map[string]any{"include_usage": true}
if b, err := json.Marshal(m); err == nil {
upBody = b
}
// 定点追加,避免整包 unmarshal/marshal(大请求体上是几百微秒 + 一倍分配)
if b, ok := injectStreamOptions(upBody, `{"include_usage":true}`); ok {
upBody = b
}
}
@@ -546,7 +609,7 @@ func (g *Gateway) streamHoldCopy(c *gin.Context, sh *streamHold, ch *store.Chann
return true, false, false, 0, nil
}
if trySink != nil {
trySink.outputText += sseContentText(out)
trySink.addContent(out)
}
}
if !releaseSeen && sseEventHasContent(out, clientProto) {
@@ -773,10 +836,37 @@ func extractServedModel(data []byte) string {
// outputText 累积已转发内容文本,用于流式中断时按 tiktoken 估算输出 token。
type usageSink struct {
us usageShape
outputText string
servedModel string // 上游实际应答的模型名(模型映射/聚合商场景下与请求的不同)
// text 累积已转发的正文,仅用于「上游没报输出用量」时按 tiktoken 估算。
// 用 Builder 而非字符串拼接:后者每块整体拷贝,长回答是 O(n²)
// (实测 10KB/500 块要 949µs、7.9MB 分配)。
text strings.Builder
// sawOutput:上游已经报过输出 token,估算不再需要,后面不再累积正文。
sawOutput bool
}
// textCap 正文累积上限:超过后不再累积(估算本来只是近似,不值得为它吃满内存)。
const textCap = 512 << 10
// addContent 从一条已转发的 SSE 行里取出正文并累积。
// 先做廉价的子串预过滤,只有可能携带正文的行才走 JSON 解析——一份流里
// role/usage/finish 之类的元事件占比不低,逐行 JSON 解析是纯浪费。
func (u *usageSink) addContent(line []byte) {
if u.sawOutput || u.text.Len() >= textCap {
return
}
if !bytes.Contains(line, []byte(`"content"`)) && !bytes.Contains(line, []byte(`"delta"`)) && !bytes.Contains(line, []byte(`"text"`)) {
return
}
if s := fastContentText(line); s != "" {
u.text.WriteString(s)
}
}
// textOf 已累积的正文。
func (u *usageSink) textOf() string { return u.text.String() }
func (u *usageSink) push(raw json.RawMessage) {
if len(raw) == 0 {
return
@@ -793,6 +883,11 @@ func (u *usageSink) push(raw json.RawMessage) {
u.us = t
return
}
// 一旦上游报过输出 token,正文估算就没有意义了(也就不用再累积)
if t.CompletionTokens > 0 || t.OutputTokens > 0 {
u.sawOutput = true
u.text.Reset()
}
// 零值不覆盖:不同事件携带不同字段
if t.PromptTokens > 0 {
u.us.PromptTokens = t.PromptTokens
@@ -850,7 +945,10 @@ func (u *usageSink) merge(o *usageSink) {
if o.us.CacheCreationInputTokens > 0 {
u.us.CacheCreationInputTokens = o.us.CacheCreationInputTokens
}
u.outputText += o.outputText
if !u.sawOutput && !o.sawOutput {
u.text.WriteString(o.text.String())
}
u.sawOutput = u.sawOutput || o.sawOutput
if o.servedModel != "" && u.servedModel == "" {
u.servedModel = o.servedModel
}
@@ -907,21 +1005,23 @@ func (g *Gateway) finishUsage(c *gin.Context, ch *store.Channel, start time.Time
// 非流式上游必返回 usage,此处 in/out 非 0 不受影响。
if status == store.UsageStatusCanceled || status == store.UsageStatusSuccess {
if in == 0 && cacheRead == 0 && cacheCreate == 0 {
if est, ok := c.Get("est_input_text"); ok {
if v, ok2 := est.(string); ok2 && v != "" {
in = int64(tokenizer.Count(v, mn))
if v, ok := c.Get("req_body"); ok {
if raw, ok2 := v.([]byte); ok2 && len(raw) > 0 {
if text := requestText(raw); text != "" {
in = int64(tokenizer.Count(text, mn))
}
}
}
}
if out == 0 && sink != nil && sink.outputText != "" {
out = int64(tokenizer.Count(sink.outputText, mn))
if out == 0 && sink != nil && sink.textOf() != "" {
out = int64(tokenizer.Count(sink.textOf(), mn))
}
}
var model store.Model
var cost float64
var modelID uint64
_ = g.db.Where("name = ?", mn).First(&model).Error
model, _ = g.modelByName(mn)
if model.ID > 0 {
modelID = model.ID
// 计价口径:in=非缓存输入、cacheRead=缓存读、cacheCreate=缓存写(Anthropic 语义,
+80
View File
@@ -0,0 +1,80 @@
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
}
+4 -1
View File
@@ -44,8 +44,11 @@ func Open(driver, dsn string) (*gorm.DB, error) {
dialector = sqlite.Open(dsn)
}
// PrepareStmt 让 GORM 缓存预编译语句:实测请求路径上 sqlite3Prepare(SQL 解析)
// 占 CPU 的 15%,每个请求十来个查询全在重复解析同几条 SQL。
db, err := gorm.Open(dialector, &gorm.Config{
Logger: gormlogger.Default.LogMode(gormlogger.Warn),
Logger: gormlogger.Default.LogMode(gormlogger.Warn),
PrepareStmt: true,
})
if err != nil {
return nil, err
+126 -62
View File
@@ -127,6 +127,10 @@ func (r *Recorder) flushRetry(logs []*store.UsageLog) {
// flush 批量插入用量明细,并同步更新余额、余额流水与日聚合。
// 记账口径:单次成本 = in×in_price + out×out_price + cache_read×cache_read_price(每百万 token)。
//
// 按批聚合而不是逐行:每行原本 5 条语句(查用户、查流水幂等、扣余额、写流水、
// upsert 日聚合),32 条一批就是 160 条;同一用户的并发请求还会在 users 行锁上排队。
// 现在按用户聚合一次扣款、批量写流水、按 (user, model, 日期) 聚合一次日表。
func (r *Recorder) flush(logs []*store.UsageLog) error {
if len(logs) == 0 {
return nil
@@ -135,77 +139,137 @@ func (r *Recorder) flush(logs []*store.UsageLog) error {
if err := tx.Create(logs).Error; err != nil {
return err
}
for _, l := range logs {
// 计费范围:success(正常完成)与 canceled(流式中断,按已生成部分收费)
if (l.Status != store.UsageStatusSuccess && l.Status != store.UsageStatusCanceled) || l.Cost <= 0 {
continue
}
if err := r.billUser(tx, l); err != nil {
// 明细已插入但扣费失败:宁可吵也不能静默漏账
log.Printf("usage: billing failed trace=%s user=%d cost=%.6f: %v", l.TraceID, l.UserID, l.Cost, err)
}
}
return nil
return r.bill(tx, logs)
})
}
// billUser 对单条明细扣余额、写流水、更新日聚合。
// 扣费用 SQL 侧原子递减(balance = balance - cost):SQLite 驱动不支持 SELECT ... FOR UPDATE,
// Go 侧读-改-写在并发下会互相覆盖丢扣费。
func (r *Recorder) billUser(tx *gorm.DB, l *store.UsageLog) error {
var user store.User
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).First(&user, l.UserID).Error; err != nil {
return err
// bill 对一批明细计费:按用户聚合扣款,批量写流水,按 (user, model, 日期) 聚合日表。
func (r *Recorder) bill(tx *gorm.DB, logs []*store.UsageLog) error {
// 1. 可计费行:success/canceled 且有成本
billable := make([]*store.UsageLog, 0, len(logs))
refs := make([]string, 0, len(logs))
for _, l := range logs {
if (l.Status != store.UsageStatusSuccess && l.Status != store.UsageStatusCanceled) || l.Cost <= 0 {
continue
}
billable = append(billable, l)
refs = append(refs, usageRefID(l.TraceID))
}
// 余额允许为负(流式请求不中断;后续请求被 402 拒),所以 WHERE 只按 id 匹配
// 同一 trace 已记过账(重试/重复 flush)则跳过:流水 ref_id 有唯一索引,
// 但扣款本身没有,不挡住会重复扣钱
var billed int64
tx.Model(&store.BalanceLog{}).Where("ref_id = ?", usageRefID(l.TraceID)).Count(&billed)
if billed > 0 {
if len(billable) == 0 {
return nil
}
res := tx.Model(&store.User{}).Where("id = ?", l.UserID).Update("balance", gorm.Expr("balance - ?", l.Cost))
if res.Error != nil {
return res.Error
}
if res.RowsAffected == 0 {
return gorm.ErrRecordNotFound
}
newBalance := user.Balance - l.Cost
if err := tx.Create(&store.BalanceLog{
UserID: l.UserID,
Change: -l.Cost,
BalanceAfter: newBalance,
Type: store.BalanceTypeUsage,
RefID: usageRefID(l.TraceID),
Remark: "usage: " + l.ModelName,
}).Error; err != nil {
// 2. 已记过账的 ref(重试/重复 flush):一条查询取代每行一次 COUNT
already := map[string]bool{}
var billed []string
if err := tx.Model(&store.BalanceLog{}).Where("ref_id IN ?", refs).Pluck("ref_id", &billed).Error; err != nil {
return err
}
for _, ref := range billed {
already[ref] = true
}
// 日聚合 upsert
date := l.CreatedAt.UTC().Format("2006-01-02")
if err := tx.Clauses(clause.OnConflict{
Columns: []clause.Column{{Name: "user_id"}, {Name: "model_id"}, {Name: "date"}},
DoUpdates: clause.Assignments(map[string]any{
"requests": gorm.Expr("requests + 1"),
"input_tokens": gorm.Expr("input_tokens + ?", l.InputTokens),
"output_tokens": gorm.Expr("output_tokens + ?", l.OutputTokens),
"cache_read_tokens": gorm.Expr("cache_read_tokens + ?", l.CacheReadTokens),
"cost": gorm.Expr("cost + ?", l.Cost),
}),
}).Create(&store.UsageDaily{
UserID: l.UserID,
ModelID: l.ModelID,
Date: date,
Requests: 1,
InputTokens: l.InputTokens,
OutputTokens: l.OutputTokens,
CacheReadTokens: l.CacheReadTokens,
Cost: l.Cost,
}).Error; err != nil {
return err
// 3. 过滤重复(含同批内同 trace),按用户分组
byUser := map[uint64][]*store.UsageLog{}
order := make([]uint64, 0, len(byUser))
seen := map[string]bool{}
for _, l := range billable {
ref := usageRefID(l.TraceID)
if already[ref] || seen[ref] {
continue
}
seen[ref] = true
if _, ok := byUser[l.UserID]; !ok {
order = append(order, l.UserID)
}
byUser[l.UserID] = append(byUser[l.UserID], l)
}
if len(byUser) == 0 {
return nil
}
// 4. 每个用户一次扣款,流水批量插入
ledger := make([]store.BalanceLog, 0, len(billable))
for _, uid := range order {
rows := byUser[uid]
total := 0.0
for _, l := range rows {
total += l.Cost
}
var user store.User
if err := tx.First(&user, uid).Error; err != nil {
// 用户被删了:明细已插入,扣费无处可扣,吵一声但不能拖垮整批
log.Printf("usage: user %d not found, %d records unbilled: %v", uid, len(rows), err)
continue
}
// 扣费用 SQL 侧原子递减(balance = balance - ?):SQLite 驱动不支持
// SELECT ... FOR UPDATE,Go 侧读-改-写在并发下会互相覆盖丢扣费。
// 余额允许为负(流式请求不中断;后续请求被 402 拒)。
if err := tx.Model(&store.User{}).Where("id = ?", uid).
Update("balance", gorm.Expr("balance - ?", total)).Error; err != nil {
return err
}
// 流水里的 balance_after 是批内累计值,与逐行扣款时的语义一致
running := user.Balance
for _, l := range rows {
running -= l.Cost
ledger = append(ledger, store.BalanceLog{
UserID: uid,
Change: -l.Cost,
BalanceAfter: running,
Type: store.BalanceTypeUsage,
RefID: usageRefID(l.TraceID),
Remark: "usage: " + l.ModelName,
})
}
}
if len(ledger) > 0 {
// ref_id 唯一索引兜底:并发下别处已写入同 ref 时跳过,不整批失败也不重复扣款
if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&ledger).Error; err != nil {
return err
}
}
// 5. 日聚合:按 (user, model, 日期) 汇总后一次 upsert
type dailyKey struct {
user uint64
model uint64
date string
}
agg := map[dailyKey]*store.UsageDaily{}
dailyOrder := make([]dailyKey, 0, len(billable))
for _, l := range billable {
ref := usageRefID(l.TraceID)
if !seen[ref] {
continue // 重复 trace,不计入日聚合
}
k := dailyKey{user: l.UserID, model: l.ModelID, date: l.CreatedAt.UTC().Format("2006-01-02")}
d, ok := agg[k]
if !ok {
d = &store.UsageDaily{UserID: k.user, ModelID: k.model, Date: k.date}
agg[k] = d
dailyOrder = append(dailyOrder, k)
}
d.Requests++
d.InputTokens += l.InputTokens
d.OutputTokens += l.OutputTokens
d.CacheReadTokens += l.CacheReadTokens
d.Cost += l.Cost
}
for _, k := range dailyOrder {
d := agg[k]
if err := tx.Clauses(clause.OnConflict{
Columns: []clause.Column{{Name: "user_id"}, {Name: "model_id"}, {Name: "date"}},
DoUpdates: clause.Assignments(map[string]any{
"requests": gorm.Expr("requests + ?", d.Requests),
"input_tokens": gorm.Expr("input_tokens + ?", d.InputTokens),
"output_tokens": gorm.Expr("output_tokens + ?", d.OutputTokens),
"cache_read_tokens": gorm.Expr("cache_read_tokens + ?", d.CacheReadTokens),
"cost": gorm.Expr("cost + ?", d.Cost),
}),
}).Create(d).Error; err != nil {
return err
}
}
return nil
}
+48
View File
@@ -170,3 +170,51 @@ func TestBillingRefIDImpotency(t *testing.T) {
t.Errorf("balance logs = %d, want 1", n)
}
}
// 一批同用户请求:一次扣款、一条日聚合(而不是逐行 upsert)。
func TestBillingAggregatesBatch(t *testing.T) {
db := openTestDB(t)
u := seedUser(t, db, 100)
m := seedModel(t, db)
rec := NewRecorder(db)
const n = 32
for i := 0; i < n; i++ {
rec.Record(&store.UsageLog{
TraceID: fmt.Sprintf("agg-%d", i), UserID: u, ModelID: m, ModelName: "m1",
InputTokens: 10, OutputTokens: 5, Cost: 0.01,
Status: store.UsageStatusSuccess, CreatedAt: time.Now().UTC(),
})
}
rec.Close()
var user store.User
if err := db.First(&user, u).Error; err != nil {
t.Fatal(err)
}
if want := 100 - 0.01*n; user.Balance-want > 1e-9 || user.Balance-want < -1e-9 {
t.Fatalf("余额 = %v,期望 %v", user.Balance, want)
}
var ledgers int64
db.Model(&store.BalanceLog{}).Where("user_id = ?", u).Count(&ledgers)
if ledgers != n {
t.Fatalf("流水条数 = %d,期望 %d(每请求一条明细)", ledgers, n)
}
var daily []store.UsageDaily
db.Where("user_id = ?", u).Find(&daily)
if len(daily) != 1 {
t.Fatalf("日聚合应只有 1 行,实际 %d", len(daily))
}
if daily[0].Requests != n {
t.Fatalf("日聚合 requests = %d,期望 %d", daily[0].Requests, n)
}
if d := daily[0].Cost - 0.01*n; d > 1e-9 || d < -1e-9 {
t.Fatalf("日聚合 cost = %v,期望 %v", daily[0].Cost, 0.01*n)
}
// 流水里的 balance_after 应与扣款后余额一致(批内累计)
var last store.BalanceLog
db.Where("user_id = ?", u).Order("id DESC").First(&last)
if last.BalanceAfter-user.Balance > 1e-9 || user.Balance-last.BalanceAfter > 1e-9 {
t.Fatalf("最后一条流水 balance_after = %v,实际余额 %v", last.BalanceAfter, user.Balance)
}
}
+2
View File
@@ -125,6 +125,8 @@ export interface UsageLog {
error_code: string | null
raw_request?: string
raw_response?: string
/** 列表接口不再携带原始体,需要时用 /admin/usage/:id 取 */
has_raw?: boolean
created_at: string
user?: string
user_id?: number
+19 -4
View File
@@ -20,6 +20,7 @@ const pageSize = 15
const viewing = ref<UsageLog | null>(null)
const viewTab = ref<'request' | 'response'>('request')
const rawLoading = ref(false)
// cacheLine 缓存读/写的紧凑摘要(换行会碎成多行,只放两个数);无缓存流量返回空串。
function cacheLine(l: UsageLog): string {
@@ -65,9 +66,22 @@ function servedModel(l: UsageLog): string {
return l.served_model && l.served_model !== l.model ? l.served_model : ''
}
function openRaw(l: UsageLog) {
// openRaw 打开原始体弹窗:列表接口不返回原始体(可能几十 MB),
// 点开时才按 id 拉取详情。
async function openRaw(l: UsageLog) {
viewing.value = l
viewTab.value = 'request'
if (l.raw_request || l.raw_response) return
rawLoading.value = true
try {
const { data } = await http.get(`/admin/usage/${l.id}`)
viewing.value = { ...l, ...data.data }
Object.assign(l, data.data) // 回填列表行,再点开不必重复请求
} catch (e) {
toast.err(errMsg(e))
} finally {
rawLoading.value = false
}
}
// 用户自动补全候选(防抖 300ms)
@@ -198,7 +212,7 @@ onMounted(load)
<span v-if="cacheExtra(l)" class="mono-num">{{ cacheExtra(l) }}</span>
<span class="mono-num">{{ l.latency_ms }}ms</span>
<span class="mono-num w-full">{{ fmtTime(l.created_at) }}</span>
<Button v-if="l.raw_request" size="sm" variant="ghost" class="ml-auto" @click="openRaw(l)">查看原始</Button>
<Button v-if="l.has_raw || l.raw_request" size="sm" variant="ghost" class="ml-auto" @click="openRaw(l)">查看原始</Button>
</div>
</div>
<p v-if="logs.length === 0" class="card px-4 py-8 text-center text-sm text-muted">暂无请求记录</p>
@@ -248,7 +262,7 @@ onMounted(load)
{{ fmtClock(l.created_at) }}
</td>
<td class="px-4 py-2.5">
<Button v-if="l.raw_request" size="sm" variant="ghost" @click="openRaw(l)">查看</Button>
<Button v-if="l.has_raw || l.raw_request" size="sm" variant="ghost" @click="openRaw(l)">查看</Button>
<span v-else class="text-xs text-muted">-</span>
</td>
</tr>
@@ -296,8 +310,9 @@ onMounted(load)
</button>
</div>
</div>
<p v-if="rawLoading" class="rounded-md border border-edge bg-surface p-4 text-center text-xs text-muted">正在读取…</p>
<pre
v-if="viewTab === 'request' ? viewing.raw_request : viewing.raw_response"
v-else-if="viewTab === 'request' ? viewing.raw_request : viewing.raw_response"
class="max-h-[60vh] overflow-auto rounded-md border border-edge bg-surface p-3 font-mono text-xs leading-relaxed text-ink whitespace-pre-wrap break-all"
>{{ viewTab === 'request' ? viewing.raw_request : viewing.raw_response }}</pre>
<p v-else class="rounded-md border border-edge bg-surface p-4 text-center text-xs text-muted">