From 378256a53882a94c46707c579ae031d8bf18b890 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BC=A0=E5=90=8D=E9=94=90?= <1565842059@qq.com> Date: Fri, 28 Aug 2026 19:14:48 +0800 Subject: [PATCH] feat(relay): add in-memory adaptive circuit breaker and auto-fallback retry mechanism --- controller/relay.go | 10 ++ go.mod | 2 +- model/channel_cache.go | 19 ++++ service/channel_fallback_test.go | 23 ++++ service/channel_select.go | 62 +++++++---- service/circuit_breaker.go | 184 +++++++++++++++++++++++++++++++ service/circuit_breaker_test.go | 82 ++++++++++++++ 7 files changed, 358 insertions(+), 24 deletions(-) create mode 100644 service/channel_fallback_test.go create mode 100644 service/circuit_breaker.go create mode 100644 service/circuit_breaker_test.go diff --git a/controller/relay.go b/controller/relay.go index 8dccfe76..3129d024 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -230,12 +230,22 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { if newAPIError == nil { relayInfo.LastError = nil + service.GlobalCircuitBreaker.RecordSuccess(channel.Id) + useChannel := c.GetStringSlice("use_channel") + if len(useChannel) > 1 { + c.Header("X-New-API-Fallback-Count", fmt.Sprintf("%d", len(useChannel)-1)) + logger.LogInfo(c, fmt.Sprintf("[Auto-Fallback] 自动无感降级重试成功: 路径 %s", strings.Trim(strings.Join(strings.Fields(fmt.Sprint(useChannel)), " -> "), "[]"))) + } return } newAPIError = service.NormalizeViolationFeeError(newAPIError) relayInfo.LastError = newAPIError + // 记录失败至自适应熔断器，并将故障渠道加入本次请求的排除列表以实现无感自动降级 + service.GlobalCircuitBreaker.RecordFailure(channel.Id, newAPIError.StatusCode) + retryParam.AddExcludedChannel(channel.Id) + processChannelError(c, *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()), newAPIError) if !shouldRetry(c, newAPIError, common.RetryTimes-retryParam.GetRetry()) { diff --git a/go.mod b/go.mod index b0642f16..a0609149 100644 --- a/go.mod +++ b/go.mod @@ -90,7 +90,7 @@ require ( require ( github.com/DmitriyVTitov/size v1.5.0 // indirect github.com/anknown/darts v0.0.0-20151216065714-83ff685239e6 // indirect - github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.8 // indirect + github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.8 github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.21 // indirect github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.21 // indirect github.com/beorn7/perks v1.0.1 // indirect diff --git a/model/channel_cache.go b/model/channel_cache.go index 86c59438..0dc9aa3c 100644 --- a/model/channel_cache.go +++ b/model/channel_cache.go @@ -112,6 +112,10 @@ func SyncChannelCache(frequency int) { } func GetRandomSatisfiedChannel(group string, model string, retry int, requestPath string) (*Channel, error) { + return GetRandomSatisfiedChannelWithExcluded(group, model, retry, requestPath, nil) +} + +func GetRandomSatisfiedChannelWithExcluded(group string, model string, retry int, requestPath string, excludedIDs []int) (*Channel, error) { // if memory cache is disabled, get channel directly from database if !common.MemoryCacheEnabled { return GetChannel(group, model, retry, requestPath) @@ -129,6 +133,21 @@ func GetRandomSatisfiedChannel(group string, model string, retry int, requestPat channels = filterChannelsByRequestPathAndModel(group2model2channels[group][normalizedModel], requestPath, model) } + // Filter out excluded channels (channels that already failed in the current request) + if len(excludedIDs) > 0 && len(channels) > 0 { + excludedMap := make(map[int]bool, len(excludedIDs)) + for _, id := range excludedIDs { + excludedMap[id] = true + } + filtered := make([]int, 0, len(channels)) + for _, id := range channels { + if !excludedMap[id] { + filtered = append(filtered, id) + } + } + channels = filtered + } + if len(channels) == 0 { return nil, nil } diff --git a/service/channel_fallback_test.go b/service/channel_fallback_test.go new file mode 100644 index 00000000..80ac1f0e --- /dev/null +++ b/service/channel_fallback_test.go @@ -0,0 +1,23 @@ +package service + +import ( + "testing" +) + +func TestRetryParam_AddExcludedChannel(t *testing.T) { + param := &RetryParam{} + + param.AddExcludedChannel(10) + param.AddExcludedChannel(20) + param.AddExcludedChannel(10) // duplicate, should be ignored + param.AddExcludedChannel(0) // invalid, should be ignored + param.AddExcludedChannel(-1) // invalid, should be ignored + + if len(param.ExcludedChannelIDs) != 2 { + t.Fatalf("expected 2 excluded channels, got %d (%v)", len(param.ExcludedChannelIDs), param.ExcludedChannelIDs) + } + + if param.ExcludedChannelIDs[0] != 10 || param.ExcludedChannelIDs[1] != 20 { + t.Fatalf("expected [10, 20], got %v", param.ExcludedChannelIDs) + } +} diff --git a/service/channel_select.go b/service/channel_select.go index 0ab88dc8..a8546554 100644 --- a/service/channel_select.go +++ b/service/channel_select.go @@ -2,6 +2,7 @@ package service import ( "errors" + "fmt" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" @@ -11,12 +12,25 @@ import ( ) type RetryParam struct { - Ctx *gin.Context - TokenGroup string - ModelName string - RequestPath string - Retry *int - resetNextTry bool + Ctx *gin.Context + TokenGroup string + ModelName string + RequestPath string + Retry *int + ExcludedChannelIDs []int + resetNextTry bool +} + +func (p *RetryParam) AddExcludedChannel(channelID int) { + if channelID <= 0 { + return + } + for _, id := range p.ExcludedChannelIDs { + if id == channelID { + return + } + } + p.ExcludedChannelIDs = append(p.ExcludedChannelIDs, channelID) } func (p *RetryParam) GetRetry() int { @@ -45,6 +59,23 @@ func (p *RetryParam) ResetRetryNextTry() { p.resetNextTry = true } +func selectChannelWithCircuitBreaker(group string, modelName string, retry int, requestPath string, excludedIDs []int) (*model.Channel, error) { + channel, err := model.GetRandomSatisfiedChannelWithExcluded(group, modelName, retry, requestPath, excludedIDs) + if err != nil || channel == nil { + return channel, err + } + // If the chosen channel is currently in circuit breaker cooldown, attempt to pick an alternate available channel + if !GlobalCircuitBreaker.IsAvailable(channel.Id) { + combinedExcluded := append([]int{channel.Id}, excludedIDs...) + altChannel, altErr := model.GetRandomSatisfiedChannelWithExcluded(group, modelName, retry, requestPath, combinedExcluded) + if altChannel != nil && altErr == nil { + logger.LogInfo(nil, fmt.Sprintf("[Auto-Fallback] 渠道 #%d 处于熔断冷却中，自动切换至健康备用渠道 #%d", channel.Id, altChannel.Id)) + return altChannel, nil + } + } + return channel, nil +} + // CacheGetRandomSatisfiedChannel tries to get a random channel that satisfies the requirements. // 尝试获取一个满足要求的随机渠道。 // @@ -65,21 +96,6 @@ func (p *RetryParam) ResetRetryNextTry() { // // - When GetRandomSatisfiedChannel returns nil (priorities exhausted), moves to next group. // 当 GetRandomSatisfiedChannel 返回 nil（优先级用完）时，切换到下一个分组。 -// -// Example flow (2 groups, each with 2 priorities, RetryTimes=3): -// 示例流程（2个分组，每个有2个优先级，RetryTimes=3）： -// -// Retry=0: GroupA, priority0 (startRetryIndex=0, priorityRetry=0) -// 分组A, 优先级0 -// -// Retry=1: GroupA, priority1 (startRetryIndex=0, priorityRetry=1) -// 分组A, 优先级1 -// -// Retry=2: GroupA exhausted → GroupB, priority0 (startRetryIndex=2, priorityRetry=0) -// 分组A用完 → 分组B, 优先级0 -// -// Retry=3: GroupB, priority1 (startRetryIndex=2, priorityRetry=1) -// 分组B, 优先级1 func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, error) { var channel *model.Channel var err error @@ -115,7 +131,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, _ = selectChannelWithCircuitBreaker(autoGroup, param.ModelName, priorityRetry, param.RequestPath, param.ExcludedChannelIDs) if channel == nil { // Current group has no available channel for this model, try next group // 当前分组没有该模型的可用渠道，尝试下一个分组 @@ -153,7 +169,7 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, break } } else { - channel, err = model.GetRandomSatisfiedChannel(param.TokenGroup, param.ModelName, param.GetRetry(), param.RequestPath) + channel, err = selectChannelWithCircuitBreaker(param.TokenGroup, param.ModelName, param.GetRetry(), param.RequestPath, param.ExcludedChannelIDs) if err != nil { return nil, param.TokenGroup, err } diff --git a/service/circuit_breaker.go b/service/circuit_breaker.go new file mode 100644 index 00000000..95d7e6bf --- /dev/null +++ b/service/circuit_breaker.go @@ -0,0 +1,184 @@ +package service + +import ( + "fmt" + "net/http" + "sync" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/setting/operation_setting" +) + +// CircuitBreaker manages in-memory temporary failure cooldowns for channels. +// 自适应内存熔断器：在渠道出现持续故障（如 429、5xx、超时）时提供临时冷却隔离，防止流量持续撞墙并支持超时自动半开探测。 +type CircuitBreaker struct { + mu sync.RWMutex + breakers map[int]*channelBreaker + failureThreshold int // 触发熔断的连续失败阈值 (默认 3 次) + baseCooldown time.Duration // 基础冷却时长 (默认 30 秒) + maxCooldown time.Duration // 最大冷却时长 (默认 5 分钟) +} + +type channelBreaker struct { + consecutiveFailures int + lastFailureTime time.Time + cooldownDuration time.Duration + cooldownUntil time.Time + probing bool // 是否处于 Half-Open 探测中 +} + +var ( + GlobalCircuitBreaker = NewCircuitBreaker(3, 30*time.Second, 5*time.Minute) +) + +func NewCircuitBreaker(threshold int, baseCooldown time.Duration, maxCooldown time.Duration) *CircuitBreaker { + return &CircuitBreaker{ + breakers: make(map[int]*channelBreaker), + failureThreshold: threshold, + baseCooldown: baseCooldown, + maxCooldown: maxCooldown, + } +} + +// RecordSuccess records a successful request for a channel and resets its failure state. +// 记录请求成功：重置该渠道的连续失败计数与熔断状态。 +func (cb *CircuitBreaker) RecordSuccess(channelID int) { + if channelID <= 0 { + return + } + cb.mu.Lock() + defer cb.mu.Unlock() + + b, exists := cb.breakers[channelID] + if !exists { + return + } + if b.consecutiveFailures > 0 || !b.cooldownUntil.IsZero() { + common.SysLog(fmt.Sprintf("[CircuitBreaker] 渠道 #%d 恢复健康，重置熔断计数", channelID)) + } + delete(cb.breakers, channelID) +} + +// RecordFailure records a failure for a channel and triggers cooldown if threshold is met. +// 记录请求失败：根据错误状态码累计失败次数或直接触发临时熔断。 +func (cb *CircuitBreaker) RecordFailure(channelID int, statusCode int) { + if channelID <= 0 { + return + } + + // 仅对可重试或上游服务异常进行熔断统计 (429, 500, 502, 503, 504 或 网络超时) + if !isSevereOrRetryableStatus(statusCode) { + return + } + + cb.mu.Lock() + defer cb.mu.Unlock() + + b, exists := cb.breakers[channelID] + now := time.Now() + + if !exists { + b = &channelBreaker{ + cooldownDuration: cb.baseCooldown, + } + cb.breakers[channelID] = b + } + + b.consecutiveFailures++ + b.lastFailureTime = now + b.probing = false + + // 遇到 429 (Too Many Requests) 或 503 (Service Unavailable) 或 达到失败阈值，立即触发熔断 + if statusCode == http.StatusTooManyRequests || statusCode == http.StatusServiceUnavailable || b.consecutiveFailures >= cb.failureThreshold { + if b.cooldownDuration == 0 { + b.cooldownDuration = cb.baseCooldown + } else if b.consecutiveFailures > cb.failureThreshold { + // 指数退避增长冷却时间，最高不超过 maxCooldown + b.cooldownDuration *= 2 + if b.cooldownDuration > cb.maxCooldown { + b.cooldownDuration = cb.maxCooldown + } + } + b.cooldownUntil = now.Add(b.cooldownDuration) + common.SysLog(fmt.Sprintf("[CircuitBreaker] 渠道 #%d 触发临时熔断冷却，连续失败次数: %d, 状态码: %d, 冷却至: %s (%v)", + channelID, b.consecutiveFailures, statusCode, b.cooldownUntil.Format("15:04:05"), b.cooldownDuration)) + } +} + +// IsAvailable checks if the channel is currently available (not in cooldown or half-open ready for probe). +// 检查渠道当前是否可用（未熔断或已到冷却期可进行半开探测）。 +func (cb *CircuitBreaker) IsAvailable(channelID int) bool { + if channelID <= 0 { + return true + } + cb.mu.Lock() + defer cb.mu.Unlock() + + b, exists := cb.breakers[channelID] + if !exists { + return true + } + + now := time.Now() + // 如果仍在冷却期内 + if now.Before(b.cooldownUntil) { + return false + } + + // 冷却期已过，进入 Half-Open 状态，放行单次探测请求 + if !b.cooldownUntil.IsZero() { + b.probing = true + b.cooldownUntil = time.Time{} // 清空冷却期 + common.SysLog(fmt.Sprintf("[CircuitBreaker] 渠道 #%d 冷却期已过，进入半开 (Half-Open) 状态开始探测恢复", channelID)) + } + return true +} + +// FilterAvailableChannels filters out channels currently in cooldown. +// If ALL candidate channels are in cooldown, it gracefully returns all of them to prevent total starvation. +// 过滤候选渠道列表：优先剔除处于熔断冷却期的渠道；如果所有渠道均处于冷却中，则优雅降级返回全量渠道。 +func (cb *CircuitBreaker) FilterAvailableChannels(channelIDs []int) []int { + if len(channelIDs) <= 1 { + return channelIDs + } + + cb.mu.RLock() + now := time.Now() + available := make([]int, 0, len(channelIDs)) + + for _, id := range channelIDs { + b, exists := cb.breakers[id] + if !exists || now.After(b.cooldownUntil) { + available = append(available, id) + } + } + cb.mu.RUnlock() + + // 如果全部都在熔断冷却中，退化返回所有渠道（避免完全无渠道可用） + if len(available) == 0 { + return channelIDs + } + return available +} + +// Reset clears all circuit breaker records. +func (cb *CircuitBreaker) Reset() { + cb.mu.Lock() + defer cb.mu.Unlock() + cb.breakers = make(map[int]*channelBreaker) +} + +func isSevereOrRetryableStatus(statusCode int) bool { + if statusCode == 0 || statusCode < 100 || statusCode > 599 { + return true + } + if statusCode == http.StatusTooManyRequests || // 429 + statusCode == http.StatusBadGateway || // 502 + statusCode == http.StatusServiceUnavailable || // 503 + statusCode == http.StatusGatewayTimeout || // 504 + statusCode == http.StatusInternalServerError { // 500 + return true + } + return operation_setting.ShouldRetryByStatusCode(statusCode) +} diff --git a/service/circuit_breaker_test.go b/service/circuit_breaker_test.go new file mode 100644 index 00000000..290d6bb1 --- /dev/null +++ b/service/circuit_breaker_test.go @@ -0,0 +1,82 @@ +package service + +import ( + "net/http" + "testing" + "time" +) + +func TestCircuitBreaker_BasicFlow(t *testing.T) { + cb := NewCircuitBreaker(3, 100*time.Millisecond, 1*time.Second) + + channelID := 101 + + // Initial state: available + if !cb.IsAvailable(channelID) { + t.Fatalf("channel %d should be available initially", channelID) + } + + // 1st failure (500) + cb.RecordFailure(channelID, http.StatusInternalServerError) + if !cb.IsAvailable(channelID) { + t.Fatalf("channel %d should still be available after 1 failure (threshold=3)", channelID) + } + + // 2nd failure (500) + cb.RecordFailure(channelID, http.StatusInternalServerError) + if !cb.IsAvailable(channelID) { + t.Fatalf("channel %d should still be available after 2 failures", channelID) + } + + // 3rd failure (500) -> should trigger cooldown + cb.RecordFailure(channelID, http.StatusInternalServerError) + if cb.IsAvailable(channelID) { + t.Fatalf("channel %d should be in cooldown after 3 consecutive failures", channelID) + } + + // Wait for cooldown to expire (100ms) + time.Sleep(120 * time.Millisecond) + + // Now should be available (entering Half-Open probing) + if !cb.IsAvailable(channelID) { + t.Fatalf("channel %d should be available for probing after cooldown expires", channelID) + } + + // Success probe -> recovers + cb.RecordSuccess(channelID) + if !cb.IsAvailable(channelID) { + t.Fatalf("channel %d should be healthy after successful probe", channelID) + } +} + +func TestCircuitBreaker_Immediate429Cooldown(t *testing.T) { + cb := NewCircuitBreaker(3, 100*time.Millisecond, 1*time.Second) + channelID := 202 + + // 429 Too Many Requests should immediately trigger cooldown even with 1 failure + cb.RecordFailure(channelID, http.StatusTooManyRequests) + if cb.IsAvailable(channelID) { + t.Fatalf("channel %d should immediately enter cooldown upon 429", channelID) + } +} + +func TestCircuitBreaker_FilterAvailableChannels(t *testing.T) { + cb := NewCircuitBreaker(3, 200*time.Millisecond, 1*time.Second) + + c1, c2, c3 := 1, 2, 3 + cb.RecordFailure(c1, http.StatusTooManyRequests) // c1 in cooldown + + available := cb.FilterAvailableChannels([]int{c1, c2, c3}) + if len(available) != 2 || available[0] != c2 || available[1] != c3 { + t.Fatalf("expected [2, 3], got %v", available) + } + + // If all channels are in cooldown, graceful fallback returns all + cb.RecordFailure(c2, http.StatusTooManyRequests) + cb.RecordFailure(c3, http.StatusTooManyRequests) + + fallbackAll := cb.FilterAvailableChannels([]int{c1, c2, c3}) + if len(fallbackAll) != 3 { + t.Fatalf("expected all 3 channels returned when all are in cooldown, got %v", fallbackAll) + } +}