初始化提交

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

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