fix: address deep code review issues for RPM limiting
- Move IncrementRPM after Forward success to prevent phantom RPM consumption during account switch retries - Add base_rpm input sanitization (clamp to 0-10000) in Create/Update - Add WindowCost scheduling checks to legacy path sticky sessions (4 check sites + 4 prefetch sites), fixing pre-existing gap - Clean up rpm_strategy/rpm_sticky_buffer when disabling RPM in BulkEditModal (JSONB merge cannot delete keys, use empty values) - Add json.Number test cases to TestGetBaseRPM/TestGetRPMStickyBuffer - Document TOCTOU race as accepted soft-limit design trade-off
This commit is contained in:
@@ -483,6 +483,8 @@ func (h *AccountHandler) Create(c *gin.Context) {
|
||||
response.BadRequest(c, "rate_multiplier must be >= 0")
|
||||
return
|
||||
}
|
||||
// base_rpm 输入校验:负值归零,超过 10000 截断
|
||||
sanitizeExtraBaseRPM(req.Extra)
|
||||
|
||||
// 确定是否跳过混合渠道检查
|
||||
skipCheck := req.ConfirmMixedChannelRisk != nil && *req.ConfirmMixedChannelRisk
|
||||
@@ -552,6 +554,8 @@ func (h *AccountHandler) Update(c *gin.Context) {
|
||||
response.BadRequest(c, "rate_multiplier must be >= 0")
|
||||
return
|
||||
}
|
||||
// base_rpm 输入校验:负值归零,超过 10000 截断
|
||||
sanitizeExtraBaseRPM(req.Extra)
|
||||
|
||||
// 确定是否跳过混合渠道检查
|
||||
skipCheck := req.ConfirmMixedChannelRisk != nil && *req.ConfirmMixedChannelRisk
|
||||
@@ -1736,3 +1740,44 @@ func (h *AccountHandler) BatchRefreshTier(c *gin.Context) {
|
||||
func (h *AccountHandler) GetAntigravityDefaultModelMapping(c *gin.Context) {
|
||||
response.Success(c, domain.DefaultAntigravityModelMapping)
|
||||
}
|
||||
|
||||
// sanitizeExtraBaseRPM 对 extra map 中的 base_rpm 值进行范围校验和归一化。
|
||||
// 负值归零,超过 10000 截断为 10000。extra 为 nil 或不含 base_rpm 时无操作。
|
||||
func sanitizeExtraBaseRPM(extra map[string]any) {
|
||||
if extra == nil {
|
||||
return
|
||||
}
|
||||
raw, ok := extra["base_rpm"]
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
v := parseExtraIntForValidation(raw)
|
||||
if v < 0 {
|
||||
v = 0
|
||||
} else if v > 10000 {
|
||||
v = 10000
|
||||
}
|
||||
extra["base_rpm"] = v
|
||||
}
|
||||
|
||||
// parseExtraIntForValidation 从 extra 字段的 any 值解析为 int,用于输入校验。
|
||||
// 支持 int, int64, float64, json.Number, string 类型。
|
||||
func parseExtraIntForValidation(value any) int {
|
||||
switch v := value.(type) {
|
||||
case int:
|
||||
return v
|
||||
case int64:
|
||||
return int(v)
|
||||
case float64:
|
||||
return int(v)
|
||||
case json.Number:
|
||||
if i, err := v.Int64(); err == nil {
|
||||
return int(i)
|
||||
}
|
||||
case string:
|
||||
if i, err := strconv.Atoi(strings.TrimSpace(v)); err == nil {
|
||||
return i
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
@@ -366,13 +366,6 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
// 账号槽位/等待计数需要在超时或断开时安全回收
|
||||
accountReleaseFunc = wrapReleaseOnDone(c.Request.Context(), accountReleaseFunc)
|
||||
|
||||
// RPM 计数递增(调度成功后、Forward 前)
|
||||
if account.IsAnthropicOAuthOrSetupToken() && account.GetBaseRPM() > 0 {
|
||||
if err := h.gatewayService.IncrementAccountRPM(c.Request.Context(), account.ID); err != nil {
|
||||
reqLog.Warn("gateway.rpm_increment_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
// 转发请求 - 根据账号平台分流
|
||||
var result *service.ForwardResult
|
||||
requestCtx := c.Request.Context()
|
||||
@@ -410,6 +403,15 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// RPM 计数递增(Forward 成功后)
|
||||
// 注意:TOCTOU 竞态是已知且可接受的设计权衡,与 WindowCost 一致的 soft-limit 模式。
|
||||
// 在高并发下可能短暂超出 RPM 限制,但不会导致请求失败。
|
||||
if account.IsAnthropicOAuthOrSetupToken() && account.GetBaseRPM() > 0 {
|
||||
if err := h.gatewayService.IncrementAccountRPM(c.Request.Context(), account.ID); err != nil {
|
||||
reqLog.Warn("gateway.rpm_increment_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
// 捕获请求信息(用于异步记录,避免在 goroutine 中访问 gin.Context)
|
||||
userAgent := c.GetHeader("User-Agent")
|
||||
clientIP := ip.GetClientIP(c)
|
||||
@@ -556,13 +558,6 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
// 账号槽位/等待计数需要在超时或断开时安全回收
|
||||
accountReleaseFunc = wrapReleaseOnDone(c.Request.Context(), accountReleaseFunc)
|
||||
|
||||
// RPM 计数递增(调度成功后、Forward 前)
|
||||
if account.IsAnthropicOAuthOrSetupToken() && account.GetBaseRPM() > 0 {
|
||||
if err := h.gatewayService.IncrementAccountRPM(c.Request.Context(), account.ID); err != nil {
|
||||
reqLog.Warn("gateway.rpm_increment_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
// 转发请求 - 根据账号平台分流
|
||||
var result *service.ForwardResult
|
||||
requestCtx := c.Request.Context()
|
||||
@@ -609,7 +604,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
h.handleStreamingAwareError(c, status, code, message, streamStarted)
|
||||
return
|
||||
}
|
||||
// 兜底重试按“直接请求兜底分组”处理:清除强制平台,允许按分组平台调度
|
||||
// 兜底重试按"直接请求兜底分组"处理:清除强制平台,允许按分组平台调度
|
||||
ctx := context.WithValue(c.Request.Context(), ctxkey.ForcePlatform, "")
|
||||
c.Request = c.Request.WithContext(ctx)
|
||||
currentAPIKey = fallbackAPIKey
|
||||
@@ -643,6 +638,15 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// RPM 计数递增(Forward 成功后)
|
||||
// 注意:TOCTOU 竞态是已知且可接受的设计权衡,与 WindowCost 一致的 soft-limit 模式。
|
||||
// 在高并发下可能短暂超出 RPM 限制,但不会导致请求失败。
|
||||
if account.IsAnthropicOAuthOrSetupToken() && account.GetBaseRPM() > 0 {
|
||||
if err := h.gatewayService.IncrementAccountRPM(c.Request.Context(), account.ID); err != nil {
|
||||
reqLog.Warn("gateway.rpm_increment_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
// 捕获请求信息(用于异步记录,避免在 goroutine 中访问 gin.Context)
|
||||
userAgent := c.GetHeader("User-Agent")
|
||||
clientIP := ip.GetClientIP(c)
|
||||
|
||||
Reference in New Issue
Block a user