From b605a9cd06a5dea8d04334eff3437f272f51d4f4 Mon Sep 17 00:00:00 2001 From: yuanjia Date: Sun, 19 Jul 2026 02:40:59 +0800 Subject: [PATCH 1/2] feat(channel): optional adaptive balance with circuit breaker (default off) Add score-based channel selection, local circuit breaker, and EWMA metrics. Flags default false so behavior matches upstream until operators enable ADAPTIVE_BALANCE_ENABLED / SHADOW / CHANNEL_CIRCUIT_BREAKER_ENABLED. --- .env.example | 10 + common/env.go | 13 + common/init.go | 7 + constant/env.go | 9 + controller/relay.go | 130 +++++++- controller/relay_origin_test.go | 43 +++ controller/relay_retry_test.go | 37 +++ model/ability.go | 10 + model/channel_cache.go | 64 +++- model/channel_selection_exclusion_test.go | 54 +++ service/channel_adaptive.go | 326 +++++++++++++++++++ service/channel_adaptive_test.go | 275 ++++++++++++++++ service/channel_affinity.go | 111 ++++++- service/channel_affinity_template_test.go | 26 +- service/channel_affinity_usage_cache_test.go | 29 +- service/channel_circuit.go | 245 ++++++++++++++ service/channel_metrics.go | 255 +++++++++++++++ service/channel_score.go | 166 ++++++++++ service/channel_select.go | 16 +- 19 files changed, 1782 insertions(+), 44 deletions(-) create mode 100644 controller/relay_origin_test.go create mode 100644 controller/relay_retry_test.go create mode 100644 model/channel_selection_exclusion_test.go create mode 100644 service/channel_adaptive.go create mode 100644 service/channel_adaptive_test.go create mode 100644 service/channel_circuit.go create mode 100644 service/channel_metrics.go create mode 100644 service/channel_score.go diff --git a/.env.example b/.env.example index a63ed7668e98..089993684027 100644 --- a/.env.example +++ b/.env.example @@ -97,3 +97,13 @@ LINUX_DO_USER_ENDPOINT=https://connect.linux.do/api/user # 用于验证支付成功/取消回调URL的域名安全性 # 示例: example.com,myapp.io 将允许 example.com, sub.example.com, myapp.io 等 # TRUSTED_REDIRECT_DOMAINS=example.com,myapp.io + +# --- Adaptive channel balance (default OFF) --- +# When enabled without shadow, selection uses score+circuit. Prefer shadow first. +# ADAPTIVE_BALANCE_ENABLED=false +# ADAPTIVE_BALANCE_SHADOW_MODE=false +# CHANNEL_CIRCUIT_BREAKER_ENABLED=false +# EWMA_ALPHA=0.1 +# MAX_CHANNEL_CONCURRENCY=10 +# CHANNEL_COOLDOWN_SECONDS=30 +# MAX_RETRY_CHANNELS=0 diff --git a/common/env.go b/common/env.go index 1aa340f85ea1..1cc9fdbb7de0 100644 --- a/common/env.go +++ b/common/env.go @@ -36,3 +36,16 @@ func GetEnvOrDefaultBool(env string, defaultValue bool) bool { } return b } + +func GetEnvOrDefaultFloat(env string, defaultValue float64) float64 { + if env == "" || os.Getenv(env) == "" { + return defaultValue + } + num, err := strconv.ParseFloat(os.Getenv(env), 64) + if err != nil { + SysError(fmt.Sprintf("failed to parse %s: %s, using default value: %g", env, err.Error(), defaultValue)) + return defaultValue + } + return num +} + diff --git a/common/init.go b/common/init.go index 88b2dc3e62e1..c7ab141d2104 100644 --- a/common/init.go +++ b/common/init.go @@ -110,6 +110,13 @@ func InitEnv() { RelayIdleConnTimeout = GetEnvOrDefault("RELAY_IDLE_CONN_TIMEOUT", 90) RelayMaxIdleConns = GetEnvOrDefault("RELAY_MAX_IDLE_CONNS", 500) RelayMaxIdleConnsPerHost = GetEnvOrDefault("RELAY_MAX_IDLE_CONNS_PER_HOST", 100) + constant.AdaptiveBalanceEnabled = GetEnvOrDefaultBool("ADAPTIVE_BALANCE_ENABLED", false) + constant.AdaptiveBalanceShadowMode = GetEnvOrDefaultBool("ADAPTIVE_BALANCE_SHADOW_MODE", false) + constant.ChannelCircuitBreakerEnabled = GetEnvOrDefaultBool("CHANNEL_CIRCUIT_BREAKER_ENABLED", false) + constant.MaxRetryChannels = GetEnvOrDefault("MAX_RETRY_CHANNELS", 0) + constant.ChannelCooldownSeconds = GetEnvOrDefault("CHANNEL_COOLDOWN_SECONDS", 30) + constant.EwmaAlpha = GetEnvOrDefaultFloat("EWMA_ALPHA", 0.1) + constant.MaxChannelConcurrency = GetEnvOrDefault("MAX_CHANNEL_CONCURRENCY", 10) // Initialize string variables with GetEnvOrDefaultString GeminiSafetySetting = GetEnvOrDefaultString("GEMINI_SAFETY_SETTING", "BLOCK_NONE") diff --git a/constant/env.go b/constant/env.go index 512bfc31126b..0dfd794bb590 100644 --- a/constant/env.go +++ b/constant/env.go @@ -25,3 +25,12 @@ var TaskPricePatches []string // TrustedRedirectDomains is a list of trusted domains for redirect URL validation. // Domains support subdomain matching (e.g., "example.com" matches "sub.example.com"). var TrustedRedirectDomains []string + +// Adaptive channel balance settings +var AdaptiveBalanceEnabled bool +var AdaptiveBalanceShadowMode bool +var ChannelCircuitBreakerEnabled bool +var MaxRetryChannels int +var ChannelCooldownSeconds int +var EwmaAlpha float64 +var MaxChannelConcurrency int diff --git a/controller/relay.go b/controller/relay.go index 6e91ccb60506..2e01d339dad1 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -6,6 +6,7 @@ import ( "io" "log" "net/http" + "net/url" "strings" "time" @@ -200,6 +201,7 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { addUsedChannel(c, channel.Id) bodyStorage, bodyErr := common.GetBodyStorage(c) if bodyErr != nil { + service.ReleaseAdaptiveCircuitPermit(c, channel.Id) // Ensure consistent 413 for oversized bodies even when error occurs later (e.g., retry path) if common.IsRequestBodyTooLargeError(bodyErr) || errors.Is(bodyErr, common.ErrRequestBodyTooLarge) { newAPIError = types.NewErrorWithStatusCode(bodyErr, types.ErrorCodeReadRequestBodyFailed, http.StatusRequestEntityTooLarge, types.ErrOptionWithSkipRetry()) @@ -210,15 +212,46 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { } c.Request.Body = io.NopCloser(bodyStorage) - switch relayFormat { - case types.RelayFormatOpenAIRealtime: - newAPIError = relay.WssHelper(c, relayInfo) - case types.RelayFormatClaude: - newAPIError = relay.ClaudeHelper(c, relayInfo) - case types.RelayFormatGemini: - newAPIError = geminiRelayHandler(c, relayInfo) - default: - newAPIError = relayHandler(c, relayInfo) + attemptStart := time.Now() + service.IncChannelConcurrency(channel.Id) + // Always dec even if helper panics (CustomRecovery still runs after). + func() { + defer service.DecChannelConcurrency(channel.Id) + switch relayFormat { + case types.RelayFormatOpenAIRealtime: + newAPIError = relay.WssHelper(c, relayInfo) + case types.RelayFormatClaude: + newAPIError = relay.ClaudeHelper(c, relayInfo) + case types.RelayFormatGemini: + newAPIError = geminiRelayHandler(c, relayInfo) + default: + newAPIError = relayHandler(c, relayInfo) + } + }() + { + statusCode := http.StatusOK + var recErr error + if newAPIError != nil { + statusCode = newAPIError.StatusCode + if statusCode == 0 { + statusCode = http.StatusInternalServerError + } + recErr = newAPIError + } + // Prefer UsingGroup (resolved auto group) so score buckets match selection. + metricGroup := relayInfo.UsingGroup + if metricGroup == "" { + metricGroup = relayInfo.TokenGroup + } + service.RecordAdaptiveResult( + c, + channel.Id, + metricGroup, + relayInfo.OriginModelName, + statusCode, + time.Since(attemptStart), + recErr, + ) } if newAPIError == nil { @@ -250,12 +283,30 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { var upgrader = websocket.Upgrader{ Subprotocols: []string{"realtime"}, // WS 握手支持的协议，如果有使用 Sec-WebSocket-Protocol，则必须在此声明对应的 Protocol TODO add other protocol - CheckOrigin: func(r *http.Request) bool { - return true // 允许跨域 - }, + CheckOrigin: isRealtimeWebSocketOriginAllowed, +} + +func isRealtimeWebSocketOriginAllowed(r *http.Request) bool { + if r == nil { + return false + } + originValue := strings.TrimSpace(r.Header.Get("Origin")) + if originValue == "" { + return true + } + + origin, err := url.Parse(originValue) + if err != nil || origin.Host == "" || (origin.Scheme != "http" && origin.Scheme != "https") { + return false + } + if strings.EqualFold(origin.Host, r.Host) { + return true + } + return common.ValidateRedirectURL(originValue) == nil } func addUsedChannel(c *gin.Context, channelId int) { + service.MarkChannelUsed(c, channelId) useChannel := c.GetStringSlice("use_channel") useChannel = append(useChannel, fmt.Sprintf("%d", channelId)) c.Set("use_channel", useChannel) @@ -317,6 +368,7 @@ func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service newAPIError := middleware.SetupContextForSelectedChannel(c, channel, info.OriginModelName) if newAPIError != nil { + service.ReleaseAdaptiveCircuitPermit(c, channel.Id) return nil, newAPIError } return channel, nil @@ -329,16 +381,22 @@ func shouldRetry(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) b if service.ShouldSkipRetryAfterChannelAffinityFailure(c) { return false } - if types.IsChannelError(openaiErr) { - return true + if retryTimes <= 0 { + return false } - if types.IsSkipRetryError(openaiErr) { + if _, ok := c.Get("specific_channel_id"); ok { return false } - if retryTimes <= 0 { + if openaiErr.GetErrorCode() == types.ErrorCodeGetChannelFailed { return false } - if _, ok := c.Get("specific_channel_id"); ok { + if isUpstreamChannelQuotaError(openaiErr) { + return true + } + if types.IsChannelError(openaiErr) { + return true + } + if types.IsSkipRetryError(openaiErr) { return false } code := openaiErr.StatusCode @@ -354,6 +412,44 @@ func shouldRetry(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) b return operation_setting.ShouldRetryByStatusCode(code) } +func isUpstreamChannelQuotaError(err *types.NewAPIError) bool { + if err == nil { + return false + } + code := strings.ToLower(strings.TrimSpace(string(err.GetErrorCode()))) + if code == string(types.ErrorCodeInsufficientUserQuota) || code == string(types.ErrorCodePreConsumeTokenQuotaFailed) { + return false + } + if err.StatusCode == http.StatusPaymentRequired { + return true + } + for _, marker := range []string{ + "insufficient_quota", + "quota_exceeded", + "billing_hard_limit_reached", + "insufficient_balance", + "insufficient_credits", + } { + if strings.Contains(code, marker) { + return true + } + } + message := strings.ToLower(err.Error()) + for _, marker := range []string{ + "insufficient quota", + "quota exceeded", + "insufficient balance", + "insufficient credit", + "额度不足", + "余额不足", + } { + if strings.Contains(message, marker) { + return true + } + } + return false +} + func processChannelError(c *gin.Context, channelError types.ChannelError, err *types.NewAPIError) { logger.LogError(c, fmt.Sprintf("channel error (channel #%d, status code: %d): %s", channelError.ChannelId, err.StatusCode, common.LocalLogPreview(err.Error()))) // 不要使用context获取渠道信息，异步处理时可能会出现渠道信息不一致的情况 diff --git a/controller/relay_origin_test.go b/controller/relay_origin_test.go new file mode 100644 index 000000000000..ef0e44c072fa --- /dev/null +++ b/controller/relay_origin_test.go @@ -0,0 +1,43 @@ +package controller + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/constant" + "github.com/stretchr/testify/require" +) + +func TestRealtimeWebSocketOriginAllowed(t *testing.T) { + originalDomains := append([]string(nil), constant.TrustedRedirectDomains...) + constant.TrustedRedirectDomains = []string{"example.com"} + t.Cleanup(func() { + constant.TrustedRedirectDomains = originalDomains + }) + + tests := []struct { + name string + origin string + host string + want bool + }{ + {name: "missing origin", host: "api.internal", want: true}, + {name: "same origin", origin: "https://api.internal", host: "api.internal", want: true}, + {name: "trusted exact domain", origin: "https://example.com", host: "api.internal", want: true}, + {name: "trusted subdomain", origin: "https://console.example.com", host: "api.internal", want: true}, + {name: "untrusted domain", origin: "https://evil.example.net", host: "api.internal", want: false}, + {name: "suffix spoof", origin: "https://fakeexample.com", host: "api.internal", want: false}, + {name: "invalid scheme", origin: "file://example.com", host: "api.internal", want: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + request := httptest.NewRequest(http.MethodGet, "https://"+tt.host+"/v1/realtime", nil) + if tt.origin != "" { + request.Header.Set("Origin", tt.origin) + } + require.Equal(t, tt.want, isRealtimeWebSocketOriginAllowed(request)) + }) + } +} diff --git a/controller/relay_retry_test.go b/controller/relay_retry_test.go new file mode 100644 index 000000000000..53d6f9b0ba4d --- /dev/null +++ b/controller/relay_retry_test.go @@ -0,0 +1,37 @@ +package controller + +import ( + "errors" + "net/http" + "testing" + + "github.com/QuantumNous/new-api/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestShouldRetryStopsAfterChannelSelectionFailure(t *testing.T) { + ctx, _ := gin.CreateTestContext(nil) + err := types.NewError(errors.New("no eligible channel"), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()) + require.False(t, shouldRetry(ctx, err, 2)) +} + +func TestShouldRetrySwitchesChannelOnUpstreamQuotaExhaustion(t *testing.T) { + ctx, _ := gin.CreateTestContext(nil) + err := types.WithOpenAIError(types.OpenAIError{ + Message: "upstream account has insufficient balance", + Code: "insufficient_quota", + }, http.StatusTooManyRequests) + require.True(t, shouldRetry(ctx, err, 2)) +} + +func TestShouldRetryDoesNotSwitchForLocalUserQuota(t *testing.T) { + ctx, _ := gin.CreateTestContext(nil) + err := types.NewErrorWithStatusCode( + errors.New("user quota insufficient"), + types.ErrorCodeInsufficientUserQuota, + http.StatusForbidden, + types.ErrOptionWithSkipRetry(), + ) + require.False(t, shouldRetry(ctx, err, 2)) +} diff --git a/model/ability.go b/model/ability.go index e67b28301e02..216b22912dd7 100644 --- a/model/ability.go +++ b/model/ability.go @@ -106,6 +106,10 @@ func getChannelQuery(group string, model string, retry int) (*gorm.DB, error) { } func GetChannel(group string, model string, retry int, requestPath string) (*Channel, error) { + return GetChannelExcluding(group, model, retry, requestPath, nil) +} + +func GetChannelExcluding(group string, model string, retry int, requestPath string, excluded map[int]struct{}) (*Channel, error) { var abilities []Ability var err error = nil @@ -122,6 +126,12 @@ func GetChannel(group string, model string, retry int, requestPath string) (*Cha return nil, err } abilities = filterAbilitiesByRequestPathAndModel(abilities, requestPath, model) + if len(excluded) > 0 { + abilities = lo.Filter(abilities, func(ability Ability, _ int) bool { + _, skip := excluded[ability.ChannelId] + return !skip + }) + } channel := Channel{} if len(abilities) > 0 { // Randomly choose one diff --git a/model/channel_cache.go b/model/channel_cache.go index 81923017d79c..9dcb095fa9eb 100644 --- a/model/channel_cache.go +++ b/model/channel_cache.go @@ -111,10 +111,60 @@ func SyncChannelCache(frequency int) { } } +// GetSatisfiedChannels returns all enabled channels for group+model (path-aware), +// highest priority first. Used by adaptive balance candidate collection. +// When memory cache is off, falls back to a single DB-selected channel. +func GetSatisfiedChannels(group string, modelName string, requestPath string) ([]*Channel, error) { + if !common.MemoryCacheEnabled { + ch, err := GetChannel(group, modelName, 0, requestPath) + if err != nil { + return nil, err + } + if ch == nil { + return nil, nil + } + return []*Channel{ch}, nil + } + + channelSyncLock.RLock() + defer channelSyncLock.RUnlock() + + ids := filterChannelsByRequestPathAndModel(group2model2channels[group][modelName], requestPath, modelName) + if len(ids) == 0 { + normalizedModel := ratio_setting.FormatMatchingModelName(modelName) + ids = filterChannelsByRequestPathAndModel(group2model2channels[group][normalizedModel], requestPath, normalizedModel) + } + if len(ids) == 0 { + return nil, nil + } + + out := make([]*Channel, 0, len(ids)) + seen := make(map[int]struct{}, len(ids)) + for _, id := range ids { + if _, ok := seen[id]; ok { + continue + } + seen[id] = struct{}{} + ch, ok := channelsIDM[id] + if !ok || ch == nil { + continue + } + if ch.Status != common.ChannelStatusEnabled { + continue + } + out = append(out, ch) + } + return out, nil +} + func GetRandomSatisfiedChannel(group string, model string, retry int, requestPath string) (*Channel, error) { + return GetRandomSatisfiedChannelExcluding(group, model, retry, requestPath, nil) +} + +func GetRandomSatisfiedChannelExcluding(group string, model string, retry int, requestPath string, excluded map[int]struct{}) (*Channel, error) { // if memory cache is disabled, get channel directly from database if !common.MemoryCacheEnabled { - return GetChannel(group, model, retry, requestPath) + return GetChannelExcluding(group, model, retry, requestPath, excluded) } channelSyncLock.RLock() @@ -132,6 +182,18 @@ func GetRandomSatisfiedChannel(group string, model string, retry int, requestPat if len(channels) == 0 { return nil, nil } + if len(excluded) > 0 { + filtered := make([]int, 0, len(channels)) + for _, channelID := range channels { + if _, skip := excluded[channelID]; !skip { + filtered = append(filtered, channelID) + } + } + channels = filtered + if len(channels) == 0 { + return nil, nil + } + } if len(channels) == 1 { if channel, ok := channelsIDM[channels[0]]; ok { diff --git a/model/channel_selection_exclusion_test.go b/model/channel_selection_exclusion_test.go new file mode 100644 index 000000000000..c50033512856 --- /dev/null +++ b/model/channel_selection_exclusion_test.go @@ -0,0 +1,54 @@ +package model + +import ( + "fmt" + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func TestGetChannelExcludingSkipsPreviouslyFailedChannel(t *testing.T) { + originalDB := DB + originalMemoryCacheEnabled := common.MemoryCacheEnabled + common.MemoryCacheEnabled = false + + dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_")) + db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&Channel{}, &Ability{})) + DB = db + t.Cleanup(func() { + DB = originalDB + common.MemoryCacheEnabled = originalMemoryCacheEnabled + if sqlDB, dbErr := db.DB(); dbErr == nil { + _ = sqlDB.Close() + } + }) + + priority := int64(100) + require.NoError(t, db.Create(&Channel{Id: 1, Name: "depleted"}).Error) + require.NoError(t, db.Create(&Channel{Id: 2, Name: "available"}).Error) + require.NoError(t, db.Create(&Ability{ + Group: "default", Model: "gpt-test", ChannelId: 1, + Enabled: true, Priority: &priority, Weight: 100, + }).Error) + require.NoError(t, db.Create(&Ability{ + Group: "default", Model: "gpt-test", ChannelId: 2, + Enabled: true, Priority: &priority, Weight: 1, + }).Error) + + channel, err := GetChannelExcluding( + "default", + "gpt-test", + 0, + "/v1/chat/completions", + map[int]struct{}{1: {}}, + ) + require.NoError(t, err) + require.NotNil(t, channel) + require.Equal(t, 2, channel.Id) +} diff --git a/service/channel_adaptive.go b/service/channel_adaptive.go new file mode 100644 index 000000000000..fe00ed1cfbf2 --- /dev/null +++ b/service/channel_adaptive.go @@ -0,0 +1,326 @@ +package service + +import ( + "fmt" + "math/rand/v2" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/model" + "github.com/gin-gonic/gin" +) + +var adaptiveLogSample = 0.01 // shadow-mode log sample rate + +// 请求上下文 key +type adaptiveContextKey string + +const ( + ctxKeyAdaptiveUsedChannels adaptiveContextKey = "adaptive_used_channels" + ctxKeyAdaptiveGroup adaptiveContextKey = "adaptive_group" + ctxKeyAdaptiveModel adaptiveContextKey = "adaptive_model" + ctxKeyAdaptiveSelected adaptiveContextKey = "adaptive_selected" + ctxKeyAdaptiveScores adaptiveContextKey = "adaptive_scores" + ctxKeyAdaptiveCircuitPermit adaptiveContextKey = "adaptive_circuit_permit" +) + +// AdaptiveSelectChannel 动态评分调度器主入口。 +// 所有回退必须调用 cacheGetRandomSatisfiedChannelLegacy，禁止再进 CacheGetRandomSatisfiedChannel。 +func AdaptiveSelectChannel(param *RetryParam) (*model.Channel, string, error) { + ctx := param.Ctx + + // 未开启完整自适应：仅 legacy（含「只开 shadow」旧行为，避免递归） + if !constant.AdaptiveBalanceEnabled { + return cacheGetRandomSatisfiedChannelLegacy(param) + } + + // 提取 group 和 model + group := common.GetContextKeyString(ctx, constant.ContextKeyUsingGroup) + if group == "" { + group = param.TokenGroup + } + modelName := param.ModelName + + // 获取该 group+model 下的可用渠道 + channels, err := getCandidateChannels(group, modelName, param) + if err != nil { + return nil, group, err + } + if len(channels) == 0 { + return cacheGetRandomSatisfiedChannelLegacy(param) + } + + // 获取亲和偏好 channel + preferredID := getPreferredChannelID(ctx, modelName, group) + + // 评分 + candidates := ScoreCandidates(channels, group, modelName, preferredID) + + usedIDs := getAdaptiveUsedChannels(ctx) + filtered, permits := filterAdaptiveCandidates( + candidates, group, modelName, preferredID, usedIDs, constant.AdaptiveBalanceShadowMode, + ) + + if len(filtered) == 0 { + if shouldFallbackToLegacy(len(candidates), len(filtered), constant.AdaptiveBalanceShadowMode) { + return cacheGetRandomSatisfiedChannelLegacy(param) + } + return nil, group, fmt.Errorf("adaptive: no available channels after circuit and retry filtering") + } + + // topK 加权随机选择 + selected := SelectTopKWeighted(filtered, 3) + if selected == nil { + releaseUnselectedCircuitPermits(permits, 0) + return nil, group, fmt.Errorf("adaptive: failed to select an eligible channel") + } + + // Shadow Mode：选择仍走旧逻辑，仅记录对比 + if constant.AdaptiveBalanceShadowMode { + oldCh, oldGroup, oldErr := cacheGetRandomSatisfiedChannelLegacy(param) + + // 采样日志 + if randFloat64() < adaptiveLogSample { + logAdaptiveCompare(ctx, modelName, group, selected, oldCh) + } + + // shadow mode never changes routing or acquires half-open permits. + if oldCh != nil { + addAdaptiveUsedChannel(ctx, oldCh.Id) + storeAdaptiveSelection(ctx, selected.Channel, group, candidates) + } + return oldCh, oldGroup, oldErr + } + + // 正常模式：使用动态选择的渠道 + selectGroup := group + ch := selected.Channel + permit := permits[ch.Id] + releaseUnselectedCircuitPermits(permits, ch.Id) + + addAdaptiveUsedChannel(ctx, ch.Id) + storeAdaptiveSelection(ctx, ch, group, candidates) + ctx.Set(string(ctxKeyAdaptiveCircuitPermit), permit) + + logger.LogDebug(ctx, "adaptive selected channel #%d (score=%.3f) for group=%s model=%s", + ch.Id, selected.Score, group, modelName) + + return ch, selectGroup, nil +} + +func filterAdaptiveCandidates( + candidates []CandidateScore, + group string, + modelName string, + preferredID int, + usedIDs []int, + shadowMode bool, +) ([]CandidateScore, map[int]CircuitPermit) { + filtered := make([]CandidateScore, 0, len(candidates)) + permits := make(map[int]CircuitPermit, len(candidates)) + for _, candidate := range candidates { + channelID := candidate.Channel.Id + if containsInt(usedIDs, channelID) { + continue + } + if shadowMode { + if IsCircuitOpen(channelID) || candidate.Score <= 0 { + continue + } + filtered = append(filtered, candidate) + continue + } + + permit, ok := AcquireCircuitPermit(channelID) + if !ok { + continue + } + if permit.HalfOpen { + candidate = scoreCandidate(candidate.Channel, group, modelName, preferredID, 0.5) + } + if candidate.Score <= 0 { + ReleaseCircuitPermit(permit) + continue + } + permits[channelID] = permit + filtered = append(filtered, candidate) + } + return filtered, permits +} + +func ReleaseAdaptiveCircuitPermit(c *gin.Context, channelID int) { + if c == nil || channelID <= 0 { + return + } + permitAny, ok := c.Get(string(ctxKeyAdaptiveCircuitPermit)) + permit, permitOK := permitAny.(CircuitPermit) + if !ok || !permitOK || permit.ChannelID != channelID { + return + } + ReleaseCircuitPermit(permit) + c.Set(string(ctxKeyAdaptiveCircuitPermit), CircuitPermit{}) +} + +func shouldFallbackToLegacy(candidateCount, filteredCount int, shadowMode bool) bool { + return candidateCount == 0 || (shadowMode && filteredCount == 0) +} + +func releaseUnselectedCircuitPermits(permits map[int]CircuitPermit, selectedChannelID int) { + for channelID, permit := range permits { + if channelID != selectedChannelID { + ReleaseCircuitPermit(permit) + } + } +} + +// getCandidateChannels 获取 group+model 全部候选（非单渠道路由） +func getCandidateChannels(group, modelName string, param *RetryParam) ([]*model.Channel, error) { + // auto 分组：优先用上下文已解析的 auto group，否则 legacy 解析一次 + if group == "auto" || param.TokenGroup == "auto" { + if g := common.GetContextKeyString(param.Ctx, constant.ContextKeyAutoGroup); g != "" { + group = g + } else { + // 用 legacy 解析 auto → 具体 group，再拉全量候选 + ch, selectGroup, err := cacheGetRandomSatisfiedChannelLegacy(param) + if err != nil { + return nil, err + } + if ch == nil { + return nil, nil + } + if selectGroup != "" { + group = selectGroup + } + // 继续用解析后的 group 拉全量；若失败至少返回当前渠道 + list, listErr := model.GetSatisfiedChannels(group, modelName, param.RequestPath) + if listErr != nil { + return []*model.Channel{ch}, nil + } + if len(list) == 0 { + return []*model.Channel{ch}, nil + } + return list, nil + } + } + + return model.GetSatisfiedChannels(group, modelName, param.RequestPath) +} + +// getPreferredChannelID 读取亲和偏好（如果有） +func getPreferredChannelID(ctx *gin.Context, modelName, group string) int { + if !common.MemoryCacheEnabled { + return 0 + } + id, found := GetPreferredChannelByAffinity(ctx, modelName, group) + if found { + return id + } + return 0 +} + +// getAdaptiveUsedChannels 获取本次请求已用过的渠道 ID 列表 +func getAdaptiveUsedChannels(c *gin.Context) []int { + v, ok := c.Get(string(ctxKeyAdaptiveUsedChannels)) + if !ok { + return nil + } + ids, _ := v.([]int) + return ids +} + +// addAdaptiveUsedChannel 记录本次请求使用过的渠道 +func addAdaptiveUsedChannel(c *gin.Context, channelID int) { + existing := getAdaptiveUsedChannels(c) + existing = append(existing, channelID) + c.Set(string(ctxKeyAdaptiveUsedChannels), existing) +} + +func MarkChannelUsed(c *gin.Context, channelID int) { + if c == nil || channelID <= 0 || containsInt(getAdaptiveUsedChannels(c), channelID) { + return + } + addAdaptiveUsedChannel(c, channelID) +} + +func adaptiveUsedChannelSet(c *gin.Context) map[int]struct{} { + used := getAdaptiveUsedChannels(c) + if len(used) == 0 { + return nil + } + excluded := make(map[int]struct{}, len(used)) + for _, channelID := range used { + excluded[channelID] = struct{}{} + } + return excluded +} + +// storeAdaptiveSelection 保存本次选择结果到上下文（供失败回写用） +func storeAdaptiveSelection(c *gin.Context, ch *model.Channel, group string, candidates []CandidateScore) { + c.Set(string(ctxKeyAdaptiveSelected), ch.Id) + c.Set(string(ctxKeyAdaptiveGroup), group) + if len(candidates) > 0 { + c.Set(string(ctxKeyAdaptiveScores), candidates) + } +} + +// logAdaptiveCompare shadow mode 日志 +func logAdaptiveCompare(c *gin.Context, modelName, group string, selected *CandidateScore, oldCh *model.Channel) { + oldID := 0 + if oldCh != nil { + oldID = oldCh.Id + } + logger.LogDebug(c, "[shadow] model=%s group=%s adaptive=#%d(%.3f) orig=#%d", + modelName, group, selected.Channel.Id, selected.Score, oldID) +} + +// RecordAdaptiveResult 请求完成后回调：更新指标 + 熔断状态 +func RecordAdaptiveResult(c *gin.Context, channelID int, group, modelName string, statusCode int, latency time.Duration, err error) { + if !constant.AdaptiveBalanceEnabled { + return + } + if channelID <= 0 { + return + } + + succeeded := err == nil && statusCode < 400 + if succeeded { + ObserveSuccess(channelID, group, modelName, latency) + } else { + ObserveFailure(channelID, group, modelName, statusCode, latency) + } + + if constant.AdaptiveBalanceShadowMode { + return + } + permitAny, ok := c.Get(string(ctxKeyAdaptiveCircuitPermit)) + permit, permitOK := permitAny.(CircuitPermit) + if !ok || !permitOK || permit.ChannelID != channelID { + return + } + if succeeded { + RecordCircuitSuccessWithPermit(permit) + } else if statusCode >= 500 || statusCode == 429 { + RecordCircuitFailureWithPermit(permit, fmt.Sprintf("HTTP %d", statusCode)) + } else { + // A client/input error still proves the upstream is reachable. Do not + // leave a half-open permit stuck or preserve an old failure streak. + RecordCircuitSuccessWithPermit(permit) + } +} + +// containsInt 检查 int 切片是否包含某值 +func containsInt(slice []int, val int) bool { + for _, v := range slice { + if v == val { + return true + } + } + return false +} + +// randFloat64 生成 [0,1) 随机数 +var randFloat64 = func() float64 { + return rand.Float64() +} diff --git a/service/channel_adaptive_test.go b/service/channel_adaptive_test.go new file mode 100644 index 000000000000..16f6ba202bb5 --- /dev/null +++ b/service/channel_adaptive_test.go @@ -0,0 +1,275 @@ +package service + +import ( + "fmt" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/gin-gonic/gin" + + "github.com/stretchr/testify/require" +) + +// 测试辅助：构造测试渠道 +func testChannel(id int, weight uint, priority int64) *model.Channel { + w := weight + return &model.Channel{ + Id: id, + Weight: &w, + Priority: &priority, + Name: fmt.Sprintf("ch-%d", id), + Status: common.ChannelStatusEnabled, + } +} + +func init() { + // 测试前重置全局状态 + globalSnapshot.mu.Lock() + globalSnapshot.metrics = make(map[metricsKey]*ChannelMetrics) + globalSnapshot.mu.Unlock() + + constant.EwmaAlpha = 0.3 + constant.MaxChannelConcurrency = 10 + constant.ChannelCircuitBreakerEnabled = true + constant.AdaptiveBalanceEnabled = true + constant.AdaptiveBalanceShadowMode = false +} + +// 测试1：高延迟渠道会被降权 +func TestHighLatencyDowngraded(t *testing.T) { + ch1 := testChannel(1, 10, 100) + ch2 := testChannel(2, 10, 100) + + // ch1 低延迟，ch2 高延迟 + ObserveSuccess(1, "test", "gpt-4", 200*time.Millisecond) + ObserveSuccess(1, "test", "gpt-4", 150*time.Millisecond) + ObserveSuccess(1, "test", "gpt-4", 180*time.Millisecond) + ObserveSuccess(2, "test", "gpt-4", 5*time.Second) + ObserveSuccess(2, "test", "gpt-4", 6*time.Second) + ObserveSuccess(2, "test", "gpt-4", 4*time.Second) + + channels := []*model.Channel{ch1, ch2} + candidates := ScoreCandidates(channels, "test", "gpt-4", 0) + + require.GreaterOrEqual(t, len(candidates), 2) + // ch1（低延迟）的分数应显著高于 ch2 + require.Equal(t, 1, candidates[0].Channel.Id, "expected ch1 (low latency) to rank first") + scoreDiff := candidates[0].Score - candidates[1].Score + require.Greater(t, scoreDiff, 0.1, "expected significant score difference") + t.Logf("ch1 (low latency) score=%.4f, ch2 (high latency) score=%.4f", candidates[0].Score, candidates[1].Score) +} + +// 测试2：429 渠道不会完全排除但会被降权 +func TestRateLimitedChannelDowngraded(t *testing.T) { + ch1 := testChannel(10, 10, 100) + ch2 := testChannel(11, 10, 100) + + // ch1 正常，ch2 有 429 + ObserveSuccess(10, "test", "gpt-4", 300*time.Millisecond) + ObserveSuccess(10, "test", "gpt-4", 250*time.Millisecond) + ObserveFailure(11, "test", "gpt-4", 429, 100*time.Millisecond) + ObserveFailure(11, "test", "gpt-4", 429, 100*time.Millisecond) + ObserveFailure(11, "test", "gpt-4", 429, 100*time.Millisecond) + + channels := []*model.Channel{ch1, ch2} + candidates := ScoreCandidates(channels, "test", "gpt-4", 0) + + require.GreaterOrEqual(t, len(candidates), 2) + // 正常渠道应排名更高 + require.Equal(t, 10, candidates[0].Channel.Id, "expected ch10 (normal) to rank first") + t.Logf("ch10 (normal) score=%.4f rate_limit_factor=%.4f", candidates[0].Score, candidates[0].RateLimitFactor) + t.Logf("ch11 (429) score=%.4f rate_limit_factor=%.4f", candidates[1].Score, candidates[1].RateLimitFactor) + require.Less(t, candidates[1].RateLimitFactor, candidates[0].RateLimitFactor) +} + +// 测试3：熔断渠道不会被选中 +func TestCircuitBreakerChannelExcluded(t *testing.T) { + ch := testChannel(20, 10, 100) + + // 模拟三次连续失败，触发熔断 + RecordCircuitFailure(20, "500 Internal Server Error") + RecordCircuitFailure(20, "500 Internal Server Error") + RecordCircuitFailure(20, "500 Internal Server Error") + + require.True(t, IsCircuitOpen(20), "expected circuit to be open after 3 consecutive failures") + + // 评分中应过滤掉熔断渠道 + channels := []*model.Channel{ch} + _ = ScoreCandidates(channels, "test", "gpt-4", 0) + + // channel_adaptive.go 中过滤逻辑会跳过 open 渠道 + require.True(t, IsCircuitOpen(20)) +} + +// 测试4：多渠道重试不会重复选择同一个渠道 +func TestNoDuplicateChannelInRetry(t *testing.T) { + ch1 := testChannel(30, 10, 100) + ch2 := testChannel(31, 10, 100) + ch3 := testChannel(32, 10, 100) + + // 全部成功 + for _, id := range []int{30, 31, 32} { + ObserveSuccess(id, "test", "gpt-4", 200*time.Millisecond) + } + + channels := []*model.Channel{ch1, ch2, ch3} + candidates := ScoreCandidates(channels, "test", "gpt-4", 0) + + // 模拟已使用的渠道 + usedIDs := []int{30} + + // 过滤掉已使用的渠道 + var filtered []CandidateScore + for _, c := range candidates { + if containsInt(usedIDs, c.Channel.Id) { + continue + } + filtered = append(filtered, c) + } + + require.Len(t, filtered, 2) + selected := SelectTopKWeighted(filtered, 3) + require.NotNil(t, selected) + require.NotEqual(t, 30, selected.Channel.Id) + t.Logf("selected ch%d (ch30 excluded)", selected.Channel.Id) +} + +// 测试5：TopK 加权随机不会全部集中在最高分渠道 +func TestTopKWeightedRandomFairness(t *testing.T) { + channels := make([]*model.Channel, 10) + for i := 0; i < 10; i++ { + channels[i] = testChannel(100+i, 10, 100) + ObserveSuccess(100+i, "test", "gpt-4", time.Duration(200+(i*100))*time.Millisecond) + } + + candidates := ScoreCandidates(channels, "test", "gpt-4", 0) + require.GreaterOrEqual(t, len(candidates), 10) + + // 模拟多次选择，统计分布 + selectionCount := make(map[int]int) + trials := 1000 + for i := 0; i < trials; i++ { + selected := SelectTopKWeighted(candidates, 3) + if selected != nil { + selectionCount[selected.Channel.Id]++ + } + } + + // TopK 的前三名应该都有一定比例 + for _, c := range candidates[:3] { + count := selectionCount[c.Channel.Id] + ratio := float64(count) / float64(trials) + t.Logf("ch%d selection rate: %.2f%% (score=%.3f)", c.Channel.Id, ratio*100, c.Score) + require.GreaterOrEqual(t, ratio, 0.05, "ch%d selected too few times", c.Channel.Id) + } +} + +// 测试6：EWMA 计算正确 +func TestEwmaUpdate(t *testing.T) { + alpha := 0.3 + + // 初始 1.0，观察到 0.5 + result := EwmaUpdate(1.0, 0.5, alpha) + expected := 0.3*0.5 + 0.7*1.0 // = 0.85 + require.InDelta(t, expected, result, 0.001) + + // 再次衰减 + result2 := EwmaUpdate(result, 0.5, alpha) + expected2 := 0.3*0.5 + 0.7*0.85 // = 0.745 + require.InDelta(t, expected2, result2, 0.001) +} + +// 测试7：延迟桶边界 + +func TestAdaptiveFilteringRecomputesHalfOpenProbeScore(t *testing.T) { + channelID := 220 + circuitBreakers.Delete(channelID) + t.Cleanup(func() { circuitBreakers.Delete(channelID) }) + + cb := getCircuitBreaker(channelID) + cb.mu.Lock() + cb.State = CircuitOpen + cb.OpenUntil = time.Now().Add(-time.Second) + cb.ConsecutiveFailure = 3 + cb.mu.Unlock() + + ch := testChannel(channelID, 10, 100) + candidates := ScoreCandidates([]*model.Channel{ch}, "test", "gpt-4", 0) + require.Len(t, candidates, 1) + require.Zero(t, candidates[0].Score) + + filtered, permits := filterAdaptiveCandidates( + candidates, + "test", + "gpt-4", + 0, + nil, + false, + ) + require.Len(t, filtered, 1) + require.Greater(t, filtered[0].Score, 0.0) + require.True(t, permits[channelID].HalfOpen) + ReleaseCircuitPermit(permits[channelID]) +} + +func TestAdaptiveLegacyFallbackPolicyDoesNotBypassFiltering(t *testing.T) { + require.True(t, shouldFallbackToLegacy(0, 0, false)) + require.False(t, shouldFallbackToLegacy(2, 0, false)) + require.True(t, shouldFallbackToLegacy(2, 0, true)) +} + +func TestOpenCircuitIgnoresLateSuccess(t *testing.T) { + channelID := 221 + circuitBreakers.Delete(channelID) + t.Cleanup(func() { circuitBreakers.Delete(channelID) }) + + RecordCircuitFailure(channelID, "500") + RecordCircuitFailure(channelID, "500") + RecordCircuitFailure(channelID, "500") + require.True(t, IsCircuitOpen(channelID)) + + RecordCircuitSuccess(channelID) + state, failures, _ := GetCircuitState(channelID) + require.Equal(t, CircuitOpen, state) + require.Equal(t, 3, failures) +} + +func TestGetMetricsReturnsImmutableSnapshot(t *testing.T) { + channelID := 222 + ObserveSuccess(channelID, "snapshot", "gpt-4", 250*time.Millisecond) + + first := GetMetrics(channelID, "snapshot", "gpt-4") + first.SuccessRate = 0 + first.AvgLatency = 99 * time.Second + + second := GetMetrics(channelID, "snapshot", "gpt-4") + require.Greater(t, second.SuccessRate, 0.0) + require.NotEqual(t, 99*time.Second, second.AvgLatency) +} + +func TestHalfOpenClientErrorReleasesProbeAndClosesCircuit(t *testing.T) { + channelID := 223 + circuitBreakers.Delete(channelID) + t.Cleanup(func() { circuitBreakers.Delete(channelID) }) + + cb := getCircuitBreaker(channelID) + cb.mu.Lock() + cb.State = CircuitOpen + cb.OpenUntil = time.Now().Add(-time.Second) + cb.ConsecutiveFailure = 3 + cb.mu.Unlock() + permit, ok := AcquireCircuitPermit(channelID) + require.True(t, ok) + require.True(t, permit.HalfOpen) + + ctx, _ := gin.CreateTestContext(nil) + ctx.Set(string(ctxKeyAdaptiveCircuitPermit), permit) + RecordAdaptiveResult(ctx, channelID, "test", "gpt-4", 400, time.Millisecond, fmt.Errorf("bad request")) + + state, failures, _ := GetCircuitState(channelID) + require.Equal(t, CircuitClosed, state) + require.Zero(t, failures) +} diff --git a/service/channel_affinity.go b/service/channel_affinity.go index 96ec13e248cc..076ac468cfaf 100644 --- a/service/channel_affinity.go +++ b/service/channel_affinity.go @@ -1,6 +1,7 @@ package service import ( + "context" "fmt" "hash/fnv" "regexp" @@ -27,6 +28,7 @@ const ( ginKeyChannelAffinitySkipRetry = "channel_affinity_skip_retry_on_failure" channelAffinityCacheNamespace = "new-api:channel_affinity:v1" + channelAffinityRedisLRUIndex = "new-api:channel_affinity_lru:v1" channelAffinityUsageCacheStatsNamespace = "new-api:channel_affinity_usage_cache_stats:v1" ) @@ -207,6 +209,13 @@ func ClearChannelAffinityCacheAll() int { common.SysError(fmt.Sprintf("channel affinity cache delete many failed: err=%v", err)) } } + if common.RedisEnabled && common.RDB != nil { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + if err := common.RDB.Del(ctx, channelAffinityRedisLRUIndex).Err(); err != nil { + common.SysError(fmt.Sprintf("channel affinity LRU index clear failed: err=%v", err)) + } + } return len(keys) } @@ -238,10 +247,42 @@ func ClearChannelAffinityCacheByRuleName(ruleName string) (int, error) { } cache := getChannelAffinityCache() - deleted, err := cache.DeleteByPrefix(ruleName) + if !common.RedisEnabled || common.RDB == nil { + return cache.DeleteByPrefix(ruleName) + } + keys, err := cache.Keys() if err != nil { return 0, err } + prefix := cache.FullKey(ruleName) + ":" + matched := make([]string, 0, len(keys)) + for _, key := range keys { + if strings.HasPrefix(key, prefix) { + matched = append(matched, key) + } + } + if len(matched) == 0 { + return 0, nil + } + deletedMap, err := cache.DeleteMany(matched) + if err != nil { + return 0, err + } + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + members := make([]interface{}, 0, len(matched)) + for _, key := range matched { + members = append(members, key) + } + if err := common.RDB.ZRem(ctx, channelAffinityRedisLRUIndex, members...).Err(); err != nil { + return 0, err + } + deleted := 0 + for _, ok := range deletedMap { + if ok { + deleted++ + } + } return deleted, nil } @@ -334,18 +375,35 @@ func extractChannelAffinityValue(c *gin.Context, src operation_setting.ChannelAf } } -func buildChannelAffinityCacheKeySuffix(rule operation_setting.ChannelAffinityRule, modelName string, usingGroup string, affinityValue string) string { +func fixedAffinityKeyPart(value string) string { + return fmt.Sprintf("%x", common.Sha256Raw([]byte(value))) +} + +func channelAffinityCredentialScope(c *gin.Context) string { + if c != nil { + if tokenID := c.GetInt("token_id"); tokenID > 0 { + return fmt.Sprintf("token:%d", tokenID) + } + if userID := c.GetInt("id"); userID > 0 { + return fmt.Sprintf("user:%d", userID) + } + } + return "anonymous" +} + +func buildChannelAffinityCacheKeySuffix(c *gin.Context, rule operation_setting.ChannelAffinityRule, modelName string, usingGroup string, affinityValue string) string { parts := make([]string, 0, 4) if rule.IncludeRuleName && rule.Name != "" { parts = append(parts, rule.Name) } if rule.IncludeModelName && modelName != "" { - parts = append(parts, modelName) + parts = append(parts, fixedAffinityKeyPart(modelName)[:16]) } if rule.IncludeUsingGroup && usingGroup != "" { - parts = append(parts, usingGroup) + parts = append(parts, fixedAffinityKeyPart(usingGroup)[:16]) } - parts = append(parts, affinityValue) + scopedValue := channelAffinityCredentialScope(c) + "\x00" + affinityValue + parts = append(parts, fixedAffinityKeyPart(scopedValue)) return strings.Join(parts, ":") } @@ -591,7 +649,7 @@ func GetPreferredChannelByAffinity(c *gin.Context, modelName string, usingGroup if ttlSeconds <= 0 { ttlSeconds = setting.DefaultTTLSeconds } - cacheKeySuffix := buildChannelAffinityCacheKeySuffix(rule, modelName, usingGroup, affinityValue) + cacheKeySuffix := buildChannelAffinityCacheKeySuffix(c, rule, modelName, usingGroup, affinityValue) cacheKeyFull := channelAffinityCacheNamespace + ":" + cacheKeySuffix setChannelAffinityContext(c, channelAffinityMeta{ CacheKey: cacheKeyFull, @@ -734,11 +792,50 @@ func RecordChannelAffinity(c *gin.Context, channelID int) { ttlSeconds = 3600 } cache := getChannelAffinityCache() - if err := cache.SetWithTTL(cacheKey, channelID, time.Duration(ttlSeconds)*time.Second); err != nil { + if err := setChannelAffinityWithLimit(cache, cacheKey, channelID, time.Duration(ttlSeconds)*time.Second, setting.MaxEntries); err != nil { common.SysError(fmt.Sprintf("channel affinity cache set failed: key=%s, err=%v", cacheKey, err)) } } +const channelAffinityRedisSetScript = ` +redis.call('SET', KEYS[1], ARGV[1], 'PX', ARGV[2]) +redis.call('ZADD', KEYS[2], ARGV[3], KEYS[1]) +local max_entries = tonumber(ARGV[4]) +if max_entries and max_entries > 0 then + local count = redis.call('ZCARD', KEYS[2]) + local overflow = count - max_entries + if overflow > 0 then + local victims = redis.call('ZRANGE', KEYS[2], 0, overflow - 1) + for _, key in ipairs(victims) do + redis.call('DEL', key) + end + redis.call('ZREMRANGEBYRANK', KEYS[2], 0, overflow - 1) + end +end +return 1 +` + +func setChannelAffinityWithLimit(cache *cachex.HybridCache[int], key string, channelID int, ttl time.Duration, maxEntries int) error { + if !common.RedisEnabled || common.RDB == nil { + return cache.SetWithTTL(key, channelID, ttl) + } + if maxEntries <= 0 { + maxEntries = 100_000 + } + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + _, err := common.RDB.Eval( + ctx, + channelAffinityRedisSetScript, + []string{cache.FullKey(key), channelAffinityRedisLRUIndex}, + strconv.Itoa(channelID), + strconv.FormatInt(ttl.Milliseconds(), 10), + strconv.FormatInt(time.Now().UnixNano(), 10), + strconv.Itoa(maxEntries), + ).Result() + return err +} + type ChannelAffinityUsageCacheStats struct { RuleName string `json:"rule_name"` UsingGroup string `json:"using_group"` diff --git a/service/channel_affinity_template_test.go b/service/channel_affinity_template_test.go index fb703a24e720..ed148431f7ed 100644 --- a/service/channel_affinity_template_test.go +++ b/service/channel_affinity_template_test.go @@ -205,7 +205,7 @@ func TestGetPreferredChannelByAffinity_RequestHeaderKeySource(t *testing.T) { } affinityValue := fmt.Sprintf("header-hit-%d", time.Now().UnixNano()) - cacheKeySuffix := buildChannelAffinityCacheKeySuffix(rule, "gpt-5", "default", affinityValue) + cacheKeySuffix := buildChannelAffinityCacheKeySuffix(nil, rule, "gpt-5", "default", affinityValue) cache := getChannelAffinityCache() require.NoError(t, cache.SetWithTTL(cacheKeySuffix, 9528, time.Minute)) @@ -236,6 +236,28 @@ func TestGetPreferredChannelByAffinity_RequestHeaderKeySource(t *testing.T) { require.Equal(t, buildChannelAffinityKeyHint(affinityValue), meta.KeyHint) } +func TestChannelAffinityCacheKeyIsScopedAndBounded(t *testing.T) { + rule := operation_setting.ChannelAffinityRule{ + Name: "trace-affinity", + IncludeRuleName: true, + IncludeModelName: true, + IncludeUsingGroup: true, + } + longTrace := strings.Repeat("attacker-controlled-trace-", 100) + + ctx1, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx1.Set("token_id", 101) + ctx2, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx2.Set("token_id", 202) + + key1 := buildChannelAffinityCacheKeySuffix(ctx1, rule, strings.Repeat("model", 100), "default", longTrace) + key2 := buildChannelAffinityCacheKeySuffix(ctx2, rule, strings.Repeat("model", 100), "default", longTrace) + + require.NotEqual(t, key1, key2) + require.NotContains(t, key1, longTrace) + require.Less(t, len(key1), 160) +} + func TestClearCurrentChannelAffinityCache(t *testing.T) { gin.SetMode(gin.TestMode) @@ -280,7 +302,7 @@ func TestChannelAffinityHitCodexTemplatePassHeadersEffective(t *testing.T) { require.NotNil(t, codexRule) affinityValue := fmt.Sprintf("pc-hit-%d", time.Now().UnixNano()) - cacheKeySuffix := buildChannelAffinityCacheKeySuffix(*codexRule, "gpt-5", "default", affinityValue) + cacheKeySuffix := buildChannelAffinityCacheKeySuffix(nil, *codexRule, "gpt-5", "default", affinityValue) cache := getChannelAffinityCache() require.NoError(t, cache.SetWithTTL(cacheKeySuffix, 9527, time.Minute)) diff --git a/service/channel_affinity_usage_cache_test.go b/service/channel_affinity_usage_cache_test.go index 64d3d715b547..2af84f6da178 100644 --- a/service/channel_affinity_usage_cache_test.go +++ b/service/channel_affinity_usage_cache_test.go @@ -4,7 +4,6 @@ import ( "fmt" "net/http/httptest" "testing" - "time" "github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/types" @@ -12,7 +11,16 @@ import ( "github.com/stretchr/testify/require" ) -func buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP string) *gin.Context { +func buildChannelAffinityStatsContextForTest(t *testing.T) (*gin.Context, string, string, string) { + t.Helper() + ruleName := fmt.Sprintf("rule_%s", t.Name()) + usingGroup := "default" + keyFP := fmt.Sprintf("fp_%s", t.Name()) + entryKey := channelAffinityUsageCacheEntryKey(ruleName, usingGroup, keyFP) + t.Cleanup(func() { + _, _ = getChannelAffinityUsageCacheStatsCache().DeleteMany([]string{entryKey}) + }) + rec := httptest.NewRecorder() ctx, _ := gin.CreateTestContext(rec) setChannelAffinityContext(ctx, channelAffinityMeta{ @@ -22,14 +30,11 @@ func buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP string) UsingGroup: usingGroup, KeyFingerprint: keyFP, }) - return ctx + return ctx, ruleName, usingGroup, keyFP } func TestObserveChannelAffinityUsageCacheByRelayFormat_ClaudeMode(t *testing.T) { - ruleName := fmt.Sprintf("rule_%d", time.Now().UnixNano()) - usingGroup := "default" - keyFP := fmt.Sprintf("fp_%d", time.Now().UnixNano()) - ctx := buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP) + ctx, ruleName, usingGroup, keyFP := buildChannelAffinityStatsContextForTest(t) usage := &dto.Usage{ PromptTokens: 100, @@ -53,10 +58,7 @@ func TestObserveChannelAffinityUsageCacheByRelayFormat_ClaudeMode(t *testing.T) } func TestObserveChannelAffinityUsageCacheByRelayFormat_MixedMode(t *testing.T) { - ruleName := fmt.Sprintf("rule_%d", time.Now().UnixNano()) - usingGroup := "default" - keyFP := fmt.Sprintf("fp_%d", time.Now().UnixNano()) - ctx := buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP) + ctx, ruleName, usingGroup, keyFP := buildChannelAffinityStatsContextForTest(t) openAIUsage := &dto.Usage{ PromptTokens: 100, @@ -83,10 +85,7 @@ func TestObserveChannelAffinityUsageCacheByRelayFormat_MixedMode(t *testing.T) { } func TestObserveChannelAffinityUsageCacheByRelayFormat_UnsupportedModeKeepsEmpty(t *testing.T) { - ruleName := fmt.Sprintf("rule_%d", time.Now().UnixNano()) - usingGroup := "default" - keyFP := fmt.Sprintf("fp_%d", time.Now().UnixNano()) - ctx := buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP) + ctx, ruleName, usingGroup, keyFP := buildChannelAffinityStatsContextForTest(t) usage := &dto.Usage{ PromptTokens: 100, diff --git a/service/channel_circuit.go b/service/channel_circuit.go new file mode 100644 index 000000000000..4058932a912d --- /dev/null +++ b/service/channel_circuit.go @@ -0,0 +1,245 @@ +package service + +import ( + "sync" + "time" + + "github.com/QuantumNous/new-api/constant" +) + +// CircuitState 熔断器状态 +type CircuitState string + +const ( + CircuitClosed CircuitState = "closed" // 正常 + CircuitOpen CircuitState = "open" // 熔断打开，不选 + CircuitHalfOpen CircuitState = "half_open" // 半开，允许探测 +) + +// ChannelCircuitBreaker 渠道熔断器（本地状态 + Redis 同步） +// 约定：请求路径上只读本地状态，不访问 Redis。 +type ChannelCircuitBreaker struct { + mu sync.RWMutex + + State CircuitState + ConsecutiveFailure int // 连续失败计数 + OpenUntil time.Time // open 状态过期时间 + HalfOpenLimit int // half-open 最大探测数 + HalfOpenInFlight int // half-open 进行中的探测数 + HalfOpenSince time.Time // when current half-open probe started + LastError string // 最近一次错误信息 + Generation uint64 // invalidates results from requests started before a transition +} + +type CircuitPermit struct { + ChannelID int + Generation uint64 + HalfOpen bool +} + +var ( + circuitBreakers sync.Map // map[int]*ChannelCircuitBreaker, key=channelID +) + +// getCircuitBreaker 获取或创建渠道熔断器 +func getCircuitBreaker(channelID int) *ChannelCircuitBreaker { + v, _ := circuitBreakers.LoadOrStore(channelID, &ChannelCircuitBreaker{ + State: CircuitClosed, + HalfOpenLimit: 1, + }) + return v.(*ChannelCircuitBreaker) +} + +// IsCircuitOpen 判断渠道是否熔断（请求路径使用，读本地状态） +func IsCircuitOpen(channelID int) bool { + if !constant.ChannelCircuitBreakerEnabled { + return false + } + + cb := getCircuitBreaker(channelID) + cb.mu.RLock() + defer cb.mu.RUnlock() + + if cb.State == CircuitClosed { + return false + } + + if cb.State == CircuitOpen && time.Now().After(cb.OpenUntil) { + // Cooldown elapsed: still report open so selector must call ProbeHalfOpen. + return true + } + + return true +} + + +// RecordSuccess 成功调用 -> 重置熔断状态 +func RecordCircuitSuccess(channelID int) { + if !constant.ChannelCircuitBreakerEnabled { + return + } + + cb := getCircuitBreaker(channelID) + cb.mu.Lock() + defer cb.mu.Unlock() + + // A success without a selection-time permit may be a late result from a + // request that started before the circuit opened. It must not close an + // open/half-open circuit. + if cb.State != CircuitClosed { + return + } + cb.ConsecutiveFailure = 0 + cb.OpenUntil = time.Time{} + cb.LastError = "" +} + +func RecordCircuitSuccessWithPermit(permit CircuitPermit) { + if !constant.ChannelCircuitBreakerEnabled || permit.ChannelID <= 0 { + return + } + + cb := getCircuitBreaker(permit.ChannelID) + cb.mu.Lock() + defer cb.mu.Unlock() + + if permit.Generation != cb.Generation { + return + } + if permit.HalfOpen { + if cb.State != CircuitHalfOpen { + return + } + if cb.HalfOpenInFlight > 0 { + cb.HalfOpenInFlight-- + } + cb.State = CircuitClosed + cb.Generation++ + cb.HalfOpenSince = time.Time{} + } else if cb.State != CircuitClosed { + return + } + + cb.ConsecutiveFailure = 0 + cb.OpenUntil = time.Time{} + cb.LastError = "" +} + +// RecordFailure 失败调用 -> 可能触发熔断 +func RecordCircuitFailure(channelID int, errMsg string) { + if !constant.ChannelCircuitBreakerEnabled { + return + } + cb := getCircuitBreaker(channelID) + cb.mu.RLock() + permit := CircuitPermit{ChannelID: channelID, Generation: cb.Generation} + cb.mu.RUnlock() + RecordCircuitFailureWithPermit(permit, errMsg) +} + +func RecordCircuitFailureWithPermit(permit CircuitPermit, errMsg string) { + if !constant.ChannelCircuitBreakerEnabled || permit.ChannelID <= 0 { + return + } + + cb := getCircuitBreaker(permit.ChannelID) + cb.mu.Lock() + defer cb.mu.Unlock() + + if permit.Generation != cb.Generation { + return + } + if permit.HalfOpen { + if cb.State != CircuitHalfOpen { + return + } + if cb.HalfOpenInFlight > 0 { + cb.HalfOpenInFlight-- + } + cb.HalfOpenSince = time.Time{} + cb.State = CircuitOpen + cb.Generation++ + cb.ConsecutiveFailure++ + cb.LastError = errMsg + cb.OpenUntil = time.Now().Add(time.Duration(constant.ChannelCooldownSeconds) * time.Second) + return + } + if cb.State != CircuitClosed { + return + } + + cb.ConsecutiveFailure++ + cb.LastError = errMsg + + // closed 状态下连续失败达到阈值 -> open + threshold := 3 + if cb.ConsecutiveFailure >= threshold { + cb.State = CircuitOpen + cb.Generation++ + cb.OpenUntil = time.Now().Add(time.Duration(constant.ChannelCooldownSeconds) * time.Second) + } + + // 如果配置了熔断但未启用，不做任何事 +} + +func AcquireCircuitPermit(channelID int) (CircuitPermit, bool) { + if !constant.ChannelCircuitBreakerEnabled { + return CircuitPermit{ChannelID: channelID}, true + } + + cb := getCircuitBreaker(channelID) + cb.mu.Lock() + defer cb.mu.Unlock() + + if cb.State == CircuitClosed { + return CircuitPermit{ChannelID: channelID, Generation: cb.Generation}, true + } + if cb.State == CircuitOpen && time.Now().After(cb.OpenUntil) { + cb.State = CircuitHalfOpen + cb.HalfOpenInFlight = 0 + cb.HalfOpenSince = time.Time{} + } + if cb.State != CircuitHalfOpen { + return CircuitPermit{}, false + } + + const halfOpenProbeTimeout = 60 * time.Second + if cb.HalfOpenInFlight > 0 && !cb.HalfOpenSince.IsZero() && time.Since(cb.HalfOpenSince) > halfOpenProbeTimeout { + cb.HalfOpenInFlight = 0 + cb.HalfOpenSince = time.Time{} + } + if cb.HalfOpenInFlight >= cb.HalfOpenLimit { + return CircuitPermit{}, false + } + + cb.HalfOpenInFlight++ + cb.HalfOpenSince = time.Now() + return CircuitPermit{ChannelID: channelID, Generation: cb.Generation, HalfOpen: true}, true +} + +func ReleaseCircuitPermit(permit CircuitPermit) { + if !constant.ChannelCircuitBreakerEnabled || !permit.HalfOpen || permit.ChannelID <= 0 { + return + } + cb := getCircuitBreaker(permit.ChannelID) + cb.mu.Lock() + defer cb.mu.Unlock() + if cb.State != CircuitHalfOpen || cb.Generation != permit.Generation { + return + } + if cb.HalfOpenInFlight > 0 { + cb.HalfOpenInFlight-- + } + if cb.HalfOpenInFlight == 0 { + cb.HalfOpenSince = time.Time{} + } +} + + +// GetCircuitState 读取熔断状态（供日志/观测使用） +func GetCircuitState(channelID int) (CircuitState, int, string) { + cb := getCircuitBreaker(channelID) + cb.mu.RLock() + defer cb.mu.RUnlock() + return cb.State, cb.ConsecutiveFailure, cb.LastError +} diff --git a/service/channel_metrics.go b/service/channel_metrics.go new file mode 100644 index 000000000000..0d9c09f03b9a --- /dev/null +++ b/service/channel_metrics.go @@ -0,0 +1,255 @@ +package service + +import ( + "fmt" + "sync" + "sync/atomic" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" +) + +// ChannelMetrics 渠道运行时指标，按 (channelID, group, model) 分桶 +type ChannelMetrics struct { + mu sync.Mutex `json:"-"` // 保护并发写 + SuccessRate float64 `json:"success_rate"` // EWMA + ErrorRate float64 `json:"error_rate"` // EWMA + RateLimitRate float64 `json:"rate_limit_rate"` // EWMA 429 率 + Status5xxRate float64 `json:"status_5xx_rate"` // EWMA 5xx 率 + AvgLatency time.Duration `json:"avg_latency"` // EWMA 平均延迟 + SampleCount int64 `json:"sample_count"` // 总样本数 + LastSeen time.Time `json:"last_seen"` + +} + +// LocalMetricsSnapshot 进程内本地指标快照，定期从 Redis sync 或直接从本地累加 +type LocalMetricsSnapshot struct { + mu sync.RWMutex + metrics map[metricsKey]*ChannelMetrics + updatedAt time.Time +} + +type metricsKey struct { + ChannelID int + Group string + Model string +} + +var globalSnapshot = &LocalMetricsSnapshot{ + metrics: make(map[metricsKey]*ChannelMetrics), +} + +// ensureKey 获取或创建指定 key 的指标桶 +func (s *LocalMetricsSnapshot) ensureKey(key metricsKey) *ChannelMetrics { + s.mu.Lock() + defer s.mu.Unlock() + m, ok := s.metrics[key] + if !ok { + m = &ChannelMetrics{ + SuccessRate: 1.0, // 冷启动默认信任 + AvgLatency: 500 * time.Millisecond, + } + s.metrics[key] = m + } + return m +} + +// EwmaUpdate 更新指标的 EWMA 值 +func EwmaUpdate(current, observed, alpha float64) float64 { + if alpha <= 0 || alpha > 1 { + alpha = constant.EwmaAlpha + } + return alpha*observed + (1-alpha)*current +} + +// ObserveSuccess 记录一次成功调用 +func ObserveSuccess(channelID int, group, model string, latency time.Duration) { + alpha := constant.EwmaAlpha + key := metricsKey{channelID, group, model} + m := globalSnapshot.ensureKey(key) + + m.mu.Lock() + defer m.mu.Unlock() + + m.SampleCount++ + m.LastSeen = time.Now() + + // 更新延迟 EWMA + if m.AvgLatency == 0 { + m.AvgLatency = latency + } else { + m.AvgLatency = time.Duration(EwmaUpdate(float64(m.AvgLatency), float64(latency), alpha)) + } + + + // 更新成功率 EWMA + m.SuccessRate = EwmaUpdate(m.SuccessRate, 1.0, alpha) + m.ErrorRate = EwmaUpdate(m.ErrorRate, 0, alpha) + m.RateLimitRate = EwmaUpdate(m.RateLimitRate, 0, alpha) + m.Status5xxRate = EwmaUpdate(m.Status5xxRate, 0, alpha) +} + +// ObserveFailure 记录一次失败 +func ObserveFailure(channelID int, group, model string, statusCode int, latency time.Duration) { + alpha := constant.EwmaAlpha + key := metricsKey{channelID, group, model} + m := globalSnapshot.ensureKey(key) + + m.mu.Lock() + defer m.mu.Unlock() + + m.SampleCount++ + m.LastSeen = time.Now() + + // 更新延迟 + if m.AvgLatency == 0 { + m.AvgLatency = latency + } else { + m.AvgLatency = time.Duration(EwmaUpdate(float64(m.AvgLatency), float64(latency), alpha)) + } + + m.SuccessRate = EwmaUpdate(m.SuccessRate, 0, alpha) + m.ErrorRate = EwmaUpdate(m.ErrorRate, 1, alpha) + + switch { + case statusCode == 429: + m.RateLimitRate = EwmaUpdate(m.RateLimitRate, 1, alpha) + m.Status5xxRate = EwmaUpdate(m.Status5xxRate, 0, alpha) + case statusCode >= 500: + m.Status5xxRate = EwmaUpdate(m.Status5xxRate, 1, alpha) + m.RateLimitRate = EwmaUpdate(m.RateLimitRate, 0, alpha) + default: + m.RateLimitRate = EwmaUpdate(m.RateLimitRate, 0, alpha) + m.Status5xxRate = EwmaUpdate(m.Status5xxRate, 0, alpha) + } +} + +// GetMetrics 读取指定 (channelID, group, model) 的指标快照 +func GetMetrics(channelID int, group, model string) *ChannelMetrics { + key := metricsKey{channelID, group, model} + globalSnapshot.mu.RLock() + m, ok := globalSnapshot.metrics[key] + globalSnapshot.mu.RUnlock() + if ok { + return snapshotChannelMetrics(m) + } + + // 尝试回退到 (channelID, group) + key2 := metricsKey{channelID, group, ""} + globalSnapshot.mu.RLock() + m2, ok2 := globalSnapshot.metrics[key2] + globalSnapshot.mu.RUnlock() + if ok2 { + return snapshotChannelMetrics(m2) + } + + // 回退到 (channelID) + key3 := metricsKey{channelID, "", ""} + globalSnapshot.mu.RLock() + m3, ok3 := globalSnapshot.metrics[key3] + globalSnapshot.mu.RUnlock() + if ok3 { + return snapshotChannelMetrics(m3) + } + + // 无数据，返回中性默认值 + return &ChannelMetrics{ + SuccessRate: 1.0, + AvgLatency: 500 * time.Millisecond, + } +} + +func snapshotChannelMetrics(m *ChannelMetrics) *ChannelMetrics { + if m == nil { + return nil + } + m.mu.Lock() + defer m.mu.Unlock() + return &ChannelMetrics{ + SuccessRate: m.SuccessRate, + ErrorRate: m.ErrorRate, + RateLimitRate: m.RateLimitRate, + Status5xxRate: m.Status5xxRate, + AvgLatency: m.AvgLatency, + SampleCount: m.SampleCount, + LastSeen: m.LastSeen, + } +} + + +// CurrentConcurrencyTracker 本地并发计数器（原子操作，零网络开销） +type CurrentConcurrencyTracker struct { + counters sync.Map // map[int]*atomic.Int64 +} + +var globalConcurrency = &CurrentConcurrencyTracker{} + +func (t *CurrentConcurrencyTracker) Inc(channelID int) int64 { + v, _ := t.counters.LoadOrStore(channelID, new(atomic.Int64)) + return v.(*atomic.Int64).Add(1) +} + +func (t *CurrentConcurrencyTracker) Dec(channelID int) int64 { + v, ok := t.counters.Load(channelID) + if !ok { + return 0 + } + return v.(*atomic.Int64).Add(-1) +} + +func (t *CurrentConcurrencyTracker) Get(channelID int) int64 { + v, ok := t.counters.Load(channelID) + if !ok { + return 0 + } + return v.(*atomic.Int64).Load() +} + +// IncChannelConcurrency 增加并发计数 +func IncChannelConcurrency(channelID int) int64 { + return globalConcurrency.Inc(channelID) +} + +// DecChannelConcurrency 减少并发计数 +func DecChannelConcurrency(channelID int) int64 { + return globalConcurrency.Dec(channelID) +} + +// GetChannelConcurrency 获取当前并发 +func GetChannelConcurrency(channelID int) int64 { + return globalConcurrency.Get(channelID) +} + + +// SyncAdaptiveMetricsToRedis publishes a compact snapshot for multi-instance +// sticky-or-shared observation. Best-effort; failures are silent. +func SyncAdaptiveMetricsToRedis() { + if !common.RedisEnabled || common.RDB == nil { + return + } + globalSnapshot.mu.RLock() + defer globalSnapshot.mu.RUnlock() + type row struct { + ChannelID int `json:"c"` + Group string `json:"g"` + Model string `json:"m"` + SuccessRate float64 `json:"s"` + SampleCount int64 `json:"n"` + } + out := make([]row, 0, len(globalSnapshot.metrics)) + for k, m := range globalSnapshot.metrics { + m.mu.Lock() + out = append(out, row{k.ChannelID, k.Group, k.Model, m.SuccessRate, m.SampleCount}) + m.mu.Unlock() + } + b, err := common.Marshal(out) + if err != nil { + return + } + key := fmt.Sprintf("newapi:adaptive:metrics:%s", common.NodeName) + if key == "newapi:adaptive:metrics:" { + key = "newapi:adaptive:metrics:default" + } + _ = common.RedisSet(key, string(b), 2*time.Minute) +} diff --git a/service/channel_score.go b/service/channel_score.go new file mode 100644 index 000000000000..6fc034d7fe14 --- /dev/null +++ b/service/channel_score.go @@ -0,0 +1,166 @@ +package service + +import ( + "github.com/QuantumNous/new-api/constant" + "math" + "math/rand" + "sort" + "time" + + "github.com/QuantumNous/new-api/model" +) + +// CandidateScore 评分候选结果 +type CandidateScore struct { + Channel *model.Channel + Score float64 + + // 各因子明细（供日志/观测用） + BaseWeight float64 `json:"base_weight"` + SuccessFactor float64 `json:"success_factor"` + LatencyFactor float64 `json:"latency_factor"` + RateLimitFactor float64 `json:"rate_limit_factor"` + ConcurrencyFactor float64 `json:"concurrency_factor"` + CircuitFactor float64 `json:"circuit_factor"` + AffinityFactor float64 `json:"affinity_factor"` +} + +// ScoreCandidates 对一组渠道进行动态评分，返回排序后的候选列表 +// 不传 group, model 时会回退到 channel 级别指标 +func ScoreCandidates(channels []*model.Channel, group, model string, preferredChannelID int) []CandidateScore { + if len(channels) == 0 { + return nil + } + + candidates := make([]CandidateScore, 0, len(channels)) + + for _, ch := range channels { + circuitFactor := 1.0 + if constant.ChannelCircuitBreakerEnabled { + // One state read is enough for scoring; open -> 0, half-open -> 0.5. + state, _, _ := GetCircuitState(ch.Id) + switch state { + case CircuitOpen: + circuitFactor = 0.0 + case CircuitHalfOpen: + circuitFactor = 0.5 + } + } + + candidates = append(candidates, scoreCandidate(ch, group, model, preferredChannelID, circuitFactor)) + } + + // 按分数降序排序 + sort.Slice(candidates, func(i, j int) bool { + return candidates[i].Score > candidates[j].Score + }) + + return candidates +} + +func scoreCandidate(ch *model.Channel, group, model string, preferredChannelID int, circuitFactor float64) CandidateScore { + metrics := GetMetrics(ch.Id, group, model) + baseWeight := float64(ch.GetWeight()) + if baseWeight <= 0 { + baseWeight = 1.0 + } + successFactor := math.Pow(metrics.SuccessRate, 2) + latencyMs := float64(metrics.AvgLatency) / float64(time.Millisecond) + latencyFactor := latencyToScore(latencyMs) + rateLimitFactor := math.Max(0, 1.0-metrics.RateLimitRate) + + currentConcurrency := GetChannelConcurrency(ch.Id) + maxConcurrency := int64(constant.MaxChannelConcurrency) + concurrencyFactor := 1.0 + if maxConcurrency > 0 && currentConcurrency >= maxConcurrency { + concurrencyFactor = 0.1 + } else if maxConcurrency > 0 { + concurrencyFactor = 1.0 - float64(currentConcurrency)/float64(maxConcurrency)*0.5 + } + + affinityFactor := 1.0 + if preferredChannelID > 0 && ch.Id == preferredChannelID { + affinityFactor = 1.5 + } + score := baseWeight * successFactor * latencyFactor * rateLimitFactor * + concurrencyFactor * circuitFactor * affinityFactor + score *= 0.95 + rand.Float64()*0.1 + + return CandidateScore{ + Channel: ch, + Score: score, + BaseWeight: baseWeight, + SuccessFactor: successFactor, + LatencyFactor: latencyFactor, + RateLimitFactor: rateLimitFactor, + ConcurrencyFactor: concurrencyFactor, + CircuitFactor: circuitFactor, + AffinityFactor: affinityFactor, + } +} + +// SelectTopKWeighted 从候选列表中取 topK 然后按 score 加权随机选一个 +func SelectTopKWeighted(candidates []CandidateScore, k int) *CandidateScore { + if len(candidates) == 0 { + return nil + } + + if len(candidates) == 1 { + return &candidates[0] + } + + // 取 topK + if k <= 0 { + k = 3 + } + if k > len(candidates) { + k = len(candidates) + } + top := candidates[:k] + + // 加权随机 + var totalWeight float64 + for _, c := range top { + if c.Score > 0 { + totalWeight += c.Score + } + } + + if totalWeight <= 0 { + // 所有分数为 0，均匀随机 + idx := rand.Intn(len(top)) + return &top[idx] + } + + r := rand.Float64() * totalWeight + var cumulative float64 + for i, c := range top { + cumulative += c.Score + if r < cumulative { + return &top[i] + } + } + + return &top[len(top)-1] +} + +// latencyToScore 将延迟（毫秒）映射到 [0, 1] 分数 +// 500ms → 1.0, 1s → 0.8, 2s → 0.5, 5s → 0.2, 10s+ → 0.05 +func latencyToScore(ms float64) float64 { + switch { + case ms <= 0: + return 1.0 + case ms <= 500: + return 1.0 + case ms <= 1000: + return 0.8 + case ms <= 2000: + return 0.5 + case ms <= 5000: + return 0.2 + case ms <= 10000: + return 0.1 + default: + return 0.05 + } +} diff --git a/service/channel_select.go b/service/channel_select.go index 24c4e252bfb3..4de77641e314 100644 --- a/service/channel_select.go +++ b/service/channel_select.go @@ -82,6 +82,18 @@ func (p *RetryParam) ResetRetryNextTry() { // Retry=3: GroupB, priority1 (startRetryIndex=2, priorityRetry=1) // 分组B, 优先级1 func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, error) { + // Adaptive entry. AdaptiveSelectChannel must only call + // cacheGetRandomSatisfiedChannelLegacy — never this function — or flags + // cause infinite recursion / stack overflow. + if constant.AdaptiveBalanceEnabled || constant.AdaptiveBalanceShadowMode { + return AdaptiveSelectChannel(param) + } + return cacheGetRandomSatisfiedChannelLegacy(param) +} + +// cacheGetRandomSatisfiedChannelLegacy is the original random / auto-group picker. +// Safe to call from adaptive fallbacks and candidate collection. +func cacheGetRandomSatisfiedChannelLegacy(param *RetryParam) (*model.Channel, string, error) { var channel *model.Channel var err error selectGroup := param.TokenGroup @@ -116,7 +128,7 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, } logger.LogDebug(param.Ctx, "Auto selecting group: %s, priorityRetry: %d", autoGroup, priorityRetry) - channel, _ = model.GetRandomSatisfiedChannel(autoGroup, param.ModelName, priorityRetry, param.RequestPath) + channel, _ = model.GetRandomSatisfiedChannelExcluding(autoGroup, param.ModelName, priorityRetry, param.RequestPath, adaptiveUsedChannelSet(param.Ctx)) if channel == nil { // Current group has no available channel for this model, try next group // 当前分组没有该模型的可用渠道，尝试下一个分组 @@ -154,7 +166,7 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, break } } else { - channel, err = model.GetRandomSatisfiedChannel(param.TokenGroup, param.ModelName, param.GetRetry(), param.RequestPath) + channel, err = model.GetRandomSatisfiedChannelExcluding(param.TokenGroup, param.ModelName, param.GetRetry(), param.RequestPath, adaptiveUsedChannelSet(param.Ctx)) if err != nil { return nil, param.TokenGroup, err } From 063b4815565d65f1d2c05940578083ae6f959559 Mon Sep 17 00:00:00 2001 From: yuanjia Date: Sun, 19 Jul 2026 11:10:11 +0800 Subject: [PATCH 2/2] =?UTF-8?q?fix(channel):=20adaptive=20P0=20=E2=80=94?= =?UTF-8?q?=20shadow=20observability,=20permits,=20candidates?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CodeRabbit #6301 P0: - allow shadow-only path to score and emit [shadow] LogInfo (5% sample) - treat non-429 4xx as healthy for metrics (align with circuit) - release half-open permit on relay helper panic; bump generation on probe timeout - load full DB candidate set when memory cache is off - compute retry priority tiers before exclusions - honor IsSkipRetryError before generic channel-error retry - validate EWMA/cooldown/concurrency env knobs --- common/init.go | 13 +++++ controller/relay.go | 12 +++-- model/channel_cache.go | 96 ++++++++++++++++++++++++++----------- service/channel_adaptive.go | 20 +++++--- service/channel_circuit.go | 3 ++ 5 files changed, 106 insertions(+), 38 deletions(-) diff --git a/common/init.go b/common/init.go index c7ab141d2104..827aa6329caf 100644 --- a/common/init.go +++ b/common/init.go @@ -1,6 +1,7 @@ package common import ( + "math" "flag" "fmt" "log" @@ -117,6 +118,18 @@ func InitEnv() { constant.ChannelCooldownSeconds = GetEnvOrDefault("CHANNEL_COOLDOWN_SECONDS", 30) constant.EwmaAlpha = GetEnvOrDefaultFloat("EWMA_ALPHA", 0.1) constant.MaxChannelConcurrency = GetEnvOrDefault("MAX_CHANNEL_CONCURRENCY", 10) + if constant.ChannelCooldownSeconds <= 0 { + constant.ChannelCooldownSeconds = 30 + } + if math.IsNaN(constant.EwmaAlpha) || constant.EwmaAlpha <= 0 || constant.EwmaAlpha > 1 { + constant.EwmaAlpha = 0.1 + } + if constant.MaxChannelConcurrency < 0 { + constant.MaxChannelConcurrency = 10 + } + if constant.MaxRetryChannels < 0 { + constant.MaxRetryChannels = 0 + } // Initialize string variables with GetEnvOrDefaultString GeminiSafetySetting = GetEnvOrDefaultString("GEMINI_SAFETY_SETTING", "BLOCK_NONE") diff --git a/controller/relay.go b/controller/relay.go index 2e01d339dad1..f329daca720f 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -217,6 +217,12 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { // Always dec even if helper panics (CustomRecovery still runs after). func() { defer service.DecChannelConcurrency(channel.Id) + defer func() { + if r := recover(); r != nil { + service.ReleaseAdaptiveCircuitPermit(c, channel.Id) + panic(r) + } + }() switch relayFormat { case types.RelayFormatOpenAIRealtime: newAPIError = relay.WssHelper(c, relayInfo) @@ -390,15 +396,15 @@ func shouldRetry(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) b if openaiErr.GetErrorCode() == types.ErrorCodeGetChannelFailed { return false } + if types.IsSkipRetryError(openaiErr) { + return false + } if isUpstreamChannelQuotaError(openaiErr) { return true } if types.IsChannelError(openaiErr) { return true } - if types.IsSkipRetryError(openaiErr) { - return false - } code := openaiErr.StatusCode if code >= 200 && code < 300 { return false diff --git a/model/channel_cache.go b/model/channel_cache.go index 9dcb095fa9eb..339790744d00 100644 --- a/model/channel_cache.go +++ b/model/channel_cache.go @@ -116,14 +116,8 @@ func SyncChannelCache(frequency int) { // When memory cache is off, falls back to a single DB-selected channel. func GetSatisfiedChannels(group string, modelName string, requestPath string) ([]*Channel, error) { if !common.MemoryCacheEnabled { - ch, err := GetChannel(group, modelName, 0, requestPath) - if err != nil { - return nil, err - } - if ch == nil { - return nil, nil - } - return []*Channel{ch}, nil + // Adaptive scoring needs the full candidate set, not one random row. + return getSatisfiedChannelsFromDB(group, modelName, requestPath) } channelSyncLock.RLock() @@ -182,26 +176,10 @@ func GetRandomSatisfiedChannelExcluding(group string, model string, retry int, r if len(channels) == 0 { return nil, nil } - if len(excluded) > 0 { - filtered := make([]int, 0, len(channels)) - for _, channelID := range channels { - if _, skip := excluded[channelID]; !skip { - filtered = append(filtered, channelID) - } - } - channels = filtered - if len(channels) == 0 { - return nil, nil - } - } - - if len(channels) == 1 { - if channel, ok := channelsIDM[channels[0]]; ok { - return channel, nil - } - return nil, fmt.Errorf("数据库一致性错误，渠道# %d 不存在，请联系管理员修复", channels[0]) - } + // Compute priority tiers from the FULL set first, then apply exclusions within + // the chosen tier. Filtering first reindexes tiers (e.g. 100/50/0 with 100 + // excluded would make retry=1 pick 0 instead of 50). uniquePriorities := make(map[int]bool) for _, channelId := range channels { if channel, ok := channelsIDM[channelId]; ok { @@ -221,10 +199,15 @@ func GetRandomSatisfiedChannelExcluding(group string, model string, retry int, r } targetPriority := int64(sortedUniquePriorities[retry]) - // get the priority for the given retry number + // get the priority for the given retry number, then exclude within the tier var sumWeight = 0 var targetChannels []*Channel for _, channelId := range channels { + if len(excluded) > 0 { + if _, skip := excluded[channelId]; skip { + continue + } + } if channel, ok := channelsIDM[channelId]; ok { if channel.GetPriority() == targetPriority { sumWeight += channel.GetWeight() @@ -389,3 +372,60 @@ func CacheUpdateChannel(channel *Channel) { channelSyncLock.Unlock() InvalidatePricingCache() } + + +// getSatisfiedChannelsFromDB loads every enabled channel for group+model when the +// in-memory channel cache is off. Adaptive balance needs the full set so circuit +// filtering and top-K scoring remain meaningful. +func getSatisfiedChannelsFromDB(group string, modelName string, requestPath string) ([]*Channel, error) { + var abilities []Ability + err := DB.Where(commonGroupCol+" = ? and model = ? and enabled = ?", group, modelName, true). + Order("priority DESC").Find(&abilities).Error + if err != nil { + return nil, err + } + if len(abilities) == 0 { + normalizedModel := ratio_setting.FormatMatchingModelName(modelName) + if normalizedModel != modelName { + err = DB.Where(commonGroupCol+" = ? and model = ? and enabled = ?", group, normalizedModel, true). + Order("priority DESC").Find(&abilities).Error + if err != nil { + return nil, err + } + modelName = normalizedModel + } + } + if len(abilities) == 0 { + return nil, nil + } + abilities = filterAbilitiesByRequestPathAndModel(abilities, requestPath, modelName) + if len(abilities) == 0 { + return nil, nil + } + + seen := make(map[int]struct{}, len(abilities)) + ids := make([]int, 0, len(abilities)) + for _, ability := range abilities { + if _, ok := seen[ability.ChannelId]; ok { + continue + } + seen[ability.ChannelId] = struct{}{} + ids = append(ids, ability.ChannelId) + } + if len(ids) == 0 { + return nil, nil + } + + channels, err := GetChannelsByIds(ids) + if err != nil { + return nil, err + } + out := make([]*Channel, 0, len(channels)) + for _, ch := range channels { + if ch == nil || ch.Status != common.ChannelStatusEnabled { + continue + } + out = append(out, ch) + } + return out, nil +} diff --git a/service/channel_adaptive.go b/service/channel_adaptive.go index fe00ed1cfbf2..430dd93368d6 100644 --- a/service/channel_adaptive.go +++ b/service/channel_adaptive.go @@ -12,7 +12,7 @@ import ( "github.com/gin-gonic/gin" ) -var adaptiveLogSample = 0.01 // shadow-mode log sample rate +var adaptiveLogSample = 0.05 // shadow-mode log sample rate (Info-level) // 请求上下文 key type adaptiveContextKey string @@ -31,8 +31,9 @@ const ( func AdaptiveSelectChannel(param *RetryParam) (*model.Channel, string, error) { ctx := param.Ctx - // 未开启完整自适应：仅 legacy（含「只开 shadow」旧行为，避免递归） - if !constant.AdaptiveBalanceEnabled { + // Neither adaptive nor shadow: pure legacy. Shadow-only still runs scoring so + // comparison logs are produced without changing the routed channel. + if !constant.AdaptiveBalanceEnabled && !constant.AdaptiveBalanceShadowMode { return cacheGetRandomSatisfiedChannelLegacy(param) } @@ -271,21 +272,26 @@ func logAdaptiveCompare(c *gin.Context, modelName, group string, selected *Candi if oldCh != nil { oldID = oldCh.Id } - logger.LogDebug(c, "[shadow] model=%s group=%s adaptive=#%d(%.3f) orig=#%d", - modelName, group, selected.Channel.Id, selected.Score, oldID) + // Info (not Debug): production shadow observation must be visible without + // turning on full debug logging. Sampling still limits volume. + logger.LogInfo(c, fmt.Sprintf("[shadow] model=%s group=%s adaptive=#%d(%.3f) orig=#%d", + modelName, group, selected.Channel.Id, selected.Score, oldID)) } // RecordAdaptiveResult 请求完成后回调：更新指标 + 熔断状态 func RecordAdaptiveResult(c *gin.Context, channelID int, group, modelName string, statusCode int, latency time.Duration, err error) { - if !constant.AdaptiveBalanceEnabled { + if !constant.AdaptiveBalanceEnabled && !constant.AdaptiveBalanceShadowMode { return } if channelID <= 0 { return } + // Client/input 4xx (except 429) proves the upstream is reachable; do not + // demote healthy channels for bad requests. Align metrics with circuit. succeeded := err == nil && statusCode < 400 - if succeeded { + clientSide := statusCode >= 400 && statusCode < 500 && statusCode != 429 + if succeeded || clientSide { ObserveSuccess(channelID, group, modelName, latency) } else { ObserveFailure(channelID, group, modelName, statusCode, latency) diff --git a/service/channel_circuit.go b/service/channel_circuit.go index 4058932a912d..f29ada797ad7 100644 --- a/service/channel_circuit.go +++ b/service/channel_circuit.go @@ -205,8 +205,11 @@ func AcquireCircuitPermit(channelID int) (CircuitPermit, bool) { const halfOpenProbeTimeout = 60 * time.Second if cb.HalfOpenInFlight > 0 && !cb.HalfOpenSince.IsZero() && time.Since(cb.HalfOpenSince) > halfOpenProbeTimeout { + // Drop the stuck probe AND bump generation so a late response from the + // expired permit cannot close the circuit or free a newer probe's slot. cb.HalfOpenInFlight = 0 cb.HalfOpenSince = time.Time{} + cb.Generation++ } if cb.HalfOpenInFlight >= cb.HalfOpenLimit { return CircuitPermit{}, false