perf: 修复调度器 Redis N+1、补齐数据库索引、消除请求体重复解析等性能问题

DouDOU-start · DouDOU-start · commit 827890543e0e · 2026-06-15T18:46:20.000+08:00
- scheduler: 删除每次 RPM 查询的 Redis TIME 命令，改用本地时间
- scheduler: 预获取并缓存 getCurrentLoad 结果，选号 Redis 调用量从 2N 降到 N
- schema: apikey 添加 (key_hash) 唯一索引和 (status, created_at) 复合索引
- schema: account 添加 (platform, state) 复合索引
- schema: balancelog 添加 (user_id_snapshot, created_at) 复合索引
- schema: usagelog 添加 (model, created_at) 复合索引
- schema: usersubscription 添加 (status) 索引
- plugin/request: 请求体 JSON 解析从 3 次降到 1 次，缓存 imageToolPayload
- plugin/request: buildHeaders 从全量 Clone 改为选择性拷贝，跳过 hop-by-hop 头
- store: ListTrendLogs 和 TrendEntries 添加 LIMIT 50000 防止 OOM
- plugin/extension_proxy: 流式写入检测 client disconnect，避免无效资源消耗
- server/dynamic_router: 双重 RLock 合并为单次，减少锁竞争
diff --git a/backend/ent/migrate/schema.go b/backend/ent/migrate/schema.go
diff --git a/backend/ent/schema/account.go b/backend/ent/schema/account.go
@@ -4,6 +4,7 @@ import (
 	"entgo.io/ent"
 	"entgo.io/ent/schema/edge"
 	"entgo.io/ent/schema/field"
+	"entgo.io/ent/schema/index"
 )
 
 // Account 上游 AI 账户。
@@ -47,6 +48,12 @@ func (Account) Fields() []ent.Field {
 	}
 }
 
+func (Account) Indexes() []ent.Index {
+	return []ent.Index{
+		index.Fields("platform", "state"),
+	}
+}
+
 func (Account) Edges() []ent.Edge {
 	return []ent.Edge{
 		edge.To("groups", Group.Type),
diff --git a/backend/ent/schema/apikey.go b/backend/ent/schema/apikey.go
@@ -4,6 +4,7 @@ import (
 	"entgo.io/ent"
 	"entgo.io/ent/schema/edge"
 	"entgo.io/ent/schema/field"
+	"entgo.io/ent/schema/index"
 )
 
 // APIKey API 密钥
@@ -35,6 +36,13 @@ func (APIKey) Fields() []ent.Field {
 	}
 }
 
+func (APIKey) Indexes() []ent.Index {
+	return []ent.Index{
+		index.Fields("key_hash").Unique(),
+		index.Fields("status", "created_at"),
+	}
+}
+
 func (APIKey) Edges() []ent.Edge {
 	return []ent.Edge{
 		edge.From("user", User.Type).Ref("api_keys").Unique().Required(),
diff --git a/backend/ent/schema/balancelog.go b/backend/ent/schema/balancelog.go
@@ -43,5 +43,6 @@ func (BalanceLog) Indexes() []ent.Index {
 	return []ent.Index{
 		// Postgres 唯一索引对 NULL 不互斥，仅约束显式提供的幂等键
 		index.Fields("idempotency_key").Unique(),
+		index.Fields("user_id_snapshot", "created_at"),
 	}
 }
diff --git a/backend/ent/schema/usagelog.go b/backend/ent/schema/usagelog.go
@@ -107,6 +107,8 @@ func (UsageLog) Indexes() []ent.Index {
 			StorageKey("usage_log_platform_created_at"),
 		index.Fields("user_id_snapshot", "created_at").
 			StorageKey("usage_log_user_snapshot_created_at"),
+		index.Fields("model", "created_at").
+			StorageKey("usage_log_model_created_at"),
 		index.Edges("user").
 			StorageKey("usage_log_user"),
 		index.Edges("api_key").
diff --git a/backend/ent/schema/usersubscription.go b/backend/ent/schema/usersubscription.go
@@ -4,6 +4,7 @@ import (
 	"entgo.io/ent"
 	"entgo.io/ent/schema/edge"
 	"entgo.io/ent/schema/field"
+	"entgo.io/ent/schema/index"
 )
 
 // UserSubscription 用户订阅
@@ -22,6 +23,12 @@ func (UserSubscription) Fields() []ent.Field {
 	}
 }
 
+func (UserSubscription) Indexes() []ent.Index {
+	return []ent.Index{
+		index.Fields("status"),
+	}
+}
+
 func (UserSubscription) Edges() []ent.Edge {
 	return []ent.Edge{
 		edge.From("user", User.Type).Ref("subscriptions").Unique().Required(),
diff --git a/backend/internal/infra/store/dashboard_store.go b/backend/internal/infra/store/dashboard_store.go
@@ -92,6 +92,7 @@ func (s *DashboardStore) ListTrendLogs(ctx context.Context, startTime, endTime t
 		preds = append(preds, usageUserPredicate(int64(userID)))
 	}
 
+	const trendLogLimit = 50000
 	list, err := s.db.UsageLog.Query().
 		Where(preds...).
 		Select(
@@ -106,6 +107,8 @@ func (s *DashboardStore) ListTrendLogs(ctx context.Context, startTime, endTime t
 			entusagelog.FieldTotalCost,
 			entusagelog.FieldCreatedAt,
 		).
+		Order(ent.Desc(entusagelog.FieldCreatedAt)).
+		Limit(trendLogLimit).
 		All(ctx)
 	if err != nil {
 		return nil, err
diff --git a/backend/internal/infra/store/usage_store.go b/backend/internal/infra/store/usage_store.go
@@ -384,6 +384,7 @@ func (s *UsageStore) TrendEntries(ctx context.Context, filter appusage.TrendFilt
 		query = query.Where(entusagelog.CreatedAtGTE(time.Now().Add(-time.Duration(filter.DefaultRecentHours) * time.Hour)))
 	}
 
+	const trendEntryLimit = 50000
 	logs, err := query.
 		Select(
 			entusagelog.FieldInputTokens,
@@ -395,6 +396,8 @@ func (s *UsageStore) TrendEntries(ctx context.Context, filter appusage.TrendFilt
 			entusagelog.FieldTotalCost,
 			entusagelog.FieldCreatedAt,
 		).
+		Order(ent.Desc(entusagelog.FieldCreatedAt)).
+		Limit(trendEntryLimit).
 		All(ctx)
 	if err != nil {
 		return nil, err
diff --git a/backend/internal/plugin/extension_proxy.go b/backend/internal/plugin/extension_proxy.go
@@ -239,7 +239,9 @@ func (ep *ExtensionProxy) handleStream(c *gin.Context, ext *sdkgrpc.ExtensionGRP
 		}
 
 		if len(chunk.Data) > 0 {
-			_, _ = c.Writer.Write(chunk.Data)
+			if _, err := c.Writer.Write(chunk.Data); err != nil {
+				return // 客户端已断开
+			}
 			c.Writer.Flush()
 		}
 
diff --git a/backend/internal/plugin/forwarder.go b/backend/internal/plugin/forwarder.go
@@ -129,7 +129,7 @@ func (f *Forwarder) Forward(c *gin.Context) {
 	defer releaseClientQuota()
 
 	requirements := routing.Requirements{
-		NeedsImage: requestNeedsImage(f.manager, state.requestPath, state.model, state.body),
+		NeedsImage: requestNeedsImageCached(f.manager, state),
 	}
 	routes := routesForAPIKey(state, requirements)
 	if len(routes) == 0 {
diff --git a/backend/internal/plugin/request.go b/backend/internal/plugin/request.go
@@ -61,20 +61,22 @@ func (f *Forwarder) parseRequest(c *gin.Context) (*forwardState, bool) {
 	}
 
 	return &forwardState{
-		startedAt:         startedAt,
-		requestPath:       path,
-		body:              body,
-		model:             parsed.Model,
-		schedulingModels:  schedulingModels,
-		schedulingModel:   schedulingModel,
-		stream:            parsed.Stream,
-		realtime:          parsed.Stream,
-		sessionID:         parsed.SessionID,
-		reasoningEffort:   parsed.ReasoningEffort,
-		accountReq:        accountRequirementsForRequest(f.manager, path, parsed.Model, body),
-		requestedPlatform: requestedPlatform,
-		keyInfo:           keyInfo,
-		plugin:            inst,
+		startedAt:             startedAt,
+		requestPath:           path,
+		body:                  body,
+		model:                 parsed.Model,
+		schedulingModels:      schedulingModels,
+		schedulingModel:       schedulingModel,
+		stream:                parsed.Stream,
+		realtime:              parsed.Stream,
+		sessionID:             parsed.SessionID,
+		reasoningEffort:       parsed.ReasoningEffort,
+		accountReq:            accountRequirementsForRequestCached(f.manager, path, parsed.Model, &parsed),
+		imageToolPayloadValid: parsed.imageToolPayloadValid,
+		imageToolPayload:      parsed.imageToolPayload,
+		requestedPlatform:     requestedPlatform,
+		keyInfo:               keyInfo,
+		plugin:                inst,
 	}, true
 }
 
@@ -120,12 +122,18 @@ func parseBody(body []byte, contentType string) parsedRequest {
 	var fields requestFields
 	if json.Unmarshal(body, &fields) == nil {
 		effort := extractAndNormalizeReasoningEffort(fields)
-		return parsedRequest{
+		pr := parsedRequest{
 			Model:           fields.Model,
 			Stream:          fields.Stream,
 			SessionID:       fields.Metadata.UserID,
 			ReasoningEffort: effort,
 		}
+		// 一次性提取 image tool 信息，避免后续 requestNeedsImage / accountRequirementsForRequest 重复反序列化 body
+		if payload, ok := parseImageToolPayloadFromFields(body); ok {
+			pr.imageToolPayloadValid = true
+			pr.imageToolPayload = payload
+		}
+		return pr
 	}
 	if strings.HasPrefix(contentType, "multipart/") {
 		return parseMultipartFields(body, contentType)
@@ -182,6 +190,17 @@ func requestNeedsImage(mgr *Manager, path, model string, body []byte) bool {
 	return isImageAPIPath(path) || isImageModel(mgr, model) || hasForcedImageGenerationTool(body)
 }
 
+// requestNeedsImageCached 使用 forwardState 中缓存的 imageToolPayload，避免重复反序列化 body。
+func requestNeedsImageCached(mgr *Manager, state *forwardState) bool {
+	if isImageAPIPath(state.requestPath) || isImageModel(mgr, state.model) {
+		return true
+	}
+	if state.imageToolPayloadValid {
+		return hasForcedImageGenerationToolFromPayload(state.imageToolPayload)
+	}
+	return hasForcedImageGenerationTool(state.body)
+}
+
 // requestHasImageWorkload 判断请求是否需要更长的图片工作超时。
 // 这里也保留 Responses API 的 image_generation tool 识别，用于放宽生成链路等待时间。
 func requestHasImageWorkload(mgr *Manager, path, model string, body []byte) bool {
@@ -207,6 +226,36 @@ func accountRequirementsForRequest(mgr *Manager, path, model string, body []byte
 	return scheduler.AccountRequirements{Workload: scheduler.WorkloadChat}
 }
 
+// accountRequirementsForRequestCached 复用 parsedRequest 中已缓存的 imageToolPayload，避免重复反序列化 body。
+func accountRequirementsForRequestCached(mgr *Manager, path, model string, parsed *parsedRequest) scheduler.AccountRequirements {
+	if isImageAPIPath(path) || isImageModel(mgr, model) {
+		return scheduler.AccountRequirements{
+			Workload: scheduler.WorkloadImage,
+			ImageProtocols: []scheduler.ImageProtocol{
+				scheduler.ImageProtocolImagesAPI,
+				scheduler.ImageProtocolResponsesTool,
+			},
+		}
+	}
+	if parsed.imageToolPayloadValid && hasForcedImageGenerationToolFromPayload(parsed.imageToolPayload) {
+		return scheduler.AccountRequirements{
+			Workload:       scheduler.WorkloadImage,
+			ImageProtocols: []scheduler.ImageProtocol{scheduler.ImageProtocolResponsesTool},
+		}
+	}
+	return scheduler.AccountRequirements{Workload: scheduler.WorkloadChat}
+}
+
+func hasForcedImageGenerationToolFromPayload(payload imageToolPayload) bool {
+	if len(payload.ToolChoice) == 0 {
+		return false
+	}
+	if payload.toolChoiceForcesImageGeneration() {
+		return true
+	}
+	return payload.toolChoiceRequiresOnlyImageGeneration()
+}
+
 func isImageAPIPath(path string) bool {
 	if path == "" {
 		return false
@@ -274,6 +323,10 @@ func parseImageToolPayload(body []byte) (imageToolPayload, bool) {
 	return payload, true
 }
 
+// parseImageToolPayloadFromFields 复用 parseBody 阶段已经验证可解析的 body，
+// 只提取 tools/tool_choice 字段。与 parseImageToolPayload 逻辑相同，仅命名区分调用场景。
+var parseImageToolPayloadFromFields = parseImageToolPayload
+
 func (p imageToolPayload) imageGenerationToolCount() int {
 	count := 0
 	for _, tool := range p.Tools {
@@ -420,9 +473,20 @@ func applyAccountCapabilityHeaders(headers http.Header, acc *ent.Account) {
 	}
 }
 
-// buildHeaders 克隆请求头并附加 X-Airgate-* 系列（分组级 service_tier / 强制 instructions / 插件开关）。
+// buildHeaders 选择性拷贝请求头并附加 X-Airgate-* 系列（分组级 service_tier / 强制 instructions / 插件开关）。
+// 跳过 hop-by-hop 头和认证头（插件用 Account.Credentials 鉴权上游），避免全量 Clone 的分配开销。
 func buildHeaders(source http.Header, keyInfo *auth.APIKeyInfo) http.Header {
-	headers := source.Clone()
+	headers := make(http.Header, len(source)+8)
+	for k, v := range source {
+		lower := strings.ToLower(k)
+		if lower == "authorization" || lower == "x-api-key" ||
+			lower == "connection" || lower == "keep-alive" ||
+			lower == "proxy-authorization" || lower == "te" ||
+			lower == "transfer-encoding" || lower == "upgrade" {
+			continue
+		}
+		headers[k] = v
+	}
 	if keyInfo.UserID > 0 {
 		headers.Set("X-Airgate-User-ID", strconv.Itoa(keyInfo.UserID))
 	}
diff --git a/backend/internal/plugin/types.go b/backend/internal/plugin/types.go
@@ -32,6 +32,10 @@ type forwardState struct {
 	reasoningEffort string
 	accountReq      scheduler.AccountRequirements
 
+	// 缓存的 image tool payload，避免 forwarder 热路径上重复反序列化 body
+	imageToolPayloadValid bool
+	imageToolPayload      imageToolPayload
+
 	requestedPlatform string
 	selectedRoute     routing.Candidate
 
@@ -54,6 +58,10 @@ type parsedRequest struct {
 	Stream          bool
 	SessionID       string
 	ReasoningEffort string // 推理强度档位
+
+	// 缓存 image tool payload 解析结果，避免 requestNeedsImage / accountRequirementsForRequest 重复反序列化 body
+	imageToolPayloadValid bool
+	imageToolPayload      imageToolPayload
 }
 
 // requestFields 一次性 Unmarshal 的 JSON 字段结构。
diff --git a/backend/internal/scheduler/rpm.go b/backend/internal/scheduler/rpm.go
@@ -24,14 +24,11 @@ func NewRPMCounter(rdb *redis.Client) *RPMCounter {
 	return &RPMCounter{rdb: rdb}
 }
 
-// getMinuteKey 使用 Redis 服务器时间生成分钟粒度的 Redis key，避免分布式时钟偏差
-func (r *RPMCounter) getMinuteKey(ctx context.Context, accountID int) string {
-	t, err := r.rdb.Time(ctx).Result()
-	if err != nil {
-		// fallback to local time
-		t = time.Now()
-	}
-	minute := t.Unix() / 60
+// getMinuteKey 生成分钟粒度的 Redis key。
+// 使用本地时间：分钟级 key 粒度下各机器时钟偏差（通常 <1s）完全可接受，
+// 省去每次调用一次 Redis TIME 命令的额外 RTT。
+func (r *RPMCounter) getMinuteKey(_ context.Context, accountID int) string {
+	minute := time.Now().Unix() / 60
 	return fmt.Sprintf("rpm:%d:%d", accountID, minute)
 }
 
diff --git a/backend/internal/scheduler/selection.go b/backend/internal/scheduler/selection.go
diff --git a/backend/internal/server/dynamic_router.go b/backend/internal/server/dynamic_router.go

Original file line number	Diff line number	Diff line change
`@@ -4,6 +4,7 @@ import (`
`4`	`4`	`"entgo.io/ent"`
`5`	`5`	`"entgo.io/ent/schema/edge"`
`6`	`6`	`"entgo.io/ent/schema/field"`
	`7`	`+ "entgo.io/ent/schema/index"`
`7`	`8`	`)`
`8`	`9`
`9`	`10`	`// Account 上游 AI 账户。`
`@@ -47,6 +48,12 @@ func (Account) Fields() []ent.Field {`
`47`	`48`	`}`
`48`	`49`	`}`
`49`	`50`
	`51`	`+func (Account) Indexes() []ent.Index {`
	`52`	`+ return []ent.Index{`
	`53`	`+ index.Fields("platform", "state"),`
	`54`	`+ }`
	`55`	`+}`
	`56`	`+`
`50`	`57`	`func (Account) Edges() []ent.Edge {`
`51`	`58`	`return []ent.Edge{`
`52`	`59`	`edge.To("groups", Group.Type),`
Original file line number	Diff line number	Diff line change
`@@ -43,5 +43,6 @@ func (BalanceLog) Indexes() []ent.Index {`
`43`	`43`	`return []ent.Index{`
`44`	`44`	`// Postgres 唯一索引对 NULL 不互斥，仅约束显式提供的幂等键`
`45`	`45`	`index.Fields("idempotency_key").Unique(),`
	`46`	`+ index.Fields("user_id_snapshot", "created_at"),`
`46`	`47`	`}`
`47`	`48`	`}`
Original file line number	Diff line number	Diff line change
`@@ -239,7 +239,9 @@ func (ep ExtensionProxy) handleStream(c gin.Context, ext *sdkgrpc.ExtensionGRP`
`239`	`239`	`}`
`240`	`240`
`241`	`241`	`if len(chunk.Data) > 0 {`
`242`		`- _, _ = c.Writer.Write(chunk.Data)`
	`242`	`+ if _, err := c.Writer.Write(chunk.Data); err != nil {`
	`243`	`+ return // 客户端已断开`
	`244`	`+ }`
`243`	`245`	`c.Writer.Flush()`
`244`	`246`	`}`
`245`	`247`
Original file line number	Diff line number	Diff line change
`@@ -129,7 +129,7 @@ func (f Forwarder) Forward(c gin.Context) {`
`129`	`129`	`defer releaseClientQuota()`
`130`	`130`
`131`	`131`	`requirements := routing.Requirements{`
`132`		`- NeedsImage: requestNeedsImage(f.manager, state.requestPath, state.model, state.body),`
	`132`	`+ NeedsImage: requestNeedsImageCached(f.manager, state),`
`133`	`133`	`}`
`134`	`134`	`routes := routesForAPIKey(state, requirements)`
`135`	`135`	`if len(routes) == 0 {`