frostfs-sdk-go/pool/tree/circuitbreaker.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) 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
}