hunjing/mcp/web-search/internal/balancer/keypool.go

package balancer

import (
	"math"
	"sync"
	"time"

	"github.com/chenxin/hunjing/mcp/web-search/internal/models"
)

// KeyStatus 密钥状态。
type KeyStatus int

const (
	KeyActive   KeyStatus = iota // 正常
	KeyCooldown                  // 冷却中（熔断后）
	KeyDead                      // 死亡（多次失败不再恢复）
)

// KeyState 单个密钥状态。
type KeyState struct {
	Provider models.ProviderKind
	Key      string
	Weight   int

	mu           sync.Mutex
	status       KeyStatus
	latencyEWMA  time.Duration // 指数加权滑动平均延迟
	successCount int64
	failCount    int64
	cb           *CircuitBreaker
}

func newKeyState(provider models.ProviderKind, key string, weight int, cbThreshold int, cbCooldown time.Duration) *KeyState {
	return &KeyState{
		Provider: provider,
		Key:      key,
		Weight:   weight,
		status:   KeyActive,
		cb:       NewCircuitBreaker(cbThreshold, cbCooldown),
	}
}

// Allow 检查是否允许使用此密钥。
func (ks *KeyState) Allow() bool {
	ks.mu.Lock()
	defer ks.mu.Unlock()
	if ks.status != KeyActive {
		return false
	}
	return ks.cb.Allow()
}

// RecordSuccess 记录成功调用，更新 EWMA 延迟。
func (ks *KeyState) RecordSuccess(latency time.Duration) {
	ks.mu.Lock()
	defer ks.mu.Unlock()
	ks.successCount++
	ks.cb.Success()
	// EWMA 更新：α = 0.3
	alpha := 0.3
	if ks.latencyEWMA == 0 {
		ks.latencyEWMA = latency
	} else {
		ks.latencyEWMA = time.Duration(
			(1-alpha)*float64(ks.latencyEWMA) + alpha*float64(latency),
		)
	}
	ks.status = KeyActive
}

// RecordFailure 记录失败调用。
func (ks *KeyState) RecordFailure() {
	ks.mu.Lock()
	defer ks.mu.Unlock()
	ks.failCount++
	ks.cb.Failure()
	if ks.cb.State() == StateOpen {
		ks.status = KeyCooldown
	}
}

// CombinedScore 计算综合分数（越高优先选用）。
// 综合 = 基础权重 × 延迟惩罚 × 可靠性因子
func (ks *KeyState) CombinedScore() float64 {
	ks.mu.Lock()
	defer ks.mu.Unlock()

	if ks.status != KeyActive {
		return 0
	}
	if ks.cb.State() != StateClosed {
		return 0
	}

	base := float64(ks.Weight)
	latencyPenalty := 1.0
	if ks.latencyEWMA > 0 {
		latencyPenalty = 1.0 / (1.0 + ks.latencyEWMA.Seconds())
	}
	reliability := 1.0
	total := ks.successCount + ks.failCount
	if total > 0 && ks.failCount > 0 {
		reliability = 1.0 - float64(ks.failCount)/float64(total)
	}
	return base * latencyPenalty * math.Max(0.1, reliability)
}

// KeyPool 管理一个 provider 下的多密钥。
type KeyPool struct {
	provider    models.ProviderKind
	keys        []*KeyState
	mu          sync.RWMutex
	cbThreshold int
	cbCooldown  time.Duration
}

func NewKeyPool(provider models.ProviderKind, keys []string, weight int, cbThreshold int, cbCooldown time.Duration) *KeyPool {
	kp := &KeyPool{
		provider:    provider,
		keys:        make([]*KeyState, 0, len(keys)),
		cbThreshold: cbThreshold,
		cbCooldown:  cbCooldown,
	}
	for _, k := range keys {
		kp.keys = append(kp.keys, newKeyState(provider, k, weight, cbThreshold, cbCooldown))
	}
	return kp
}

// Select 按综合分数选最佳密钥。
func (kp *KeyPool) Select() *KeyState {
	kp.mu.RLock()
	defer kp.mu.RUnlock()

	var best *KeyState
	var bestScore float64
	for _, ks := range kp.keys {
		if !ks.Allow() {
			continue
		}
		score := ks.CombinedScore()
		if score > bestScore {
			bestScore = score
			best = ks
		}
	}
	return best
}

// ActiveCount 返回活跃密钥数。
func (kp *KeyPool) ActiveCount() int {
	kp.mu.RLock()
	defer kp.mu.RUnlock()
	count := 0
	for _, ks := range kp.keys {
		if ks.status == KeyActive {
			count++
		}
	}
	return count
}

// TotalCount 返回密钥总数。
func (kp *KeyPool) TotalCount() int {
	kp.mu.RLock()
	defer kp.mu.RUnlock()
	return len(kp.keys)
}