初始化提交
This commit is contained in:
chenxin committed
commit
887c1f449e
963 files changed
+166869
No files matched your search
@@ -0,0 +1,264 @@
|
||||
package balancer
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/config"
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
|
||||
)
|
||||
|
||||
// KeySearcher 按密钥执行搜索的接口。
|
||||
// 与 SearchProvider 解耦，方便测试时注入 mock。
|
||||
type KeySearcher interface {
|
||||
Search(ctx context.Context, apiKey string, req *models.SearchRequest) (*models.SearchResult, error)
|
||||
HealthCheck(ctx context.Context) error
|
||||
Name() models.ProviderKind
|
||||
}
|
||||
|
||||
// LoadBalancer 负载均衡器（goroutine-safe）。
|
||||
type LoadBalancer struct {
|
||||
pools map[models.ProviderKind]*KeyPool
|
||||
searchers map[models.ProviderKind]KeySearcher
|
||||
disabled map[models.ProviderKind]bool // 运行时动态禁用
|
||||
retryCount int
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
func NewLoadBalancer(cfg *config.BalancerConfig) *LoadBalancer {
|
||||
return &LoadBalancer{
|
||||
pools: make(map[models.ProviderKind]*KeyPool),
|
||||
searchers: make(map[models.ProviderKind]KeySearcher),
|
||||
disabled: make(map[models.ProviderKind]bool),
|
||||
retryCount: cfg.RetryCount,
|
||||
}
|
||||
}
|
||||
|
||||
// AddProvider 注册 provider 及其密钥池。
|
||||
func (lb *LoadBalancer) AddProvider(
|
||||
provider models.ProviderKind,
|
||||
keys []string,
|
||||
weight int,
|
||||
searcher KeySearcher,
|
||||
) {
|
||||
lb.mu.Lock()
|
||||
defer lb.mu.Unlock()
|
||||
|
||||
cbThreshold := 5
|
||||
cbCooldown := 30 * time.Second
|
||||
lb.pools[provider] = NewKeyPool(provider, keys, weight, cbThreshold, cbCooldown)
|
||||
lb.searchers[provider] = searcher
|
||||
// 添加时默认启用
|
||||
delete(lb.disabled, provider)
|
||||
}
|
||||
|
||||
// RemoveProvider 移除 provider 及其密钥池。
|
||||
func (lb *LoadBalancer) RemoveProvider(provider models.ProviderKind) {
|
||||
lb.mu.Lock()
|
||||
defer lb.mu.Unlock()
|
||||
delete(lb.pools, provider)
|
||||
delete(lb.searchers, provider)
|
||||
delete(lb.disabled, provider)
|
||||
}
|
||||
|
||||
// EnableProvider 启用 provider。
|
||||
func (lb *LoadBalancer) EnableProvider(provider models.ProviderKind) {
|
||||
lb.mu.Lock()
|
||||
defer lb.mu.Unlock()
|
||||
delete(lb.disabled, provider)
|
||||
}
|
||||
|
||||
// DisableProvider 禁用 provider（搜索时跳过，不删除配置）。
|
||||
func (lb *LoadBalancer) DisableProvider(provider models.ProviderKind) {
|
||||
lb.mu.Lock()
|
||||
defer lb.mu.Unlock()
|
||||
lb.disabled[provider] = true
|
||||
}
|
||||
|
||||
// IsProviderEnabled 检查 provider 是否启用。
|
||||
func (lb *LoadBalancer) IsProviderEnabled(provider models.ProviderKind) bool {
|
||||
lb.mu.RLock()
|
||||
defer lb.mu.RUnlock()
|
||||
return !lb.disabled[provider]
|
||||
}
|
||||
|
||||
// SetRetryCount 更新重试次数。
|
||||
func (lb *LoadBalancer) SetRetryCount(n int) {
|
||||
lb.mu.Lock()
|
||||
defer lb.mu.Unlock()
|
||||
lb.retryCount = n
|
||||
}
|
||||
|
||||
// HasProvider 检查 provider 是否已注册。
|
||||
func (lb *LoadBalancer) HasProvider(provider models.ProviderKind) bool {
|
||||
lb.mu.RLock()
|
||||
defer lb.mu.RUnlock()
|
||||
_, ok := lb.pools[provider]
|
||||
return ok
|
||||
}
|
||||
|
||||
// Search 执行搜索（自动选密钥、跨 provider 重试）。
|
||||
func (lb *LoadBalancer) Search(ctx context.Context, req *models.SearchRequest) (*models.SearchResult, error) {
|
||||
providers := lb.resolveProviders(req.Providers)
|
||||
|
||||
var lastErr error
|
||||
for attempt := 0; attempt <= lb.getRetryCount(); attempt++ {
|
||||
for _, prov := range providers {
|
||||
pool, ok := lb.pools[prov]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
ks := pool.Select()
|
||||
if ks == nil {
|
||||
continue
|
||||
}
|
||||
searcher := lb.searchers[prov]
|
||||
if searcher == nil {
|
||||
continue
|
||||
}
|
||||
|
||||
start := time.Now()
|
||||
result, err := searcher.Search(ctx, ks.Key, req)
|
||||
duration := time.Since(start)
|
||||
|
||||
if err != nil {
|
||||
ks.RecordFailure()
|
||||
lastErr = err
|
||||
continue
|
||||
}
|
||||
ks.RecordSuccess(duration)
|
||||
result.Provider = prov
|
||||
return result, nil
|
||||
}
|
||||
}
|
||||
|
||||
if lastErr != nil {
|
||||
return nil, fmt.Errorf("balancer: all providers exhausted: %w", lastErr)
|
||||
}
|
||||
return nil, fmt.Errorf("balancer: no available provider or key for query: %s", req.Query)
|
||||
}
|
||||
|
||||
// SearchAll 并发搜索所有启用 provider（用于 Merge 模式）。
|
||||
func (lb *LoadBalancer) SearchAll(ctx context.Context, req *models.SearchRequest) []*models.SearchResult {
|
||||
providers := lb.resolveProviders(nil)
|
||||
type result struct {
|
||||
prov models.ProviderKind
|
||||
res *models.SearchResult
|
||||
err error
|
||||
}
|
||||
|
||||
ch := make(chan result, len(providers))
|
||||
for _, prov := range providers {
|
||||
p := prov
|
||||
go func() {
|
||||
pool, ok := lb.pools[p]
|
||||
if !ok {
|
||||
ch <- result{prov: p, err: fmt.Errorf("no pool")}
|
||||
return
|
||||
}
|
||||
ks := pool.Select()
|
||||
if ks == nil {
|
||||
ch <- result{prov: p, err: fmt.Errorf("no available key")}
|
||||
return
|
||||
}
|
||||
searcher := lb.searchers[p]
|
||||
if searcher == nil {
|
||||
ch <- result{prov: p, err: fmt.Errorf("no searcher")}
|
||||
return
|
||||
}
|
||||
start := time.Now()
|
||||
res, err := searcher.Search(ctx, ks.Key, req)
|
||||
duration := time.Since(start)
|
||||
if err != nil {
|
||||
ks.RecordFailure()
|
||||
ch <- result{prov: p, err: err}
|
||||
return
|
||||
}
|
||||
ks.RecordSuccess(duration)
|
||||
res.Provider = p
|
||||
ch <- result{prov: p, res: res}
|
||||
}()
|
||||
}
|
||||
|
||||
var results []*models.SearchResult
|
||||
for range providers {
|
||||
r := <-ch
|
||||
if r.res != nil {
|
||||
results = append(results, r.res)
|
||||
}
|
||||
}
|
||||
return results
|
||||
}
|
||||
|
||||
// ActiveKeyCount 返回每个 provider 的活跃密钥数。
|
||||
// ActiveKeyCount 返回每个 provider 的活跃密钥数。
|
||||
// ActiveKeyCount 返回每个 provider 的活跃密钥数。
|
||||
func (lb *LoadBalancer) ActiveKeyCount() map[models.ProviderKind]int {
|
||||
lb.mu.RLock()
|
||||
defer lb.mu.RUnlock()
|
||||
out := make(map[models.ProviderKind]int)
|
||||
for p, pool := range lb.pools {
|
||||
if lb.disabled[p] {
|
||||
continue
|
||||
}
|
||||
out[p] = pool.ActiveCount()
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Providers 返回所有已注册的 provider 名称（无论启用与否）。
|
||||
func (lb *LoadBalancer) Providers() []models.ProviderKind {
|
||||
lb.mu.RLock()
|
||||
defer lb.mu.RUnlock()
|
||||
out := make([]models.ProviderKind, 0, len(lb.pools))
|
||||
for p := range lb.pools {
|
||||
out = append(out, p)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// HealthCheck 对指定 provider 执行健康检查（委托给 KeySearcher）。
|
||||
func (lb *LoadBalancer) HealthCheck(ctx context.Context, provider models.ProviderKind) error {
|
||||
lb.mu.RLock()
|
||||
searcher, ok := lb.searchers[provider]
|
||||
lb.mu.RUnlock()
|
||||
if !ok {
|
||||
return fmt.Errorf("healthcheck: provider not found: %s", provider)
|
||||
}
|
||||
return searcher.HealthCheck(ctx)
|
||||
}
|
||||
|
||||
func (lb *LoadBalancer) getRetryCount() int {
|
||||
lb.mu.RLock()
|
||||
defer lb.mu.RUnlock()
|
||||
return lb.retryCount
|
||||
}
|
||||
|
||||
// resolveProviders 按用户偏好和可用性确定 provider 顺序（跳过已禁用的）。
|
||||
func (lb *LoadBalancer) resolveProviders(preferred []string) []models.ProviderKind {
|
||||
lb.mu.RLock()
|
||||
defer lb.mu.RUnlock()
|
||||
|
||||
if len(preferred) == 0 {
|
||||
out := make([]models.ProviderKind, 0, len(lb.pools))
|
||||
for p := range lb.pools {
|
||||
if !lb.disabled[p] {
|
||||
out = append(out, p)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
out := make([]models.ProviderKind, 0, len(preferred))
|
||||
seen := make(map[models.ProviderKind]bool)
|
||||
for _, name := range preferred {
|
||||
p := models.ProviderKind(name)
|
||||
if _, ok := lb.pools[p]; ok && !seen[p] && !lb.disabled[p] {
|
||||
out = append(out, p)
|
||||
seen[p] = true
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,90 @@
|
||||
package balancer
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/config"
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
|
||||
)
|
||||
|
||||
type mockSearcher struct {
|
||||
name models.ProviderKind
|
||||
failOn int // 第几次调用失败（0 = 永不）
|
||||
calls int
|
||||
}
|
||||
|
||||
func (m *mockSearcher) Search(ctx context.Context, apiKey string, req *models.SearchRequest) (*models.SearchResult, error) {
|
||||
m.calls++
|
||||
if m.failOn > 0 && m.calls >= m.failOn {
|
||||
return nil, errors.New("mock failure")
|
||||
}
|
||||
return &models.SearchResult{Provider: m.name, Query: req.Query, Results: []models.ResultItem{{Title: "mock"}}}, nil
|
||||
}
|
||||
func (m *mockSearcher) HealthCheck(ctx context.Context) error { return nil }
|
||||
func (m *mockSearcher) Name() models.ProviderKind { return m.name }
|
||||
|
||||
func TestLoadBalancer_Search_Basic(t *testing.T) {
|
||||
lb := NewLoadBalancer(&config.BalancerConfig{RetryCount: 1})
|
||||
lb.AddProvider(models.ProviderTavily, []string{"key1"}, 1, &mockSearcher{name: models.ProviderTavily})
|
||||
|
||||
result, err := lb.Search(context.Background(), &models.SearchRequest{Query: "test"})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if result.Provider != models.ProviderTavily {
|
||||
t.Errorf("expected tavily, got %s", result.Provider)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadBalancer_Search_FallbackProvider(t *testing.T) {
|
||||
lb := NewLoadBalancer(&config.BalancerConfig{RetryCount: 1})
|
||||
lb.AddProvider(models.ProviderTavily, []string{"key1"}, 1, &mockSearcher{name: models.ProviderTavily, failOn: 1})
|
||||
lb.AddProvider(models.ProviderBrave, []string{"key1"}, 1, &mockSearcher{name: models.ProviderBrave})
|
||||
|
||||
result, err := lb.Search(context.Background(), &models.SearchRequest{Query: "test"})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if string(result.Provider) != "brave" {
|
||||
t.Errorf("expected fallback to brave, got %s", result.Provider)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadBalancer_Search_AllExhausted(t *testing.T) {
|
||||
lb := NewLoadBalancer(&config.BalancerConfig{RetryCount: 0})
|
||||
lb.AddProvider(models.ProviderTavily, []string{"key1"}, 1, &mockSearcher{name: models.ProviderTavily, failOn: 1})
|
||||
|
||||
_, err := lb.Search(context.Background(), &models.SearchRequest{Query: "test"})
|
||||
if err == nil {
|
||||
t.Fatal("expected error when all keys exhausted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadBalancer_Search_PreferredProviders(t *testing.T) {
|
||||
lb := NewLoadBalancer(&config.BalancerConfig{RetryCount: 0})
|
||||
lb.AddProvider(models.ProviderTavily, []string{"key1"}, 1, &mockSearcher{name: models.ProviderTavily})
|
||||
lb.AddProvider(models.ProviderBrave, []string{"key1"}, 1, &mockSearcher{name: models.ProviderBrave})
|
||||
|
||||
result, err := lb.Search(context.Background(), &models.SearchRequest{
|
||||
Query: "test",
|
||||
Providers: []string{"brave"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if result.Provider != models.ProviderBrave {
|
||||
t.Errorf("expected bing (preferred), got %s", result.Provider)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadBalancer_ActiveKeyCount(t *testing.T) {
|
||||
lb := NewLoadBalancer(&config.BalancerConfig{RetryCount: 0})
|
||||
lb.AddProvider(models.ProviderTavily, []string{"key1", "key2"}, 1, &mockSearcher{name: models.ProviderTavily})
|
||||
|
||||
counts := lb.ActiveKeyCount()
|
||||
if counts[models.ProviderTavily] != 2 {
|
||||
t.Errorf("expected 2 active keys, got %d", counts[models.ProviderTavily])
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,117 @@
|
||||
package balancer
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/config"
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
|
||||
)
|
||||
|
||||
// HealthChecker 定期对已注册 provider 执行健康检查。
|
||||
// 连续失败超过阈值时自动禁用 provider，恢复成功后自动启用。
|
||||
type HealthChecker struct {
|
||||
lb *LoadBalancer
|
||||
interval time.Duration
|
||||
timeout time.Duration
|
||||
threshold int
|
||||
|
||||
mu sync.Mutex
|
||||
failures map[models.ProviderKind]int // 连续失败计数
|
||||
done chan struct{}
|
||||
}
|
||||
|
||||
// NewHealthChecker 创建健康检查器。cfg 为对应 provider 的健康检查配置。
|
||||
func NewHealthChecker(lb *LoadBalancer, cfg config.HealthCheckConf) *HealthChecker {
|
||||
return &HealthChecker{
|
||||
lb: lb,
|
||||
interval: cfg.Interval,
|
||||
timeout: cfg.Timeout,
|
||||
threshold: cfg.FailureThreshold,
|
||||
failures: make(map[models.ProviderKind]int),
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// Start 在后台 goroutine 中启动健康检查循环。
|
||||
// 每隔 interval 时间对所有已注册 provider 执行一次健康检查。
|
||||
func (hc *HealthChecker) Start() {
|
||||
if hc.interval <= 0 {
|
||||
return
|
||||
}
|
||||
slog.Info("健康检查已启动", "interval", hc.interval, "timeout", hc.timeout, "failure_threshold", hc.threshold)
|
||||
go hc.loop()
|
||||
}
|
||||
|
||||
// Stop 停止健康检查循环。
|
||||
func (hc *HealthChecker) Stop() {
|
||||
close(hc.done)
|
||||
}
|
||||
|
||||
func (hc *HealthChecker) loop() {
|
||||
ticker := time.NewTicker(hc.interval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ticker.C:
|
||||
hc.checkAll()
|
||||
case <-hc.done:
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (hc *HealthChecker) checkAll() {
|
||||
// 获取当前所有已注册 provider 的名称
|
||||
// 我们需要从 LoadBalancer 获取提供者列表
|
||||
// 由于 LoadBalancer 没有直接暴露 provider 列表的 API，我们遍历 ActiveKeyCount
|
||||
// 来获取所有已知的 provider
|
||||
for _, prov := range hc.getAllProviders() {
|
||||
hc.checkProvider(prov)
|
||||
}
|
||||
}
|
||||
|
||||
func (hc *HealthChecker) checkProvider(prov models.ProviderKind) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), hc.timeout)
|
||||
defer cancel()
|
||||
|
||||
err := hc.lb.HealthCheck(ctx, prov)
|
||||
|
||||
hc.mu.Lock()
|
||||
defer hc.mu.Unlock()
|
||||
|
||||
if err != nil {
|
||||
hc.failures[prov]++
|
||||
fails := hc.failures[prov]
|
||||
slog.Warn("健康检查失败",
|
||||
"provider", prov,
|
||||
"consecutive_failures", fails,
|
||||
"error", err,
|
||||
)
|
||||
if fails >= hc.threshold && hc.lb.IsProviderEnabled(prov) {
|
||||
hc.lb.DisableProvider(prov)
|
||||
slog.Warn("健康检查: provider 已自动禁用",
|
||||
"provider", prov,
|
||||
"failures", fails,
|
||||
)
|
||||
}
|
||||
} else {
|
||||
// 恢复成功
|
||||
if hc.failures[prov] > 0 {
|
||||
slog.Info("健康检查恢复", "provider", prov)
|
||||
}
|
||||
hc.failures[prov] = 0
|
||||
if !hc.lb.IsProviderEnabled(prov) {
|
||||
hc.lb.EnableProvider(prov)
|
||||
slog.Info("健康检查: provider 已自动启用", "provider", prov)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// getAllProviders 返回当前所有已注册的 provider 名称（包括已禁用的）。
|
||||
func (hc *HealthChecker) getAllProviders() []models.ProviderKind {
|
||||
return hc.lb.Providers()
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
// Package balancer 实现密钥池管理和负载均衡。
|
||||
package balancer
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// State 熔断器状态。
|
||||
type State int
|
||||
|
||||
const (
|
||||
StateClosed State = iota // 正常
|
||||
StateOpen // 熔断
|
||||
StateHalfOpen // 半开（探测）
|
||||
)
|
||||
|
||||
func (s State) String() string {
|
||||
switch s {
|
||||
case StateClosed:
|
||||
return "closed"
|
||||
case StateOpen:
|
||||
return "open"
|
||||
case StateHalfOpen:
|
||||
return "half-open"
|
||||
default:
|
||||
return "unknown"
|
||||
}
|
||||
}
|
||||
|
||||
// CircuitBreaker 熔断器。
|
||||
type CircuitBreaker struct {
|
||||
mu sync.Mutex
|
||||
state State
|
||||
failures int
|
||||
threshold int
|
||||
cooldown time.Duration
|
||||
lastFailureTime time.Time
|
||||
}
|
||||
|
||||
func NewCircuitBreaker(threshold int, cooldown time.Duration) *CircuitBreaker {
|
||||
return &CircuitBreaker{
|
||||
state: StateClosed,
|
||||
threshold: threshold,
|
||||
cooldown: cooldown,
|
||||
}
|
||||
}
|
||||
|
||||
// Allow 判断是否允许请求通过。
|
||||
func (cb *CircuitBreaker) Allow() bool {
|
||||
cb.mu.Lock()
|
||||
defer cb.mu.Unlock()
|
||||
|
||||
switch cb.state {
|
||||
case StateClosed:
|
||||
return true
|
||||
case StateOpen:
|
||||
if time.Since(cb.lastFailureTime) > cb.cooldown {
|
||||
cb.state = StateHalfOpen
|
||||
return true
|
||||
}
|
||||
return false
|
||||
case StateHalfOpen:
|
||||
return true
|
||||
default:
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
// Success 记录成功。
|
||||
func (cb *CircuitBreaker) Success() {
|
||||
cb.mu.Lock()
|
||||
defer cb.mu.Unlock()
|
||||
cb.failures = 0
|
||||
if cb.state == StateHalfOpen {
|
||||
cb.state = StateClosed
|
||||
}
|
||||
}
|
||||
|
||||
// Failure 记录失败。
|
||||
func (cb *CircuitBreaker) Failure() {
|
||||
cb.mu.Lock()
|
||||
defer cb.mu.Unlock()
|
||||
cb.failures++
|
||||
cb.lastFailureTime = time.Now()
|
||||
if cb.state == StateHalfOpen || cb.failures >= cb.threshold {
|
||||
cb.state = StateOpen
|
||||
}
|
||||
}
|
||||
|
||||
// State 返回当前状态（goroutine-safe 快照）。
|
||||
func (cb *CircuitBreaker) State() State {
|
||||
cb.mu.Lock()
|
||||
defer cb.mu.Unlock()
|
||||
return cb.state
|
||||
}
|
||||
|
||||
// Failures 返回连续失败次数。
|
||||
func (cb *CircuitBreaker) Failures() int {
|
||||
cb.mu.Lock()
|
||||
defer cb.mu.Unlock()
|
||||
return cb.failures
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -0,0 +1,123 @@
|
||||
package balancer
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
|
||||
)
|
||||
|
||||
func TestKeyPool_Select_Basic(t *testing.T) {
|
||||
kp := NewKeyPool(models.ProviderTavily, []string{"key1", "key2"}, 1, 5, 30*time.Second)
|
||||
ks := kp.Select()
|
||||
if ks == nil {
|
||||
t.Fatal("expected to select a key")
|
||||
}
|
||||
if ks.Key != "key1" && ks.Key != "key2" {
|
||||
t.Errorf("unexpected key: %s", ks.Key)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyPool_Select_SkipsDead(t *testing.T) {
|
||||
kp := NewKeyPool(models.ProviderTavily, []string{"key1", "key2"}, 1, 2, 1*time.Hour)
|
||||
ks1 := kp.Select()
|
||||
|
||||
// Fail twice to trigger circuit breaker
|
||||
ks1.RecordFailure()
|
||||
ks1.RecordFailure()
|
||||
|
||||
selected := kp.Select()
|
||||
if selected == nil {
|
||||
t.Fatal("expected fallback to second key")
|
||||
}
|
||||
if selected.Key == ks1.Key {
|
||||
t.Error("should not select the failed key")
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyState_CombinedScore(t *testing.T) {
|
||||
ks := newKeyState(models.ProviderTavily, "test-key", 2, 5, 30*time.Second)
|
||||
score := ks.CombinedScore()
|
||||
if score <= 0 {
|
||||
t.Errorf("expected positive score, got %f", score)
|
||||
}
|
||||
|
||||
// After recording a slow latency, score should drop
|
||||
ks.RecordSuccess(2 * time.Second)
|
||||
slowScore := ks.CombinedScore()
|
||||
if slowScore >= score {
|
||||
t.Error("expected slow latency to reduce score")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCircuitBreaker_OpenAndRecover(t *testing.T) {
|
||||
cb := NewCircuitBreaker(3, 50*time.Millisecond)
|
||||
|
||||
// Should be closed initially
|
||||
if !cb.Allow() {
|
||||
t.Error("expected allow initially")
|
||||
}
|
||||
|
||||
// Fail 3 times to open
|
||||
cb.Failure()
|
||||
cb.Failure()
|
||||
cb.Failure()
|
||||
if cb.Allow() {
|
||||
t.Error("expected circuit open after 3 failures")
|
||||
}
|
||||
|
||||
// Wait for cooldown
|
||||
time.Sleep(60 * time.Millisecond)
|
||||
if !cb.Allow() {
|
||||
t.Error("expected allow after cooldown")
|
||||
}
|
||||
|
||||
// Half-open: success should close
|
||||
cb.Success()
|
||||
if cb.State() != StateClosed {
|
||||
t.Error("expected circuit closed after half-open success")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCircuitBreaker_HalfOpenFailure(t *testing.T) {
|
||||
cb := NewCircuitBreaker(3, 50*time.Millisecond)
|
||||
|
||||
cb.Failure()
|
||||
cb.Failure()
|
||||
cb.Failure()
|
||||
if cb.Allow() {
|
||||
t.Error("expected open")
|
||||
}
|
||||
|
||||
time.Sleep(60 * time.Millisecond)
|
||||
if !cb.Allow() {
|
||||
t.Error("expected half-open allow")
|
||||
}
|
||||
|
||||
// Half-open: failure should re-open immediately
|
||||
cb.Failure()
|
||||
if cb.State() != StateOpen {
|
||||
t.Error("expected open after half-open failure")
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyPool_ActiveCount(t *testing.T) {
|
||||
kp := NewKeyPool(models.ProviderTavily, []string{"key1", "key2", "key3"}, 1, 3, 30*time.Second)
|
||||
if kp.TotalCount() != 3 {
|
||||
t.Errorf("expected 3 total, got %d", kp.TotalCount())
|
||||
}
|
||||
if kp.ActiveCount() != 3 {
|
||||
t.Errorf("expected 3 active, got %d", kp.ActiveCount())
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyPool_AllKeysDead(t *testing.T) {
|
||||
kp := NewKeyPool(models.ProviderTavily, []string{"key1"}, 1, 1, 1*time.Hour)
|
||||
ks := kp.Select()
|
||||
ks.RecordFailure()
|
||||
|
||||
selected := kp.Select()
|
||||
if selected != nil {
|
||||
t.Error("expected nil when all keys dead")
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user