frostfs-sdk-go/pool/tree/circuitbreaker.go

package tree

import (
	"errors"
	"sync"
	"time"
)

type (
	circuitBreaker struct {
		breakDuration time.Duration
		threshold     int

		mu    sync.Mutex
		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 (c *circuitBreaker) Do(id uint64, f func() error) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if _, ok := c.state[id]; !ok {
		c.state[id] = state{}
	}

	s := c.state[id]
	defer func() {
		c.state[id] = s
	}()

	if time.Since(s.breakTimestamp) < c.breakDuration {
		return ErrCBClosed
	}

	err := f()
	if err == nil {
		s.counter = 0
		return nil
	}

	s.counter++
	if s.counter >= c.threshold {
		s.counter = c.threshold
		s.breakTimestamp = time.Now()
	}

	return err
}