perf: 代理鉴权复用用户查询,last_used_at 按分钟节流

- Auth 加载的完整用户放进上下文,resolveUser 直接复用,省掉每请求
  一次重复的 users 查询
- last_used_at 是纯展示字段,改为距上次更新满一分钟才写,高 QPS 下
  每请求一次 DB 写降到每密钥每分钟至多一次
This commit is contained in:
Sakurasan committed 2026-10-30 01:42:23 +08:00
1 parent 9d3648ffe7
commit 99829ef745
4 files changed
+47 -34

No files matched your search

+23 -23
View File
@@ -18,10 +18,10 @@ type chatTool struct {
}
type chatMsg struct {
Role string `json:"role"`
Content json.RawMessage `json:"content"`
ToolCallID string `json:"tool_call_id"`
ToolCalls []struct {
Role string `json:"role"`
Content json.RawMessage `json:"content"`
ToolCallID string `json:"tool_call_id"`
ToolCalls []struct {
ID string `json:"id"`
Function struct {
Name string `json:"name"`
@@ -31,14 +31,14 @@ type chatMsg struct {
}
type chatReq struct {
Model string `json:"model"`
Messages []chatMsg `json:"messages"`
Tools []chatTool `json:"tools"`
Temperature *float64 `json:"temperature"`
TopP *float64 `json:"top_p"`
MaxTokens *int `json:"max_tokens"`
Stop []string `json:"stop"`
Stream bool `json:"stream"`
Model string `json:"model"`
Messages []chatMsg `json:"messages"`
Tools []chatTool `json:"tools"`
Temperature *float64 `json:"temperature"`
TopP *float64 `json:"top_p"`
MaxTokens *int `json:"max_tokens"`
Stop []string `json:"stop"`
Stream bool `json:"stream"`
}
// chatToMessagesReq 将 OpenAI Chat 请求转为 Anthropic Messages 请求。
@@ -48,7 +48,7 @@ func chatToMessagesReq(body []byte) ([]byte, error) {
return nil, err
}
out := map[string]any{
"model": req.Model,
"model": req.Model,
"max_tokens": intOrNil(req.MaxTokens, 1024), // Anthropic 必填
}
if req.Stream {
@@ -157,9 +157,9 @@ func chatMsgToAnthropic(m chatMsg) any {
// 请求:Messages → Chat
type messagesReq struct {
Model string `json:"model"`
System json.RawMessage `json:"system"`
Messages []struct {
Model string `json:"model"`
System json.RawMessage `json:"system"`
Messages []struct {
Role string `json:"role"`
Content json.RawMessage `json:"content"`
} `json:"messages"`
@@ -299,13 +299,13 @@ func anthropicMsgToChat(role string, content json.RawMessage) []any {
// 响应:Messages → Chat
type messagesResp struct {
ID string `json:"id"`
Model string `json:"model"`
Content []struct {
Type string `json:"type"`
Text string `json:"text"`
ID string `json:"id"`
Name string `json:"name"`
ID string `json:"id"`
Model string `json:"model"`
Content []struct {
Type string `json:"type"`
Text string `json:"text"`
ID string `json:"id"`
Name string `json:"name"`
Input json.RawMessage `json:"input"`
} `json:"content"`
StopReason string `json:"stop_reason"`
@@ -75,7 +75,7 @@ func responsesInputToChat(raw json.RawMessage) []any {
var args any
_ = json.Unmarshal([]byte(strField(item["arguments"])), &args)
out = append(out, map[string]any{
"role": "assistant",
"role": "assistant",
"content": "",
"tool_calls": []any{map[string]any{
"id": strField(item["call_id"]),
@@ -212,9 +212,9 @@ func chatToResponsesReq(body []byte) ([]byte, error) {
switch m.Role {
case "tool":
input = append(input, map[string]any{
"type": "function_call_output",
"type": "function_call_output",
"call_id": m.ToolCallID,
"output": str(m.Content),
"output": str(m.Content),
})
case "assistant":
if len(m.ToolCalls) > 0 {
@@ -365,11 +365,11 @@ func chatToResponsesResp(body []byte) ([]byte, error) {
}
status := finish
return json.Marshal(map[string]any{
"id": "resp_" + strings.TrimPrefix(r.ID, "chatcmpl-"),
"object": "response",
"model": r.Model,
"status": status,
"output": output,
"id": "resp_" + strings.TrimPrefix(r.ID, "chatcmpl-"),
"object": "response",
"model": r.Model,
"status": status,
"output": output,
"usage": map[string]any{
"input_tokens": r.Usage.PromptTokens,
"output_tokens": r.Usage.CompletionTokens,
+13 -2
View File
@@ -25,6 +25,7 @@ import (
const (
CtxUserID = "proxy_user_id"
CtxUser = "proxy_user" // Auth 阶段加载的完整用户(含余额),resolveUser 直接复用
CtxKeyID = "proxy_key_id"
CtxKeyModel = "proxy_key_models"
CtxTrace = "proxy_trace_id"
@@ -291,11 +292,15 @@ func (g *Gateway) Auth(c *gin.Context) {
}
}
c.Set(CtxUser, &u)
c.Set(CtxUserID, u.ID)
c.Set(CtxKeyID, k.ID)
c.Set(CtxKeyModel, k.AllowedModels)
c.Set(CtxTrace, newTraceID())
g.db.Model(&store.APIKey{}).Where("id = ?", k.ID).Update("last_used_at", time.Now())
// last_used_at 按分钟节流:纯展示字段,不值得每请求一次 DB 写(高 QPS 下是可观开销)
if k.LastUsedAt == nil || time.Since(*k.LastUsedAt) >= time.Minute {
g.db.Model(&store.APIKey{}).Where("id = ?", k.ID).Update("last_used_at", time.Now())
}
c.Next()
}
@@ -385,8 +390,14 @@ func (g *Gateway) candidateChannels(model string) []channel.Candidate {
return g.ch.Candidates(model)
}
// resolveUser 取当前用户(含余额)。
// resolveUser 取当前用户(含余额)。Auth 中间件已查过一次并放入上下文,
// handler 里直接复用,省掉每请求一次重复查询。
func (g *Gateway) resolveUser(c *gin.Context) (*store.User, bool) {
if u, ok := c.Get(CtxUser); ok {
if u2, ok := u.(*store.User); ok {
return u2, true
}
}
uid, _ := c.Get(CtxUserID)
var u store.User
if err := g.db.First(&u, uid).Error; err != nil {
+3 -1
View File
@@ -50,7 +50,9 @@ func TestSSEDataPayloadSpacedAndMultiLine(t *testing.T) {
// 复现线上 qwen(dashscope)messages 流式缓存场景:
// message_start.usage.input_tokens 是总输入,message_delta.usage.input_tokens 是非缓存输入
// 且带 cache_read/cache_creation,是最终计费口径。合并后:
// in(落库)=input+cache_read+cache_creation,计价 in 只算非缓存部分。
//
// in(落库)=input+cache_read+cache_creation,计价 in 只算非缓存部分。
//
// 此前 delta 的 input 覆盖 start 的 input 导致总输入丢失(31790 → 8)。
func TestUsageSinkMessageDeltaAuthoritative(t *testing.T) {
sink := &usageSink{}