Merge pull request #679 from DaydreamCoding/feat/account-rpm-limit
feat: 添加账号级别 RPM(每分钟请求数)限流功能
This commit is contained in:
@@ -1137,6 +1137,80 @@ func (a *Account) GetSessionIdleTimeoutMinutes() int {
|
||||
return 5
|
||||
}
|
||||
|
||||
// GetBaseRPM 获取基础 RPM 限制
|
||||
// 返回 0 表示未启用(负数视为无效配置,按 0 处理)
|
||||
func (a *Account) GetBaseRPM() int {
|
||||
if a.Extra == nil {
|
||||
return 0
|
||||
}
|
||||
if v, ok := a.Extra["base_rpm"]; ok {
|
||||
val := parseExtraInt(v)
|
||||
if val > 0 {
|
||||
return val
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// GetRPMStrategy 获取 RPM 策略
|
||||
// "tiered" = 三区模型(默认), "sticky_exempt" = 粘性豁免
|
||||
func (a *Account) GetRPMStrategy() string {
|
||||
if a.Extra == nil {
|
||||
return "tiered"
|
||||
}
|
||||
if v, ok := a.Extra["rpm_strategy"]; ok {
|
||||
if s, ok := v.(string); ok && s == "sticky_exempt" {
|
||||
return "sticky_exempt"
|
||||
}
|
||||
}
|
||||
return "tiered"
|
||||
}
|
||||
|
||||
// GetRPMStickyBuffer 获取 RPM 粘性缓冲数量
|
||||
// tiered 模式下的黄区大小,默认为 base_rpm 的 20%(至少 1)
|
||||
func (a *Account) GetRPMStickyBuffer() int {
|
||||
if a.Extra == nil {
|
||||
return 0
|
||||
}
|
||||
if v, ok := a.Extra["rpm_sticky_buffer"]; ok {
|
||||
val := parseExtraInt(v)
|
||||
if val > 0 {
|
||||
return val
|
||||
}
|
||||
}
|
||||
base := a.GetBaseRPM()
|
||||
buffer := base / 5
|
||||
if buffer < 1 && base > 0 {
|
||||
buffer = 1
|
||||
}
|
||||
return buffer
|
||||
}
|
||||
|
||||
// CheckRPMSchedulability 根据当前 RPM 计数检查调度状态
|
||||
// 复用 WindowCostSchedulability 三态:Schedulable / StickyOnly / NotSchedulable
|
||||
func (a *Account) CheckRPMSchedulability(currentRPM int) WindowCostSchedulability {
|
||||
baseRPM := a.GetBaseRPM()
|
||||
if baseRPM <= 0 {
|
||||
return WindowCostSchedulable
|
||||
}
|
||||
|
||||
if currentRPM < baseRPM {
|
||||
return WindowCostSchedulable
|
||||
}
|
||||
|
||||
strategy := a.GetRPMStrategy()
|
||||
if strategy == "sticky_exempt" {
|
||||
return WindowCostStickyOnly // 粘性豁免无红区
|
||||
}
|
||||
|
||||
// tiered: 黄区 + 红区
|
||||
buffer := a.GetRPMStickyBuffer()
|
||||
if currentRPM < baseRPM+buffer {
|
||||
return WindowCostStickyOnly
|
||||
}
|
||||
return WindowCostNotSchedulable
|
||||
}
|
||||
|
||||
// CheckWindowCostSchedulability 根据当前窗口费用检查调度状态
|
||||
// - 费用 < 阈值: WindowCostSchedulable(可正常调度)
|
||||
// - 费用 >= 阈值 且 < 阈值+预留: WindowCostStickyOnly(仅粘性会话)
|
||||
@@ -1200,6 +1274,12 @@ func parseExtraFloat64(value any) float64 {
|
||||
}
|
||||
|
||||
// parseExtraInt 从 extra 字段解析 int 值
|
||||
// ParseExtraInt 从 extra 字段的 any 值解析为 int。
|
||||
// 支持 int, int64, float64, json.Number, string 类型,无法解析时返回 0。
|
||||
func ParseExtraInt(value any) int {
|
||||
return parseExtraInt(value)
|
||||
}
|
||||
|
||||
func parseExtraInt(value any) int {
|
||||
switch v := value.(type) {
|
||||
case int:
|
||||
|
||||
120
backend/internal/service/account_rpm_test.go
Normal file
120
backend/internal/service/account_rpm_test.go
Normal file
@@ -0,0 +1,120 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestGetBaseRPM(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
extra map[string]any
|
||||
expected int
|
||||
}{
|
||||
{"nil extra", nil, 0},
|
||||
{"no key", map[string]any{}, 0},
|
||||
{"zero", map[string]any{"base_rpm": 0}, 0},
|
||||
{"int value", map[string]any{"base_rpm": 15}, 15},
|
||||
{"float value", map[string]any{"base_rpm": 15.0}, 15},
|
||||
{"string value", map[string]any{"base_rpm": "15"}, 15},
|
||||
{"negative value", map[string]any{"base_rpm": -5}, 0},
|
||||
{"int64 value", map[string]any{"base_rpm": int64(20)}, 20},
|
||||
{"json.Number value", map[string]any{"base_rpm": json.Number("25")}, 25},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
a := &Account{Extra: tt.extra}
|
||||
if got := a.GetBaseRPM(); got != tt.expected {
|
||||
t.Errorf("GetBaseRPM() = %d, want %d", got, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetRPMStrategy(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
extra map[string]any
|
||||
expected string
|
||||
}{
|
||||
{"nil extra", nil, "tiered"},
|
||||
{"no key", map[string]any{}, "tiered"},
|
||||
{"tiered", map[string]any{"rpm_strategy": "tiered"}, "tiered"},
|
||||
{"sticky_exempt", map[string]any{"rpm_strategy": "sticky_exempt"}, "sticky_exempt"},
|
||||
{"invalid", map[string]any{"rpm_strategy": "foobar"}, "tiered"},
|
||||
{"empty string fallback", map[string]any{"rpm_strategy": ""}, "tiered"},
|
||||
{"numeric value fallback", map[string]any{"rpm_strategy": 123}, "tiered"},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
a := &Account{Extra: tt.extra}
|
||||
if got := a.GetRPMStrategy(); got != tt.expected {
|
||||
t.Errorf("GetRPMStrategy() = %q, want %q", got, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestCheckRPMSchedulability(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
extra map[string]any
|
||||
currentRPM int
|
||||
expected WindowCostSchedulability
|
||||
}{
|
||||
{"disabled", map[string]any{}, 100, WindowCostSchedulable},
|
||||
{"green zone", map[string]any{"base_rpm": 15}, 10, WindowCostSchedulable},
|
||||
{"yellow zone tiered", map[string]any{"base_rpm": 15}, 15, WindowCostStickyOnly},
|
||||
{"red zone tiered", map[string]any{"base_rpm": 15}, 18, WindowCostNotSchedulable},
|
||||
{"sticky_exempt at limit", map[string]any{"base_rpm": 15, "rpm_strategy": "sticky_exempt"}, 15, WindowCostStickyOnly},
|
||||
{"sticky_exempt over limit", map[string]any{"base_rpm": 15, "rpm_strategy": "sticky_exempt"}, 100, WindowCostStickyOnly},
|
||||
{"custom buffer", map[string]any{"base_rpm": 10, "rpm_sticky_buffer": 5}, 14, WindowCostStickyOnly},
|
||||
{"custom buffer red", map[string]any{"base_rpm": 10, "rpm_sticky_buffer": 5}, 15, WindowCostNotSchedulable},
|
||||
{"base_rpm=1 green", map[string]any{"base_rpm": 1}, 0, WindowCostSchedulable},
|
||||
{"base_rpm=1 yellow (at limit)", map[string]any{"base_rpm": 1}, 1, WindowCostStickyOnly},
|
||||
{"base_rpm=1 red (at limit+buffer)", map[string]any{"base_rpm": 1}, 2, WindowCostNotSchedulable},
|
||||
{"negative currentRPM", map[string]any{"base_rpm": 15}, -1, WindowCostSchedulable},
|
||||
{"base_rpm negative disabled", map[string]any{"base_rpm": -5}, 10, WindowCostSchedulable},
|
||||
{"very high currentRPM", map[string]any{"base_rpm": 10}, 9999, WindowCostNotSchedulable},
|
||||
{"sticky_exempt very high currentRPM", map[string]any{"base_rpm": 10, "rpm_strategy": "sticky_exempt"}, 9999, WindowCostStickyOnly},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
a := &Account{Extra: tt.extra}
|
||||
if got := a.CheckRPMSchedulability(tt.currentRPM); got != tt.expected {
|
||||
t.Errorf("CheckRPMSchedulability(%d) = %d, want %d", tt.currentRPM, got, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetRPMStickyBuffer(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
extra map[string]any
|
||||
expected int
|
||||
}{
|
||||
{"nil extra", nil, 0},
|
||||
{"no keys", map[string]any{}, 0},
|
||||
{"base_rpm=0", map[string]any{"base_rpm": 0}, 0},
|
||||
{"base_rpm=1 min buffer 1", map[string]any{"base_rpm": 1}, 1},
|
||||
{"base_rpm=4 min buffer 1", map[string]any{"base_rpm": 4}, 1},
|
||||
{"base_rpm=5 buffer 1", map[string]any{"base_rpm": 5}, 1},
|
||||
{"base_rpm=10 buffer 2", map[string]any{"base_rpm": 10}, 2},
|
||||
{"base_rpm=15 buffer 3", map[string]any{"base_rpm": 15}, 3},
|
||||
{"base_rpm=100 buffer 20", map[string]any{"base_rpm": 100}, 20},
|
||||
{"custom buffer=5", map[string]any{"base_rpm": 10, "rpm_sticky_buffer": 5}, 5},
|
||||
{"custom buffer=0 fallback to default", map[string]any{"base_rpm": 10, "rpm_sticky_buffer": 0}, 2},
|
||||
{"custom buffer negative fallback", map[string]any{"base_rpm": 10, "rpm_sticky_buffer": -1}, 2},
|
||||
{"custom buffer with float", map[string]any{"base_rpm": 10, "rpm_sticky_buffer": float64(7)}, 7},
|
||||
{"json.Number base_rpm", map[string]any{"base_rpm": json.Number("10")}, 2},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
a := &Account{Extra: tt.extra}
|
||||
if got := a.GetRPMStickyBuffer(); got != tt.expected {
|
||||
t.Errorf("GetRPMStickyBuffer() = %d, want %d", got, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -520,6 +520,7 @@ type GatewayService struct {
|
||||
concurrencyService *ConcurrencyService
|
||||
claudeTokenProvider *ClaudeTokenProvider
|
||||
sessionLimitCache SessionLimitCache // 会话数量限制缓存(仅 Anthropic OAuth/SetupToken)
|
||||
rpmCache RPMCache // RPM 计数缓存(仅 Anthropic OAuth/SetupToken)
|
||||
userGroupRateCache *gocache.Cache
|
||||
userGroupRateSF singleflight.Group
|
||||
modelsListCache *gocache.Cache
|
||||
@@ -549,6 +550,7 @@ func NewGatewayService(
|
||||
deferredService *DeferredService,
|
||||
claudeTokenProvider *ClaudeTokenProvider,
|
||||
sessionLimitCache SessionLimitCache,
|
||||
rpmCache RPMCache,
|
||||
digestStore *DigestSessionStore,
|
||||
) *GatewayService {
|
||||
userGroupRateTTL := resolveUserGroupRateCacheTTL(cfg)
|
||||
@@ -574,6 +576,7 @@ func NewGatewayService(
|
||||
deferredService: deferredService,
|
||||
claudeTokenProvider: claudeTokenProvider,
|
||||
sessionLimitCache: sessionLimitCache,
|
||||
rpmCache: rpmCache,
|
||||
userGroupRateCache: gocache.New(userGroupRateTTL, time.Minute),
|
||||
modelsListCache: gocache.New(modelsListTTL, time.Minute),
|
||||
modelsListCacheTTL: modelsListTTL,
|
||||
@@ -1154,6 +1157,7 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
|
||||
return nil, errors.New("no available accounts")
|
||||
}
|
||||
ctx = s.withWindowCostPrefetch(ctx, accounts)
|
||||
ctx = s.withRPMPrefetch(ctx, accounts)
|
||||
|
||||
isExcluded := func(accountID int64) bool {
|
||||
if excludedIDs == nil {
|
||||
@@ -1229,6 +1233,10 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
|
||||
filteredWindowCost++
|
||||
continue
|
||||
}
|
||||
// RPM 检查(非粘性会话路径)
|
||||
if !s.isAccountSchedulableForRPM(ctx, account, false) {
|
||||
continue
|
||||
}
|
||||
routingCandidates = append(routingCandidates, account)
|
||||
}
|
||||
|
||||
@@ -1252,7 +1260,9 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
|
||||
s.isAccountAllowedForPlatform(stickyAccount, platform, useMixed) &&
|
||||
(requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, stickyAccount, requestedModel)) &&
|
||||
s.isAccountSchedulableForModelSelection(ctx, stickyAccount, requestedModel) &&
|
||||
s.isAccountSchedulableForWindowCost(ctx, stickyAccount, true) { // 粘性会话窗口费用检查
|
||||
s.isAccountSchedulableForWindowCost(ctx, stickyAccount, true) &&
|
||||
|
||||
s.isAccountSchedulableForRPM(ctx, stickyAccount, true) { // 粘性会话窗口费用+RPM 检查
|
||||
result, err := s.tryAcquireAccountSlot(ctx, stickyAccountID, stickyAccount.Concurrency)
|
||||
if err == nil && result.Acquired {
|
||||
// 会话数量限制检查
|
||||
@@ -1406,7 +1416,9 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
|
||||
s.isAccountAllowedForPlatform(account, platform, useMixed) &&
|
||||
(requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) &&
|
||||
s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) &&
|
||||
s.isAccountSchedulableForWindowCost(ctx, account, true) { // 粘性会话窗口费用检查
|
||||
s.isAccountSchedulableForWindowCost(ctx, account, true) &&
|
||||
|
||||
s.isAccountSchedulableForRPM(ctx, account, true) { // 粘性会话窗口费用+RPM 检查
|
||||
result, err := s.tryAcquireAccountSlot(ctx, accountID, account.Concurrency)
|
||||
if err == nil && result.Acquired {
|
||||
// 会话数量限制检查
|
||||
@@ -1472,6 +1484,10 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
|
||||
if !s.isAccountSchedulableForWindowCost(ctx, acc, false) {
|
||||
continue
|
||||
}
|
||||
// RPM 检查(非粘性会话路径)
|
||||
if !s.isAccountSchedulableForRPM(ctx, acc, false) {
|
||||
continue
|
||||
}
|
||||
candidates = append(candidates, acc)
|
||||
}
|
||||
|
||||
@@ -2155,6 +2171,88 @@ checkSchedulability:
|
||||
return true
|
||||
}
|
||||
|
||||
// rpmPrefetchContextKey is the context key for prefetched RPM counts.
|
||||
type rpmPrefetchContextKeyType struct{}
|
||||
|
||||
var rpmPrefetchContextKey = rpmPrefetchContextKeyType{}
|
||||
|
||||
func rpmFromPrefetchContext(ctx context.Context, accountID int64) (int, bool) {
|
||||
if v, ok := ctx.Value(rpmPrefetchContextKey).(map[int64]int); ok {
|
||||
count, found := v[accountID]
|
||||
return count, found
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
|
||||
// withRPMPrefetch 批量预取所有候选账号的 RPM 计数
|
||||
func (s *GatewayService) withRPMPrefetch(ctx context.Context, accounts []Account) context.Context {
|
||||
if s.rpmCache == nil {
|
||||
return ctx
|
||||
}
|
||||
|
||||
var ids []int64
|
||||
for i := range accounts {
|
||||
if accounts[i].IsAnthropicOAuthOrSetupToken() && accounts[i].GetBaseRPM() > 0 {
|
||||
ids = append(ids, accounts[i].ID)
|
||||
}
|
||||
}
|
||||
if len(ids) == 0 {
|
||||
return ctx
|
||||
}
|
||||
|
||||
counts, err := s.rpmCache.GetRPMBatch(ctx, ids)
|
||||
if err != nil {
|
||||
return ctx // 失败开放
|
||||
}
|
||||
return context.WithValue(ctx, rpmPrefetchContextKey, counts)
|
||||
}
|
||||
|
||||
// isAccountSchedulableForRPM 检查账号是否可根据 RPM 进行调度
|
||||
// 仅适用于 Anthropic OAuth/SetupToken 账号
|
||||
func (s *GatewayService) isAccountSchedulableForRPM(ctx context.Context, account *Account, isSticky bool) bool {
|
||||
if !account.IsAnthropicOAuthOrSetupToken() {
|
||||
return true
|
||||
}
|
||||
baseRPM := account.GetBaseRPM()
|
||||
if baseRPM <= 0 {
|
||||
return true
|
||||
}
|
||||
|
||||
// 尝试从预取缓存获取
|
||||
var currentRPM int
|
||||
if count, ok := rpmFromPrefetchContext(ctx, account.ID); ok {
|
||||
currentRPM = count
|
||||
} else if s.rpmCache != nil {
|
||||
if count, err := s.rpmCache.GetRPM(ctx, account.ID); err == nil {
|
||||
currentRPM = count
|
||||
}
|
||||
// 失败开放:GetRPM 错误时允许调度
|
||||
}
|
||||
|
||||
schedulability := account.CheckRPMSchedulability(currentRPM)
|
||||
switch schedulability {
|
||||
case WindowCostSchedulable:
|
||||
return true
|
||||
case WindowCostStickyOnly:
|
||||
return isSticky
|
||||
case WindowCostNotSchedulable:
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// IncrementAccountRPM increments the RPM counter for the given account.
|
||||
// 已知 TOCTOU 竞态:调度时读取 RPM 计数与此处递增之间存在时间窗口,
|
||||
// 高并发下可能短暂超出 RPM 限制。这是与 WindowCost 一致的 soft-limit
|
||||
// 设计权衡——可接受的少量超额优于加锁带来的延迟和复杂度。
|
||||
func (s *GatewayService) IncrementAccountRPM(ctx context.Context, accountID int64) error {
|
||||
if s.rpmCache == nil {
|
||||
return nil
|
||||
}
|
||||
_, err := s.rpmCache.IncrementRPM(ctx, accountID)
|
||||
return err
|
||||
}
|
||||
|
||||
// checkAndRegisterSession 检查并注册会话,用于会话数量限制
|
||||
// 仅适用于 Anthropic OAuth/SetupToken 账号
|
||||
// sessionID: 会话标识符(使用粘性会话的 hash)
|
||||
@@ -2349,7 +2447,7 @@ func sameAccountWithLoadGroup(a, b accountWithLoad) bool {
|
||||
// shuffleWithinPriorityAndLastUsed 对排序后的 []*Account 切片,按 (Priority, LastUsedAt) 分组后组内随机打乱。
|
||||
//
|
||||
// 注意:当 preferOAuth=true 时,需要保证 OAuth 账号在同组内仍然优先,否则会把排序时的偏好打散掉。
|
||||
// 因此这里采用“组内分区 + 分区内 shuffle”的方式:
|
||||
// 因此这里采用"组内分区 + 分区内 shuffle"的方式:
|
||||
// - 先把同组账号按 (OAuth / 非 OAuth) 拆成两段,保持 OAuth 段在前;
|
||||
// - 再分别在各段内随机打散,避免热点。
|
||||
func shuffleWithinPriorityAndLastUsed(accounts []*Account, preferOAuth bool) {
|
||||
@@ -2489,7 +2587,7 @@ func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context,
|
||||
if clearSticky {
|
||||
_ = s.cache.DeleteSessionAccountID(ctx, derefGroupID(groupID), sessionHash)
|
||||
}
|
||||
if !clearSticky && s.isAccountInGroup(account, groupID) && account.Platform == platform && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) {
|
||||
if !clearSticky && s.isAccountInGroup(account, groupID) && account.Platform == platform && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) && s.isAccountSchedulableForWindowCost(ctx, account, true) && s.isAccountSchedulableForRPM(ctx, account, true) {
|
||||
if s.debugModelRoutingEnabled() {
|
||||
logger.LegacyPrintf("service.gateway", "[ModelRoutingDebug] legacy routed sticky hit: group_id=%v model=%s session=%s account=%d", derefGroupID(groupID), requestedModel, shortSessionHash(sessionHash), accountID)
|
||||
}
|
||||
@@ -2512,6 +2610,10 @@ func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context,
|
||||
}
|
||||
accountsLoaded = true
|
||||
|
||||
// 提前预取窗口费用+RPM 计数,确保 routing 段内的调度检查调用能命中缓存
|
||||
ctx = s.withWindowCostPrefetch(ctx, accounts)
|
||||
ctx = s.withRPMPrefetch(ctx, accounts)
|
||||
|
||||
routingSet := make(map[int64]struct{}, len(routingAccountIDs))
|
||||
for _, id := range routingAccountIDs {
|
||||
if id > 0 {
|
||||
@@ -2539,6 +2641,12 @@ func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context,
|
||||
if !s.isAccountSchedulableForModelSelection(ctx, acc, requestedModel) {
|
||||
continue
|
||||
}
|
||||
if !s.isAccountSchedulableForWindowCost(ctx, acc, false) {
|
||||
continue
|
||||
}
|
||||
if !s.isAccountSchedulableForRPM(ctx, acc, false) {
|
||||
continue
|
||||
}
|
||||
if selected == nil {
|
||||
selected = acc
|
||||
continue
|
||||
@@ -2589,7 +2697,7 @@ func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context,
|
||||
if clearSticky {
|
||||
_ = s.cache.DeleteSessionAccountID(ctx, derefGroupID(groupID), sessionHash)
|
||||
}
|
||||
if !clearSticky && s.isAccountInGroup(account, groupID) && account.Platform == platform && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) {
|
||||
if !clearSticky && s.isAccountInGroup(account, groupID) && account.Platform == platform && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) && s.isAccountSchedulableForWindowCost(ctx, account, true) && s.isAccountSchedulableForRPM(ctx, account, true) {
|
||||
return account, nil
|
||||
}
|
||||
}
|
||||
@@ -2610,6 +2718,10 @@ func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context,
|
||||
}
|
||||
}
|
||||
|
||||
// 批量预取窗口费用+RPM 计数,避免逐个账号查询(N+1)
|
||||
ctx = s.withWindowCostPrefetch(ctx, accounts)
|
||||
ctx = s.withRPMPrefetch(ctx, accounts)
|
||||
|
||||
// 3. 按优先级+最久未用选择(考虑模型支持)
|
||||
var selected *Account
|
||||
for i := range accounts {
|
||||
@@ -2628,6 +2740,12 @@ func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context,
|
||||
if !s.isAccountSchedulableForModelSelection(ctx, acc, requestedModel) {
|
||||
continue
|
||||
}
|
||||
if !s.isAccountSchedulableForWindowCost(ctx, acc, false) {
|
||||
continue
|
||||
}
|
||||
if !s.isAccountSchedulableForRPM(ctx, acc, false) {
|
||||
continue
|
||||
}
|
||||
if selected == nil {
|
||||
selected = acc
|
||||
continue
|
||||
@@ -2697,7 +2815,7 @@ func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, g
|
||||
if clearSticky {
|
||||
_ = s.cache.DeleteSessionAccountID(ctx, derefGroupID(groupID), sessionHash)
|
||||
}
|
||||
if !clearSticky && s.isAccountInGroup(account, groupID) && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) {
|
||||
if !clearSticky && s.isAccountInGroup(account, groupID) && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) && s.isAccountSchedulableForWindowCost(ctx, account, true) && s.isAccountSchedulableForRPM(ctx, account, true) {
|
||||
if account.Platform == nativePlatform || (account.Platform == PlatformAntigravity && account.IsMixedSchedulingEnabled()) {
|
||||
if s.debugModelRoutingEnabled() {
|
||||
logger.LegacyPrintf("service.gateway", "[ModelRoutingDebug] legacy mixed routed sticky hit: group_id=%v model=%s session=%s account=%d", derefGroupID(groupID), requestedModel, shortSessionHash(sessionHash), accountID)
|
||||
@@ -2718,6 +2836,10 @@ func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, g
|
||||
}
|
||||
accountsLoaded = true
|
||||
|
||||
// 提前预取窗口费用+RPM 计数,确保 routing 段内的调度检查调用能命中缓存
|
||||
ctx = s.withWindowCostPrefetch(ctx, accounts)
|
||||
ctx = s.withRPMPrefetch(ctx, accounts)
|
||||
|
||||
routingSet := make(map[int64]struct{}, len(routingAccountIDs))
|
||||
for _, id := range routingAccountIDs {
|
||||
if id > 0 {
|
||||
@@ -2749,6 +2871,12 @@ func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, g
|
||||
if !s.isAccountSchedulableForModelSelection(ctx, acc, requestedModel) {
|
||||
continue
|
||||
}
|
||||
if !s.isAccountSchedulableForWindowCost(ctx, acc, false) {
|
||||
continue
|
||||
}
|
||||
if !s.isAccountSchedulableForRPM(ctx, acc, false) {
|
||||
continue
|
||||
}
|
||||
if selected == nil {
|
||||
selected = acc
|
||||
continue
|
||||
@@ -2799,7 +2927,7 @@ func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, g
|
||||
if clearSticky {
|
||||
_ = s.cache.DeleteSessionAccountID(ctx, derefGroupID(groupID), sessionHash)
|
||||
}
|
||||
if !clearSticky && s.isAccountInGroup(account, groupID) && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) {
|
||||
if !clearSticky && s.isAccountInGroup(account, groupID) && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) && s.isAccountSchedulableForWindowCost(ctx, account, true) && s.isAccountSchedulableForRPM(ctx, account, true) {
|
||||
if account.Platform == nativePlatform || (account.Platform == PlatformAntigravity && account.IsMixedSchedulingEnabled()) {
|
||||
return account, nil
|
||||
}
|
||||
@@ -2818,6 +2946,10 @@ func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, g
|
||||
}
|
||||
}
|
||||
|
||||
// 批量预取窗口费用+RPM 计数,避免逐个账号查询(N+1)
|
||||
ctx = s.withWindowCostPrefetch(ctx, accounts)
|
||||
ctx = s.withRPMPrefetch(ctx, accounts)
|
||||
|
||||
// 3. 按优先级+最久未用选择(考虑模型支持和混合调度)
|
||||
var selected *Account
|
||||
for i := range accounts {
|
||||
@@ -2840,6 +2972,12 @@ func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, g
|
||||
if !s.isAccountSchedulableForModelSelection(ctx, acc, requestedModel) {
|
||||
continue
|
||||
}
|
||||
if !s.isAccountSchedulableForWindowCost(ctx, acc, false) {
|
||||
continue
|
||||
}
|
||||
if !s.isAccountSchedulableForRPM(ctx, acc, false) {
|
||||
continue
|
||||
}
|
||||
if selected == nil {
|
||||
selected = acc
|
||||
continue
|
||||
@@ -5185,7 +5323,7 @@ func (s *GatewayService) isThinkingBlockSignatureError(respBody []byte) bool {
|
||||
}
|
||||
|
||||
func (s *GatewayService) shouldFailoverOn400(respBody []byte) bool {
|
||||
// 只对“可能是兼容性差异导致”的 400 允许切换,避免无意义重试。
|
||||
// 只对"可能是兼容性差异导致"的 400 允许切换,避免无意义重试。
|
||||
// 默认保守:无法识别则不切换。
|
||||
msg := strings.ToLower(strings.TrimSpace(extractUpstreamErrorMessage(respBody)))
|
||||
if msg == "" {
|
||||
|
||||
17
backend/internal/service/rpm_cache.go
Normal file
17
backend/internal/service/rpm_cache.go
Normal file
@@ -0,0 +1,17 @@
|
||||
package service
|
||||
|
||||
import "context"
|
||||
|
||||
// RPMCache RPM 计数器缓存接口
|
||||
// 用于 Anthropic OAuth/SetupToken 账号的每分钟请求数限制
|
||||
type RPMCache interface {
|
||||
// IncrementRPM 原子递增并返回当前分钟的计数
|
||||
// 使用 Redis 服务器时间确定 minute key,避免多实例时钟偏差
|
||||
IncrementRPM(ctx context.Context, accountID int64) (count int, err error)
|
||||
|
||||
// GetRPM 获取当前分钟的 RPM 计数
|
||||
GetRPM(ctx context.Context, accountID int64) (count int, err error)
|
||||
|
||||
// GetRPMBatch 批量获取多个账号的 RPM 计数(使用 Pipeline)
|
||||
GetRPMBatch(ctx context.Context, accountIDs []int64) (map[int64]int, error)
|
||||
}
|
||||
Reference in New Issue
Block a user