初始化提交

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

@@ -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()
+}