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

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
}