diff --git a/internal/logic/billing/pricing.go b/internal/logic/billing/pricing.go index faf7e13..5a26fd1 100644 --- a/internal/logic/billing/pricing.go +++ b/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 diff --git a/internal/logic/billing/ratelimit.go b/internal/logic/billing/ratelimit.go index 1143deb..da78aaf 100644 --- a/internal/logic/billing/ratelimit.go +++ b/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 级检查 // --------------------------------------------------------------------------- diff --git a/relay/handler/realtime_handler.go b/relay/handler/realtime_handler.go index 0518c25..7379551 100644 --- a/relay/handler/realtime_handler.go +++ b/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) diff --git a/relay/handler/relay_handler.go b/relay/handler/relay_handler.go index d48e762..57d32db 100644 --- a/relay/handler/relay_handler.go +++ b/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) diff --git a/relay/handler/task_handler.go b/relay/handler/task_handler.go index 95839bc..e7fcfb4 100644 --- a/relay/handler/task_handler.go +++ b/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