初始化提交
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
|
||||
}
|
||||
Reference in new issue
Block a user