From d707c1f93f906664630f4d8909252b5fb379ad9b Mon Sep 17 00:00:00 2001 From: erio Date: Sat, 7 Mar 2026 04:10:54 +0800 Subject: [PATCH] feat: client affinity reverse index and balanced scheduling MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Separate clientAffinityTTL (24h) from stickySessionTTL (1h) - Atomic Lua dual-write for forward (client→accounts) and reverse (account→clients) affinity indexes with lazy expiry cleanup - Extract Lua scripts to standalone .lua files with //go:embed - Add filterByMinAffinityCount to Layer 1 & 2 scheduling so new clients are routed to accounts with fewest existing affinities - Skip Redis query when no affinity-enabled accounts exist (zero overhead for non-affinity users) Co-Authored-By: Claude Opus 4.6 --- backend/internal/repository/gateway_cache.go | 98 ++++++++++++------- .../internal/repository/lua/get_affinity.lua | 5 + .../repository/lua/get_affinity_count.lua | 5 + .../repository/lua/remove_affinity.lua | 8 ++ .../repository/lua/update_affinity.lua | 15 +++ .../service/antigravity_smart_retry_test.go | 3 + .../gateway_hotpath_optimization_test.go | 3 + .../service/gateway_multiplatform_test.go | 3 + backend/internal/service/gateway_service.go | 89 ++++++++++++++--- .../service/gemini_multiplatform_test.go | 3 + .../service/openai_gateway_service_test.go | 3 + .../service/openai_ws_state_store_test.go | 3 + backend/internal/testutil/stubs.go | 3 + 13 files changed, 195 insertions(+), 46 deletions(-) create mode 100644 backend/internal/repository/lua/get_affinity.lua create mode 100644 backend/internal/repository/lua/get_affinity_count.lua create mode 100644 backend/internal/repository/lua/remove_affinity.lua create mode 100644 backend/internal/repository/lua/update_affinity.lua diff --git a/backend/internal/repository/gateway_cache.go b/backend/internal/repository/gateway_cache.go index 8fa6d0c401..2de0b1c739 100644 --- a/backend/internal/repository/gateway_cache.go +++ b/backend/internal/repository/gateway_cache.go @@ -2,6 +2,7 @@ package repository import ( "context" + _ "embed" "fmt" "strconv" "time" @@ -11,8 +12,25 @@ import ( ) const ( - stickySessionPrefix = "sticky_session:" - clientAffinityPrefix = "client_affinity:" + stickySessionPrefix = "sticky_session:" + clientAffinityPrefix = "client_affinity:" + clientAffinityReversePrefix = "client_affinity_rev:" +) + +var ( + //go:embed lua/get_affinity.lua + getAffinityLua string + //go:embed lua/update_affinity.lua + updateAffinityLua string + //go:embed lua/remove_affinity.lua + removeAffinityLua string + //go:embed lua/get_affinity_count.lua + getAffinityCountLua string + + getAffinityScript = redis.NewScript(getAffinityLua) + updateAffinityScript = redis.NewScript(updateAffinityLua) + removeAffinityScript = redis.NewScript(removeAffinityLua) + getAffinityCountScript = redis.NewScript(getAffinityCountLua) ) type gatewayCache struct { @@ -45,43 +63,22 @@ func (c *gatewayCache) RefreshSessionTTL(ctx context.Context, groupID int64, ses } // DeleteSessionAccountID 删除粘性会话与账号的绑定关系。 -// 当检测到绑定的账号不可用（如状态错误、禁用、不可调度等）时调用， -// 以便下次请求能够重新选择可用账号。 -// -// DeleteSessionAccountID removes the sticky session binding for the given session. -// Called when the bound account becomes unavailable (e.g., error status, disabled, -// or unschedulable), allowing subsequent requests to select a new available account. func (c *gatewayCache) DeleteSessionAccountID(ctx context.Context, groupID int64, sessionHash string) error { key := buildSessionKey(groupID, sessionHash) return c.rdb.Del(ctx, key).Err() } -// buildAffinityKey 构建客户端亲和 key +// buildAffinityKey 构建正向亲和 key（client → accounts） // 格式: client_affinity:{groupID}:{clientID} func buildAffinityKey(groupID int64, clientID string) string { return fmt.Sprintf("%s%d:%s", clientAffinityPrefix, groupID, clientID) } -// getAffinityScript: 清理过期成员后返回亲和账号列表（按最近使用降序） -// KEYS[1] = client_affinity:{groupID}:{clientID} -// ARGV[1] = 过期阈值时间戳 (now - ttl) -var getAffinityScript = redis.NewScript(` -redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1]) -return redis.call('ZREVRANGE', KEYS[1], 0, -1) -`) - -// updateAffinityScript: 清理过期成员、添加/更新亲和关系、刷新 TTL -// KEYS[1] = client_affinity:{groupID}:{clientID} -// ARGV[1] = 当前时间戳 -// ARGV[2] = TTL 秒数 -// ARGV[3] = accountID -// ARGV[4] = 过期阈值时间戳 (now - ttl) -var updateAffinityScript = redis.NewScript(` -redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[4]) -redis.call('ZADD', KEYS[1], ARGV[1], ARGV[3]) -redis.call('EXPIRE', KEYS[1], ARGV[2]) -return 1 -`) +// buildAffinityReverseKey 构建反向亲和 key（account → clients） +// 格式: client_affinity_rev:{groupID}:{accountID} +func buildAffinityReverseKey(groupID int64, accountID int64) string { + return fmt.Sprintf("%s%d:%d", clientAffinityReversePrefix, groupID, accountID) +} func (c *gatewayCache) GetClientAffinityAccounts(ctx context.Context, groupID int64, clientID string, ttl time.Duration) ([]int64, error) { key := buildAffinityKey(groupID, clientID) @@ -108,17 +105,50 @@ func (c *gatewayCache) GetClientAffinityAccounts(ctx context.Context, groupID in } func (c *gatewayCache) UpdateClientAffinity(ctx context.Context, groupID int64, clientID string, accountID int64, ttl time.Duration) error { - key := buildAffinityKey(groupID, clientID) + fwdKey := buildAffinityKey(groupID, clientID) + revKey := buildAffinityReverseKey(groupID, accountID) now := time.Now().Unix() ttlSeconds := int64(ttl.Seconds()) expireThreshold := now - ttlSeconds - return updateAffinityScript.Run(ctx, c.rdb, []string{key}, - now, ttlSeconds, accountID, expireThreshold, + return updateAffinityScript.Run(ctx, c.rdb, []string{fwdKey, revKey}, + now, ttlSeconds, accountID, expireThreshold, clientID, ).Err() } func (c *gatewayCache) RemoveClientAffinity(ctx context.Context, groupID int64, clientID string, accountID int64) error { - key := buildAffinityKey(groupID, clientID) - return c.rdb.ZRem(ctx, key, accountID).Err() + fwdKey := buildAffinityKey(groupID, clientID) + revKey := buildAffinityReverseKey(groupID, accountID) + + return removeAffinityScript.Run(ctx, c.rdb, []string{fwdKey, revKey}, + accountID, clientID, + ).Err() +} + +// GetAccountAffinityCountBatch 批量获取账号的亲和客户端数量（惰性清理过期成员） +func (c *gatewayCache) GetAccountAffinityCountBatch(ctx context.Context, groupID int64, accountIDs []int64, ttl time.Duration) (map[int64]int64, error) { + if len(accountIDs) == 0 { + return map[int64]int64{}, nil + } + + now := time.Now().Unix() + expireThreshold := now - int64(ttl.Seconds()) + + pipe := c.rdb.Pipeline() + cmds := make([]*redis.Cmd, len(accountIDs)) + for i, accID := range accountIDs { + key := buildAffinityReverseKey(groupID, accID) + cmds[i] = getAffinityCountScript.Run(ctx, pipe, []string{key}, expireThreshold) + } + _, err := pipe.Exec(ctx) + if err != nil && err != redis.Nil { + return nil, err + } + + result := make(map[int64]int64, len(accountIDs)) + for i, accID := range accountIDs { + count, _ := cmds[i].Int64() + result[accID] = count + } + return result, nil } diff --git a/backend/internal/repository/lua/get_affinity.lua b/backend/internal/repository/lua/get_affinity.lua new file mode 100644 index 0000000000..7e3971dbea --- /dev/null +++ b/backend/internal/repository/lua/get_affinity.lua @@ -0,0 +1,5 @@ +-- 清理过期成员后返回亲和账号列表（按最近使用降序） +-- KEYS[1] = client_affinity:{groupID}:{clientID} +-- ARGV[1] = 过期阈值时间戳 (now - ttl) +redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1]) +return redis.call('ZREVRANGE', KEYS[1], 0, -1) diff --git a/backend/internal/repository/lua/get_affinity_count.lua b/backend/internal/repository/lua/get_affinity_count.lua new file mode 100644 index 0000000000..4b00c22daa --- /dev/null +++ b/backend/internal/repository/lua/get_affinity_count.lua @@ -0,0 +1,5 @@ +-- 清理过期成员后返回反向索引的成员数量 +-- KEYS[1] = client_affinity_rev:{groupID}:{accountID} +-- ARGV[1] = 过期阈值时间戳 (now - ttl) +redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1]) +return redis.call('ZCARD', KEYS[1]) diff --git a/backend/internal/repository/lua/remove_affinity.lua b/backend/internal/repository/lua/remove_affinity.lua new file mode 100644 index 0000000000..717974308c --- /dev/null +++ b/backend/internal/repository/lua/remove_affinity.lua @@ -0,0 +1,8 @@ +-- 原子双删正向+反向索引 +-- KEYS[1] = client_affinity:{groupID}:{clientID} (正向) +-- KEYS[2] = client_affinity_rev:{groupID}:{accountID} (反向) +-- ARGV[1] = accountID (正向索引中要删除的成员) +-- ARGV[2] = clientID (反向索引中要删除的成员) +redis.call('ZREM', KEYS[1], ARGV[1]) +redis.call('ZREM', KEYS[2], ARGV[2]) +return 1 diff --git a/backend/internal/repository/lua/update_affinity.lua b/backend/internal/repository/lua/update_affinity.lua new file mode 100644 index 0000000000..27fe8cb217 --- /dev/null +++ b/backend/internal/repository/lua/update_affinity.lua @@ -0,0 +1,15 @@ +-- 原子双写正向+反向索引 +-- KEYS[1] = client_affinity:{groupID}:{clientID} (正向: client → accounts) +-- KEYS[2] = client_affinity_rev:{groupID}:{accountID} (反向: account → clients) +-- ARGV[1] = 当前时间戳 (score) +-- ARGV[2] = TTL 秒数 +-- ARGV[3] = accountID (正向索引的成员) +-- ARGV[4] = 过期阈值时间戳 (now - ttl) +-- ARGV[5] = clientID (反向索引的成员) +redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[4]) +redis.call('ZADD', KEYS[1], ARGV[1], ARGV[3]) +redis.call('EXPIRE', KEYS[1], ARGV[2]) +redis.call('ZREMRANGEBYSCORE', KEYS[2], '-inf', ARGV[4]) +redis.call('ZADD', KEYS[2], ARGV[1], ARGV[5]) +redis.call('EXPIRE', KEYS[2], ARGV[2]) +return 1 diff --git a/backend/internal/service/antigravity_smart_retry_test.go b/backend/internal/service/antigravity_smart_retry_test.go index 7f836bf47c..19b614d6d7 100644 --- a/backend/internal/service/antigravity_smart_retry_test.go +++ b/backend/internal/service/antigravity_smart_retry_test.go @@ -39,6 +39,9 @@ func (c *stubSmartRetryCache) UpdateClientAffinity(_ context.Context, _ int64, _ func (c *stubSmartRetryCache) RemoveClientAffinity(_ context.Context, _ int64, _ string, _ int64) error { return nil } +func (c *stubSmartRetryCache) GetAccountAffinityCountBatch(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) { + return map[int64]int64{}, nil +} // mockSmartRetryUpstream 用于 handleSmartRetry 测试的 mock upstream type mockSmartRetryUpstream struct { diff --git a/backend/internal/service/gateway_hotpath_optimization_test.go b/backend/internal/service/gateway_hotpath_optimization_test.go index 137490e8db..5b3559e081 100644 --- a/backend/internal/service/gateway_hotpath_optimization_test.go +++ b/backend/internal/service/gateway_hotpath_optimization_test.go @@ -152,6 +152,9 @@ func (s *stickyGatewayCacheHotpathStub) UpdateClientAffinity(_ context.Context, func (s *stickyGatewayCacheHotpathStub) RemoveClientAffinity(_ context.Context, _ int64, _ string, _ int64) error { return nil } +func (s *stickyGatewayCacheHotpathStub) GetAccountAffinityCountBatch(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) { + return map[int64]int64{}, nil +} func (s *modelsListAccountRepoStub) ListSchedulableByGroupID(ctx context.Context, groupID int64) ([]Account, error) { s.listByGroupCalls.Add(1) diff --git a/backend/internal/service/gateway_multiplatform_test.go b/backend/internal/service/gateway_multiplatform_test.go index 466ff294bb..a913d42c18 100644 --- a/backend/internal/service/gateway_multiplatform_test.go +++ b/backend/internal/service/gateway_multiplatform_test.go @@ -244,6 +244,9 @@ func (m *mockGatewayCacheForPlatform) UpdateClientAffinity(_ context.Context, _ func (m *mockGatewayCacheForPlatform) RemoveClientAffinity(_ context.Context, _ int64, _ string, _ int64) error { return nil } +func (m *mockGatewayCacheForPlatform) GetAccountAffinityCountBatch(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) { + return map[int64]int64{}, nil +} type mockGroupRepoForGateway struct { groups map[int64]*Group diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index bf03b8a0ec..d2efd3e1fb 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -40,7 +40,8 @@ import ( const ( claudeAPIURL = "https://api.anthropic.com/v1/messages?beta=true" claudeAPICountTokensURL = "https://api.anthropic.com/v1/messages/count_tokens?beta=true" - stickySessionTTL = time.Hour // 粘性会话TTL + stickySessionTTL = time.Hour // 粘性会话TTL + clientAffinityTTL = 24 * time.Hour // 客户端亲和TTL defaultMaxLineSize = 40 * 1024 * 1024 // Canonical Claude Code banner. Keep it EXACT (no trailing whitespace/newlines) // to match real Claude CLI traffic as closely as possible. When we need a visual @@ -68,8 +69,9 @@ type forceCacheBillingKeyType struct{} // accountWithLoad 账号与负载信息的组合，用于负载感知调度 type accountWithLoad struct { - account *Account - loadInfo *AccountLoadInfo + account *Account + loadInfo *AccountLoadInfo + affinityCount int64 // 亲和客户端数量（反向索引），越少越优先 } var ForceCacheBillingContextKey = forceCacheBillingKeyType{} @@ -373,8 +375,10 @@ type GatewayCache interface { GetClientAffinityAccounts(ctx context.Context, groupID int64, clientID string, ttl time.Duration) ([]int64, error) // UpdateClientAffinity 添加/更新客户端亲和关系（更新 score 为当前时间戳，刷新 key TTL） UpdateClientAffinity(ctx context.Context, groupID int64, clientID string, accountID int64, ttl time.Duration) error - // RemoveClientAffinity 删除单个客户端亲和关系 + // RemoveClientAffinity 删除单个客户端亲和关系（原子双删正向+反向索引） RemoveClientAffinity(ctx context.Context, groupID int64, clientID string, accountID int64) error + // GetAccountAffinityCountBatch 批量获取账号的亲和客户端数量（惰性清理过期成员） + GetAccountAffinityCountBatch(ctx context.Context, groupID int64, accountIDs []int64, ttl time.Duration) (map[int64]int64, error) } // derefGroupID safely dereferences *int64 to int64, returning 0 if nil @@ -1368,7 +1372,10 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro } if len(routingAvailable) > 0 { - // 排序：优先级 > 负载率 > 最后使用时间 + // 批量获取亲和客户端数量 + s.populateAffinityCounts(ctx, routingAvailable, derefGroupID(groupID)) + + // 排序：优先级 > 负载率 > 亲和客户端数 > 最后使用时间 sort.SliceStable(routingAvailable, func(i, j int) bool { a, b := routingAvailable[i], routingAvailable[j] if a.account.Priority != b.account.Priority { @@ -1377,6 +1384,9 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro if a.loadInfo.LoadRate != b.loadInfo.LoadRate { return a.loadInfo.LoadRate < b.loadInfo.LoadRate } + if a.affinityCount != b.affinityCount { + return a.affinityCount < b.affinityCount + } switch { case a.account.LastUsedAt == nil && b.account.LastUsedAt != nil: return true @@ -1403,7 +1413,7 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro _ = s.cache.SetSessionAccountID(ctx, derefGroupID(groupID), sessionHash, item.account.ID, stickySessionTTL) } if affinityClientID != "" && s.cache != nil && item.account.IsClientAffinityEnabled() { - _ = s.cache.UpdateClientAffinity(ctx, derefGroupID(groupID), affinityClientID, item.account.ID, stickySessionTTL) + _ = s.cache.UpdateClientAffinity(ctx, derefGroupID(groupID), affinityClientID, item.account.ID, clientAffinityTTL) } if s.debugModelRoutingEnabled() { logger.LegacyPrintf("service.gateway", "[ModelRoutingDebug] routed select: group_id=%v model=%s session=%s account=%d", derefGroupID(groupID), requestedModel, shortSessionHash(sessionHash), item.account.ID) @@ -1503,7 +1513,7 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro // ============ Layer 1.6: 客户端亲和（仅在粘性会话未命中时生效） ============ if affinityClientID != "" && s.cache != nil && stickyAccountID <= 0 { - affinityAccountIDs, err := s.cache.GetClientAffinityAccounts(ctx, derefGroupID(groupID), affinityClientID, stickySessionTTL) + affinityAccountIDs, err := s.cache.GetClientAffinityAccounts(ctx, derefGroupID(groupID), affinityClientID, clientAffinityTTL) if err == nil && len(affinityAccountIDs) > 0 { for _, affinityAccID := range affinityAccountIDs { if isExcluded(affinityAccID) { @@ -1542,7 +1552,7 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro continue } // 亲和命中：更新亲和 score + 绑定粘性会话 - _ = s.cache.UpdateClientAffinity(ctx, derefGroupID(groupID), affinityClientID, affinityAccID, stickySessionTTL) + _ = s.cache.UpdateClientAffinity(ctx, derefGroupID(groupID), affinityClientID, affinityAccID, clientAffinityTTL) if sessionHash != "" { _ = s.cache.SetSessionAccountID(ctx, derefGroupID(groupID), sessionHash, affinityAccID, stickySessionTTL) } @@ -1614,7 +1624,7 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro if err != nil { if result, ok := s.tryAcquireByLegacyOrder(ctx, candidates, groupID, sessionHash, preferOAuth); ok { if affinityClientID != "" && s.cache != nil && result.Account != nil && result.Account.IsClientAffinityEnabled() { - _ = s.cache.UpdateClientAffinity(ctx, derefGroupID(groupID), affinityClientID, result.Account.ID, stickySessionTTL) + _ = s.cache.UpdateClientAffinity(ctx, derefGroupID(groupID), affinityClientID, result.Account.ID, clientAffinityTTL) } return result, nil } @@ -1633,13 +1643,18 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro } } - // 分层过滤选择：优先级 → 负载率 → LRU + // 批量获取亲和客户端数量（用于均衡分配新客户端） + s.populateAffinityCounts(ctx, available, derefGroupID(groupID)) + + // 分层过滤选择：优先级 → 负载率 → 亲和客户端数 → LRU for len(available) > 0 { // 1. 取优先级最小的集合 candidates := filterByMinPriority(available) // 2. 取负载率最低的集合 candidates = filterByMinLoadRate(candidates) - // 3. LRU 选择最久未用的账号 + // 3. 取亲和客户端数最少的集合 + candidates = filterByMinAffinityCount(candidates) + // 4. LRU 选择最久未用的账号 selected := selectByLRU(candidates, preferOAuth) if selected == nil { break @@ -1656,7 +1671,7 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro } // 更新客户端亲和关系 if affinityClientID != "" && s.cache != nil && selected.account.IsClientAffinityEnabled() { - _ = s.cache.UpdateClientAffinity(ctx, derefGroupID(groupID), affinityClientID, selected.account.ID, stickySessionTTL) + _ = s.cache.UpdateClientAffinity(ctx, derefGroupID(groupID), affinityClientID, selected.account.ID, clientAffinityTTL) } return &AccountSelectionResult{ Account: selected.account, @@ -2415,6 +2430,36 @@ func (s *GatewayService) getSchedulableAccount(ctx context.Context, accountID in return s.accountRepo.GetByID(ctx, accountID) } +// populateAffinityCounts 批量获取账号的亲和客户端数量并填入 accountWithLoad 切片。 +// 仅当存在开启了客户端亲和的账号时才查询 Redis，否则跳过。 +func (s *GatewayService) populateAffinityCounts(ctx context.Context, accounts []accountWithLoad, groupID int64) { + if s.cache == nil || len(accounts) == 0 { + return + } + // 快速检查：是否有任何账号开启了亲和 + hasAffinity := false + for _, acc := range accounts { + if acc.account.IsClientAffinityEnabled() { + hasAffinity = true + break + } + } + if !hasAffinity { + return + } + accountIDs := make([]int64, len(accounts)) + for i, acc := range accounts { + accountIDs[i] = acc.account.ID + } + countMap, err := s.cache.GetAccountAffinityCountBatch(ctx, groupID, accountIDs, clientAffinityTTL) + if err != nil { + return // 查询失败不影响调度，affinityCount 保持 0 + } + for i := range accounts { + accounts[i].affinityCount = countMap[accounts[i].account.ID] + } +} + // filterByMinPriority 过滤出优先级最小的账号集合 func filterByMinPriority(accounts []accountWithLoad) []accountWithLoad { if len(accounts) == 0 { @@ -2455,6 +2500,26 @@ func filterByMinLoadRate(accounts []accountWithLoad) []accountWithLoad { return result } +// filterByMinAffinityCount 过滤出亲和客户端数最少的账号集合 +func filterByMinAffinityCount(accounts []accountWithLoad) []accountWithLoad { + if len(accounts) == 0 { + return accounts + } + minCount := accounts[0].affinityCount + for _, acc := range accounts[1:] { + if acc.affinityCount < minCount { + minCount = acc.affinityCount + } + } + result := make([]accountWithLoad, 0, len(accounts)) + for _, acc := range accounts { + if acc.affinityCount == minCount { + result = append(result, acc) + } + } + return result +} + // selectByLRU 从集合中选择最久未用的账号 // 如果有多个账号具有相同的最小 LastUsedAt，则随机选择一个 func selectByLRU(accounts []accountWithLoad, preferOAuth bool) *accountWithLoad { diff --git a/backend/internal/service/gemini_multiplatform_test.go b/backend/internal/service/gemini_multiplatform_test.go index 56f94d5f1f..8a3d6a5f61 100644 --- a/backend/internal/service/gemini_multiplatform_test.go +++ b/backend/internal/service/gemini_multiplatform_test.go @@ -297,6 +297,9 @@ func (m *mockGatewayCacheForGemini) UpdateClientAffinity(_ context.Context, _ in func (m *mockGatewayCacheForGemini) RemoveClientAffinity(_ context.Context, _ int64, _ string, _ int64) error { return nil } +func (m *mockGatewayCacheForGemini) GetAccountAffinityCountBatch(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) { + return map[int64]int64{}, nil +} // TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_GeminiPlatform 测试 Gemini 单平台选择 func TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_GeminiPlatform(t *testing.T) { diff --git a/backend/internal/service/openai_gateway_service_test.go b/backend/internal/service/openai_gateway_service_test.go index 7e53c8996e..490df8cd7b 100644 --- a/backend/internal/service/openai_gateway_service_test.go +++ b/backend/internal/service/openai_gateway_service_test.go @@ -274,6 +274,9 @@ func (c *stubGatewayCache) UpdateClientAffinity(_ context.Context, _ int64, _ st func (c *stubGatewayCache) RemoveClientAffinity(_ context.Context, _ int64, _ string, _ int64) error { return nil } +func (c *stubGatewayCache) GetAccountAffinityCountBatch(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) { + return map[int64]int64{}, nil +} func TestOpenAISelectAccountWithLoadAwareness_FiltersUnschedulable(t *testing.T) { now := time.Now() diff --git a/backend/internal/service/openai_ws_state_store_test.go b/backend/internal/service/openai_ws_state_store_test.go index 6f77f64356..0964db61d0 100644 --- a/backend/internal/service/openai_ws_state_store_test.go +++ b/backend/internal/service/openai_ws_state_store_test.go @@ -202,6 +202,9 @@ func (c *openAIWSStateStoreTimeoutProbeCache) UpdateClientAffinity(_ context.Con func (c *openAIWSStateStoreTimeoutProbeCache) RemoveClientAffinity(_ context.Context, _ int64, _ string, _ int64) error { return nil } +func (c *openAIWSStateStoreTimeoutProbeCache) GetAccountAffinityCountBatch(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) { + return map[int64]int64{}, nil +} func TestOpenAIWSStateStore_RedisOpsUseShortTimeout(t *testing.T) { probe := &openAIWSStateStoreTimeoutProbeCache{} diff --git a/backend/internal/testutil/stubs.go b/backend/internal/testutil/stubs.go index 669372ef51..acc32a2d3b 100644 --- a/backend/internal/testutil/stubs.go +++ b/backend/internal/testutil/stubs.go @@ -106,6 +106,9 @@ func (c StubGatewayCache) UpdateClientAffinity(_ context.Context, _ int64, _ str func (c StubGatewayCache) RemoveClientAffinity(_ context.Context, _ int64, _ string, _ int64) error { return nil } +func (c StubGatewayCache) GetAccountAffinityCountBatch(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) { + return map[int64]int64{}, nil +} // ============================================================ // StubSessionLimitCache — service.SessionLimitCache 的空实现