new-api/service/circuit_breaker.go

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)
}