初始化提交

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

@@ -0,0 +1,169 @@
+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)
+}