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:
1 parent
d34d9bd2e6
commit
8c12207478
23 files changed
+1050
-170
No files matched your search
@@ -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
|
||||
|
||||
@@ -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 延迟;渠道页直接展示近况与运行期冷却剩余时间,可一键恢复
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
})
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
@@ -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{
|
||||
|
||||
@@ -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"),
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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")
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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 语义,
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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">
|
||||
|
||||
Reference in new issue
Block a user