118 lines
3.0 KiB
Go
118 lines
3.0 KiB
Go
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()
|
|
}
|