frostfs-sdk-go/pool/tree/circuitbreaker.go
Alex Vanin c8d71c450a [#339] pool/tree: Add circuit breaker
Signed-off-by: Alex Vanin <a.vanin@yadro.com>
2025-03-04 12:20:58 +03:00

83 lines
1.3 KiB
Go

package tree
import (
"context"
"errors"
"sync"
"time"
)
const (
defaultThreshold = 10
defaultBreakDuration = 10 * time.Second
)
type (
circuitBreaker struct {
breakDuration time.Duration
threshold int
mu sync.Mutex
state map[string]state
}
state struct {
counter int
breakTimestamp time.Time
}
dialer interface {
dial(context.Context) error
endpoint() string
}
)
var ErrCBClosed = errors.New("circuit breaker is closed")
func NewCircuitBreaker(breakDuration time.Duration, threshold int) *circuitBreaker {
if threshold == 0 {
threshold = defaultThreshold
}
if breakDuration == 0 {
breakDuration = defaultBreakDuration
}
return &circuitBreaker{
breakDuration: breakDuration,
threshold: threshold,
state: make(map[string]state),
}
}
func (c *circuitBreaker) Dial(ctx context.Context, cli dialer) error {
c.mu.Lock()
defer c.mu.Unlock()
endpoint := cli.endpoint()
if _, ok := c.state[endpoint]; !ok {
c.state[endpoint] = state{}
}
s := c.state[endpoint]
defer func() {
c.state[endpoint] = s
}()
if time.Since(s.breakTimestamp) < c.breakDuration {
return ErrCBClosed
}
err := cli.dial(ctx)
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
}