improve: 优化高并发下的性能问题

internal/logic/billing/pricing.go

@@ -48,6 +48,11 @@ type PricingResult struct {
 	CacheReadPrice     float64 // 缓存读取每 1M token 价格
 	CacheCreationPrice float64 // 缓存创建每 1M token 价格
 
+	// 模型最大输出 token 数（随定价一起缓存，供预扣估算使用，
+	// 避免 EstimatePreDeductAmount 每请求单独查一次 mdl_models）。
+	// 0 表示未设置（含旧缓存条目），使用方需自带默认值兜底。
+	MaxOutputTokens int
+
 	// 租户自定义阶梯定价（JSONB 解析后的原始数据，供 CalculateCost 使用）
 	CustomTiers []pricingTierRow
 }
@@ -79,18 +84,19 @@ func GetModelPrice(ctx context.Context, tenantID int64, modelName string) (*Pric
 		return &cached, nil
 	}
 
-	// 1. 查模型基础信息
+	// 1. 查模型基础信息（max_output_tokens 一并取出，随定价缓存供预扣估算复用）
 	type modelRow struct {
-		ID      int64  `json:"id"`
-		ModelId string `json:"model_id"`
-		Status  string `json:"status"`
+		ID              int64  `json:"id"`
+		ModelId         string `json:"model_id"`
+		Status          string `json:"status"`
+		MaxOutputTokens int    `json:"max_output_tokens"`
 	}
 
 	var model *modelRow
 	err := dao.MdlModels.Ctx(ctx).
 		Where("model_id", modelName).
 		Where("status", "active").
-		Fields("id, model_id, status").
+		Fields("id, model_id, status, max_output_tokens").
 		Scan(&model)
 	if err != nil {
 		return nil, gerror.Wrapf(err, "query model price")
@@ -251,6 +257,7 @@ func GetModelPrice(ctx context.Context, tenantID int64, modelName string) (*Pric
 		Currency:           "USD",
 		CacheReadPrice:     cacheReadPrice,
 		CacheCreationPrice: cacheCreationPrice,
+		MaxOutputTokens:    model.MaxOutputTokens,
 		CustomTiers:        customTiers,
 	}
 
@@ -431,22 +438,13 @@ func EstimatePreDeductAmount(ctx context.Context, tenantID int64, modelName stri
 		return pricing.PerRequestPrice, nil
 	}
 
-	// Token 计费：估算
-	type modelRow struct {
-		MaxOutputTokens int `json:"max_output_tokens"`
-	}
-	var model *modelRow
-	err = dao.MdlModels.Ctx(ctx).
-		Where("model_id", modelName).
-		Fields("max_output_tokens").
-		Scan(&model)
-	if err != nil {
-		return 0, gerror.Wrapf(err, "estimate pre-deduct: query model output limit")
-	}
-
+	// Token 计费：估算输出上限。
+	// max_output_tokens 已随 GetModelPrice 的 600s 定价缓存一起取出，
+	// 不再单独查 mdl_models（原实现每请求一条无缓存 SQL，高并发下是纯浪费）。
+	// 为 0（模型未设置，或旧缓存条目缺此字段）时按 4096 兜底。
 	maxOutput := 4096
-	if model != nil && model.MaxOutputTokens > 0 {
-		maxOutput = model.MaxOutputTokens
+	if pricing.MaxOutputTokens > 0 {
+		maxOutput = pricing.MaxOutputTokens
 	}
 
 	estimatedOutput := requestedMaxTokens

internal/logic/billing/ratelimit.go

@@ -3,6 +3,7 @@ package billing
 import (
 	"context"
 	"fmt"
+	"sync/atomic"
 	"time"
 
 	"github.com/gogf/gf/v2/frame/g"
@@ -38,7 +39,16 @@ var DefaultRateLimitConfig = RateLimitConfig{
 }
 
 // LoadRateLimitConfig loads rate limit config from ConfigService, falling back to defaults.
+//
+// 热路径优化：CheckRateLimit（阶段 2）和 AcquireConcurrent（阶段 3）每请求各调用一次本函数，
+// 每次做 5 个配置 key 查找（各自含缓存锁与可能的 L2/DB 回源）。限流配置是全局低频变更数据，
+// 用进程级 5s 快照把高并发下的重复组装收敛为一次原子读；后台修改限流配置最多延迟 5s 生效，
+// 对限流阈值这类"软控制线"可接受。
 func LoadRateLimitConfig(ctx context.Context) RateLimitConfig {
+	if snap := rateLimitConfigSnapshot.Load(); snap != nil && time.Now().Before(snap.expireAt) {
+		return snap.config
+	}
+
 	cfg := common.Config()
 	c := DefaultRateLimitConfig
 
@@ -58,9 +68,22 @@ func LoadRateLimitConfig(ctx context.Context) RateLimitConfig {
 		c.TenantConc = v
 	}
 
+	rateLimitConfigSnapshot.Store(&rateLimitConfigEntry{
+		config:   c,
+		expireAt: time.Now().Add(5 * time.Second),
+	})
 	return c
 }
 
+// rateLimitConfigEntry 限流配置的进程级快照条目
+type rateLimitConfigEntry struct {
+	config   RateLimitConfig
+	expireAt time.Time
+}
+
+// rateLimitConfigSnapshot 限流配置快照（5s TTL，过期后由下一个请求刷新）
+var rateLimitConfigSnapshot atomic.Pointer[rateLimitConfigEntry]
+
 // ---------------------------------------------------------------------------
 // QPS 限流：单个 Lua 脚本完成 4 级检查
 // ---------------------------------------------------------------------------

relay/handler/realtime_handler.go

@@ -84,6 +84,43 @@ type RealtimeContext struct {
 func HandleRealtime(w http.ResponseWriter, r *http.Request, rc *RealtimeContext, provider common.DataProvider, billing common.BillingProvider) (*common.Usage, *BillingResult, error) {
 	ctx := r.Context()
 
+	// 0. 前置检查（scope / IP 白名单 / QPS 限流）：只依赖认证上下文，放在 WebSocket
+	// 升级之前——升级后连接已被劫持，无法再写标准 HTTP 状态码；前置后被拒请求能拿到
+	// 正常的 401/429 响应，也省掉无谓的升级握手和首消息读取。
+	// 注意：上层调用方（internal/handler/relay）对 realtime 的返回错误只记日志、不写响应
+	//（因升级后的错误无法写回），故升级前的拒绝必须在此处自行写回 HTTP 错误响应，
+	// 否则客户端会收到空 body 的 200。
+	billingResult := &BillingResult{}
+	if billing != nil {
+		if !billing.CheckScope(rc.Scope, "realtime") {
+			err := constant.NewAuthError("API key scope denied")
+			WriteRelayError(w, err)
+			return nil, billingResult, err
+		}
+		if !billing.CheckIPWhitelist(rc.KeyIpWhitelist, rc.ClientIP) {
+			err := constant.NewAuthError("IP address is not allowed")
+			WriteRelayError(w, err)
+			return nil, billingResult, err
+		}
+		allowed, limitLevel, limit, remaining, resetAt := billing.CheckRateLimit(ctx, rc.TenantID, rc.UserID, rc.ApiKeyID, rc.KeyRateLimitQps)
+		if !allowed {
+			err := &RelayErrorWithRateLimit{
+				StatusCode: 429,
+				Message:    fmt.Sprintf("rate limit exceeded at %s level", limitLevel),
+				LimitLevel: limitLevel,
+				Remaining:  remaining,
+				ResetAt:    resetAt,
+			}
+			WriteRelayError(w, err)
+			return nil, billingResult, err
+		}
+		billingResult.RateLimitInfo = &common.RateLimitInfo{
+			Limit:     limit,
+			Remaining: remaining,
+			ResetAt:   resetAt,
+		}
+	}
+
 	// 1. WebSocket 升级
 	clientConn, err := websocketUpgrader.Upgrade(w, r, nil)
 	if err != nil {
@@ -113,31 +150,6 @@ func HandleRealtime(w http.ResponseWriter, r *http.Request, rc *RealtimeContext,
 		return nil, nil, fmt.Errorf("model not found in first websocket message")
 	}
 
-	billingResult := &BillingResult{}
-	if billing != nil {
-		if !billing.CheckScope(rc.Scope, "realtime") {
-			return nil, billingResult, constant.NewAuthError("API key scope denied")
-		}
-		if !billing.CheckIPWhitelist(rc.KeyIpWhitelist, rc.ClientIP) {
-			return nil, billingResult, constant.NewAuthError("IP address is not allowed")
-		}
-		allowed, limitLevel, limit, remaining, resetAt := billing.CheckRateLimit(ctx, rc.TenantID, rc.UserID, rc.ApiKeyID, rc.KeyRateLimitQps)
-		if !allowed {
-			return nil, billingResult, &RelayErrorWithRateLimit{
-				StatusCode: 429,
-				Message:    fmt.Sprintf("rate limit exceeded at %s level", limitLevel),
-				LimitLevel: limitLevel,
-				Remaining:  remaining,
-				ResetAt:    resetAt,
-			}
-		}
-		billingResult.RateLimitInfo = &common.RateLimitInfo{
-			Limit:     limit,
-			Remaining: remaining,
-			ResetAt:   resetAt,
-		}
-	}
-
 	// 3. 验证模型
 	_, _, err = provider.GetModelMapping(ctx, modelName)
 	if err != nil {
@@ -247,13 +259,6 @@ func HandleRealtime(w http.ResponseWriter, r *http.Request, rc *RealtimeContext,
 		}
 		defer billing.ReleaseApiKeyConcurrent(ctx, rc.ApiKeyID)
 
-		if err := billing.CheckMemberQuota(ctx, rc.TenantID, rc.UserID, 0); err != nil {
-			return nil, billingResult, constant.NewQuotaError("member quota exceeded", err)
-		}
-		if err := billing.CheckApiKeyQuota(ctx, rc.ApiKeyID, 0); err != nil {
-			return nil, billingResult, constant.NewQuotaError("API key quota exceeded", err)
-		}
-
 		amt, billErr := billing.PreDeduct(ctx, rc.TenantID, modelName, 0, 0, false, rc.RequestID)
 		if billErr != nil {
 			// 未配价模型 fail-closed：返回明确的请求错误，而非误导性的"余额不足"
@@ -263,11 +268,12 @@ func HandleRealtime(w http.ResponseWriter, r *http.Request, rc *RealtimeContext,
 			return nil, nil, constant.NewQuotaError("insufficient balance", billErr)
 		}
 		preDeductAmount = amt
+		// 成员/Key 额度检查只做这一次（带实际预扣额），与 RelayHandler 一致：
+		// 预扣前的 Check(0) 快速闸门对全量请求多付 Redis 往返，收益不成比例，已移除。
 		if err := billing.CheckApiKeyQuota(ctx, rc.ApiKeyID, amt); err != nil {
 			_ = billing.SettleFailed(ctx, rc.TenantID, rc.RequestID, amt)
 			return nil, billingResult, constant.NewQuotaError("API key quota exceeded", err)
 		}
-		// 按实际预扣额复查成员额度（上方 Check(0) 只是冻结前的快速闸门）
 		if err := billing.CheckMemberQuota(ctx, rc.TenantID, rc.UserID, amt); err != nil {
 			_ = billing.SettleFailed(ctx, rc.TenantID, rc.RequestID, amt)
 			return nil, billingResult, constant.NewQuotaError("member quota exceeded", err)

relay/handler/relay_handler.go

@@ -67,7 +67,7 @@ type relayValidation struct {
 	sessionSignals  dispatch.SessionSignals // 请求体中的会话信号（影子模式/新调度用）
 }
 
-// validateRelayRequest 校验请求合法性：relay mode、模型存在性、弃用状态、成员/API Key 模型范围、QPS 限流。
+// validateRelayRequest 校验请求合法性：relay mode、QPS 限流（前置）、模型存在性、弃用状态、成员/API Key 模型范围。
 // 纯校验逻辑，无 defer 副作用。
 func validateRelayRequest(
 	ctx context.Context,
@@ -84,6 +84,29 @@ func validateRelayRequest(
 		return nil, constant.NewRequestError("unsupported endpoint: "+path, nil)
 	}
 
+	// 1.5 QPS 限流检查（尽可能前置：只依赖认证阶段就已就绪的 tenant/user/key ID，
+	// 提前到请求体解析和模型权限检查之前，让超限请求以 1 次 Redis EVAL 的成本被拒绝，
+	// 不再为其解析大 JSON、查模型缓存。副作用：无效请求（坏 body/无权限模型）也计入
+	// QPS 计数——这是期望行为，探测型流量同样应被限流兜住。
+	billingResult := &BillingResult{}
+	if billing != nil {
+		allowed, limitLevel, limit, remaining, resetAt := billing.CheckRateLimit(ctx, rc.TenantID, rc.UserID, rc.ApiKeyID, rc.KeyRateLimitQps)
+		if !allowed {
+			return nil, &RelayErrorWithRateLimit{
+				StatusCode: 429,
+				Message:    fmt.Sprintf("rate limit exceeded at %s level", limitLevel),
+				LimitLevel: limitLevel,
+				Remaining:  remaining,
+				ResetAt:    resetAt,
+			}
+		}
+		billingResult.RateLimitInfo = &common.RateLimitInfo{
+			Limit:     limit,
+			Remaining: remaining,
+			ResetAt:   resetAt,
+		}
+	}
+
 	relayModeStr := relayModeString(relayMode)
 	if billing != nil {
 		if !billing.CheckScope(rc.Scope, relayModeStr) {
@@ -143,6 +166,7 @@ func validateRelayRequest(
 			return nil, constant.NewModelGoneError(lookupModel, depInfo.SunsetDate)
 		}
 	}
+	billingResult.Deprecation = depInfo
 
 	// 3.7 检查成员模型范围
 	if allowed, err := provider.CheckMemberModelAccess(ctx, rc.TenantID, rc.UserID, lookupModel); err != nil {
@@ -158,25 +182,7 @@ func validateRelayRequest(
 		return nil, constant.NewAuthError("model not allowed for this API key")
 	}
 
-	// 5. QPS 限流检查
-	billingResult := &BillingResult{Deprecation: depInfo}
-	if billing != nil {
-		allowed, limitLevel, limit, remaining, resetAt := billing.CheckRateLimit(ctx, rc.TenantID, rc.UserID, rc.ApiKeyID, rc.KeyRateLimitQps)
-		if !allowed {
-			return nil, &RelayErrorWithRateLimit{
-				StatusCode: 429,
-				Message:    fmt.Sprintf("rate limit exceeded at %s level", limitLevel),
-				LimitLevel: limitLevel,
-				Remaining:  remaining,
-				ResetAt:    resetAt,
-			}
-		}
-		billingResult.RateLimitInfo = &common.RateLimitInfo{
-			Limit:     limit,
-			Remaining: remaining,
-			ResetAt:   resetAt,
-		}
-	}
+	// 5. QPS 限流已前置到步骤 1.5（见上），此处不再检查
 
 	// 7. 估算输入 token 数
 	estimatedInputTokens := estimateInputTokens(body)
@@ -406,15 +412,6 @@ func RelayHandler(ctx context.Context, body []byte, path string, headers http.He
 	})
 	defer monitor.UnregisterRequest(rc.RequestID)
 
-	if billing != nil {
-		if err := billing.CheckMemberQuota(ctx, rc.TenantID, rc.UserID, 0); err != nil {
-			return nil, nil, constant.NewQuotaError("member quota exceeded", err)
-		}
-		if err := billing.CheckApiKeyQuota(ctx, rc.ApiKeyID, 0); err != nil {
-			return nil, nil, constant.NewQuotaError("API key quota exceeded", err)
-		}
-	}
-
 	var preDeductAmount float64
 	if billing != nil {
 		amt, err := billing.PreDeduct(ctx, rc.TenantID, v.modelName, v.estimatedTokens, v.maxTokens, v.isStream, rc.RequestID)
@@ -427,12 +424,14 @@ func RelayHandler(ctx context.Context, body []byte, path string, headers http.He
 		}
 		preDeductAmount = amt
 		v.billingResult.PreDeductAmount = amt
+		// 成员/Key 额度检查只做这一次（带实际预扣额）：额度是"控制线"计数器而非资源冻结，
+		// 预扣前后检查的竞态语义相同。原先预扣前额外的 Check(0) 快速闸门对全量请求
+		// 各多付一次 Redis 往返，只为让"额度已耗尽"的请求少走一步预扣+退款，得不偿失
+		//（该场景已被上游 QPS 限流兜底），故移除。
 		if err := billing.CheckApiKeyQuota(ctx, rc.ApiKeyID, amt); err != nil {
 			_ = billing.SettleFailed(ctx, rc.TenantID, rc.RequestID, amt)
 			return nil, nil, constant.NewQuotaError("API key quota exceeded", err)
 		}
-		// 按实际预扣额复查成员额度（前置的 Check(0) 只是冻结前的快速闸门，
-		// 不带金额无法拦截"当前用量在限额内、本笔请求会超限"的场景）
 		if err := billing.CheckMemberQuota(ctx, rc.TenantID, rc.UserID, amt); err != nil {
 			_ = billing.SettleFailed(ctx, rc.TenantID, rc.RequestID, amt)
 			return nil, nil, constant.NewQuotaError("member quota exceeded", err)

relay/handler/task_handler.go

@@ -92,6 +92,13 @@ func HandleTaskSubmit(
 	billingProvider common.TaskBillingProvider,
 	channelMeta *common.ChannelMeta,
 ) {
+	// 0. QPS 限流检查（前置：只依赖认证上下文，超限请求在解析请求体之前被拒绝）
+	allowed, limitLevel, _, _, _ := billingProvider.CheckRateLimit(ctx, rc.TenantID, rc.UserID, rc.ApiKeyID, rc.KeyRateLimitQps)
+	if !allowed {
+		writeTaskError(rc.Writer, http.StatusTooManyRequests, fmt.Sprintf("rate limit exceeded at %s level", limitLevel), "")
+		return
+	}
+
 	// 1. 解析模型名
 	var req map[string]json.RawMessage
 	if err := json.Unmarshal(body, &req); err != nil {
@@ -121,11 +128,6 @@ func HandleTaskSubmit(
 		writeTaskError(rc.Writer, http.StatusForbidden, "IP address is not allowed", "")
 		return
 	}
-	allowed, limitLevel, _, _, _ := billingProvider.CheckRateLimit(ctx, rc.TenantID, rc.UserID, rc.ApiKeyID, rc.KeyRateLimitQps)
-	if !allowed {
-		writeTaskError(rc.Writer, http.StatusTooManyRequests, fmt.Sprintf("rate limit exceeded at %s level", limitLevel), "")
-		return
-	}
 	if !billingProvider.AcquireApiKeyConcurrent(ctx, rc.ApiKeyID, rc.KeyConcurrency) {
 		writeTaskError(rc.Writer, http.StatusTooManyRequests, "API key concurrent request limit exceeded", "")
 		return