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) breakTime(id uint64) (time.Time, bool) { cb.mu.RLock() defer cb.mu.RUnlock() if s, ok := cb.state[id]; ok { return s.breakTimestamp, true } return time.Time{}, false } func (cb *circuitBreaker) openBreak(id uint64) { cb.mu.Lock() defer cb.mu.Unlock() delete(cb.state, id) } func (cb *circuitBreaker) incError(id uint64, doTime time.Time) { cb.mu.Lock() defer cb.mu.Unlock() s := cb.state[id] s.counter++ if s.counter >= cb.threshold { s.counter = cb.threshold if s.breakTimestamp.Before(doTime) { s.breakTimestamp = doTime } } cb.state[id] = s } func (c *circuitBreaker) Do(id uint64, f func() error) error { breakTime, ok := c.breakTime(id) if ok && time.Since(breakTime) < c.breakDuration { return ErrCBClosed } // Use this timestamp to update circuit breaker in case of an error. // f() may be blocked for unpredictable duration, so concurrent calls // may update time in 'incError' endlessly and circuit will never be open doTime := time.Now() err := f() if err == nil { c.openBreak(id) } else { c.incError(id, doTime) } return err }