104 lines
1.9 KiB
Go
104 lines
1.9 KiB
Go
// 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
|
|
}
|