170 lines
3.6 KiB
Go
170 lines
3.6 KiB
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)
|
|
}
|