82 lines
1.3 KiB
Go
82 lines
1.3 KiB
Go
package tree
|
|
|
|
import (
|
|
"errors"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type (
|
|
circuitBreaker struct {
|
|
breakDuration time.Duration
|
|
threshold int
|
|
|
|
mu sync.RWMutex
|
|
state map[uint64]state
|
|
}
|
|
|
|
state struct {
|
|
counter int
|
|
breakTimestamp time.Time
|
|
}
|
|
)
|
|
|
|
var ErrCBClosed = errors.New("circuit breaker is closed")
|
|
|
|
func newCircuitBreaker(breakDuration time.Duration, threshold int) *circuitBreaker {
|
|
return &circuitBreaker{
|
|
breakDuration: breakDuration,
|
|
threshold: threshold,
|
|
state: make(map[uint64]state),
|
|
}
|
|
}
|
|
|
|
func (cb *circuitBreaker) checkBreak(id uint64) error {
|
|
cb.mu.RLock()
|
|
s, ok := cb.state[id]
|
|
cb.mu.RUnlock()
|
|
|
|
if ok && time.Since(s.breakTimestamp) < cb.breakDuration {
|
|
return ErrCBClosed
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (cb *circuitBreaker) openBreak(id uint64) {
|
|
cb.mu.Lock()
|
|
defer cb.mu.Unlock()
|
|
delete(cb.state, id)
|
|
}
|
|
|
|
func (cb *circuitBreaker) incError(id uint64) {
|
|
cb.mu.Lock()
|
|
defer cb.mu.Unlock()
|
|
|
|
s := cb.state[id]
|
|
|
|
s.counter++
|
|
if s.counter >= cb.threshold {
|
|
s.counter = cb.threshold
|
|
if time.Since(s.breakTimestamp) >= cb.breakDuration {
|
|
s.breakTimestamp = time.Now()
|
|
}
|
|
}
|
|
|
|
cb.state[id] = s
|
|
}
|
|
|
|
func (cb *circuitBreaker) Do(id uint64, f func() error) error {
|
|
if err := cb.checkBreak(id); err != nil {
|
|
return err
|
|
}
|
|
|
|
err := f()
|
|
if err == nil {
|
|
cb.openBreak(id)
|
|
} else {
|
|
cb.incError(id)
|
|
}
|
|
|
|
return err
|
|
}
|