go.charczuk.com@v0.0.0-20240327042549-bc490516bd1a/sdk/breaker/breaker.go

package breaker

import (
	"context"
	"sync"
	"sync/atomic"
	"time"

	"go.charczuk.com/sdk/errutil"
)

// New creates a new breaker with the given options.
func New[A, B any](action Action[A, B], options ...Option[A, B]) *Breaker[A, B] {
	b := Breaker[A, B]{
		Action: action,
	}
	for _, opt := range options {
		opt(&b)
	}
	return &b
}

type (
	// OnStateChangeFunc is called when the state changes.
	OnStateChangeFunc func(ctx context.Context, from, to State, generation int64)
	// ShouldOpenFunc returns if the breaker should open.
	ShouldOpenFunc[A any] func(ctx context.Context, counts Counts, args A) bool
	// NowFunc returns the current time.
	NowFunc func() time.Time
)

type Option[A, B any] func(b *Breaker[A, B])

func OptConfig[A, B any](cfg Config) Option[A, B] {
	return func(b *Breaker[A, B]) {
		b.Config = cfg
	}
}

func OptOpenAction[A, B any](action Action[A, B]) Option[A, B] {
	return func(b *Breaker[A, B]) {
		b.OpenAction = action
	}
}

func OptShouldOpen[A, B any](shouldOpenFunc ShouldOpenFunc[A]) Option[A, B] {
	return func(b *Breaker[A, B]) {
		b.ShouldOpen = shouldOpenFunc
	}
}

func OptNow[A, B any](nowFn NowFunc) Option[A, B] {
	return func(b *Breaker[A, B]) {
		b.Now = nowFn
	}
}

var _ Action[any, any] = (*Breaker[any, any])(nil)

// Breaker is a state machine to prevent performing actions that are likely to fail.
type Breaker[A, B any] struct {
	Action        Action[A, B]
	OpenAction    Action[A, B]
	OnStateChange OnStateChangeFunc
	ShouldOpen    ShouldOpenFunc[A]
	Now           NowFunc
	Config        Config

	mu                     sync.Mutex
	state                  State
	generation             int64
	openStateExpiresAt     time.Time
	closedFailuresExpireAt time.Time
	counts                 Counts
}

func (b *Breaker[A, B]) Counts() Counts { return b.counts }

// Call invokes the action and returns the result.
func (b *Breaker[A, B]) Call(ctx context.Context, args A) (res B, err error) {
	var generation int64
	generation, err = b.beforeAction(ctx)
	if err != nil {
		if b.OpenAction != nil {
			res, err = b.OpenAction.Call(ctx, args)
			return
		}
		return
	}
	defer func() {
		if r := recover(); r != nil {
			b.afterAction(ctx, generation, false /*success*/, args)
			err = errutil.New(r)
		} else {
			b.afterAction(ctx, generation, err == nil /*success*/, args)
		}
	}()
	res, err = b.Action.Call(ctx, args)
	return
}

// EvaluateState returns the current state of the CircuitBreaker.
//
// This method is a kind of idirect because can't know for sure
// what the state is at a given time without evaluating expiration
// times and potentially calling handlers if the state changes
// after an expiry.
//
// As a result this method takes a context, and may call the `OnStateChange`
// delegate value if it's set.
func (b *Breaker[A, B]) EvaluateState(ctx context.Context) State {
	b.mu.Lock()
	defer b.mu.Unlock()

	now := b.now()
	state, _ := b.evaluateStateUnsafe(ctx, now)
	return state
}

//
// internal methods
//

func (b *Breaker[A, B]) beforeAction(ctx context.Context) (int64, error) {
	b.mu.Lock()
	defer b.mu.Unlock()

	now := b.now()
	state, generation := b.evaluateStateUnsafe(ctx, now)
	if state == StateOpen {
		return generation, ErrOpenState
	} else if state == StateHalfOpen && b.counts.Requests >= b.Config.HalfOpenMaxActionsOrDefault() {
		return generation, ErrTooManyRequests
	}
	atomic.AddUint64(&b.counts.Requests, 1)
	return generation, nil
}

func (b *Breaker[A, B]) afterAction(ctx context.Context, currentGeneration int64, success bool, args A) {
	b.mu.Lock()
	defer b.mu.Unlock()

	now := b.now()
	state, generation := b.evaluateStateUnsafe(ctx, now)
	if generation != currentGeneration {
		return
	}
	if success {
		b.success(ctx, state, now)
		return
	}
	b.failure(ctx, state, now, args)
}

func (b *Breaker[A, B]) success(ctx context.Context, state State, now time.Time) {
	switch state {
	case StateClosed:
		atomic.AddUint64(&b.counts.TotalSuccesses, 1)
		atomic.AddUint64(&b.counts.ConsecutiveSuccesses, 1)
		atomic.StoreUint64(&b.counts.ConsecutiveFailures, 0)
	case StateHalfOpen:
		atomic.AddUint64(&b.counts.TotalSuccesses, 1)
		atomic.AddUint64(&b.counts.ConsecutiveSuccesses, 1)
		atomic.StoreUint64(&b.counts.ConsecutiveFailures, 0)
		if b.counts.ConsecutiveSuccesses >= b.Config.HalfOpenMaxActionsOrDefault() {
			b.setStateUnsafe(ctx, StateClosed, now)
		}
	}
}

func (b *Breaker[A, B]) failure(ctx context.Context, state State, now time.Time, args A) {
	switch state {
	case StateClosed:
		atomic.AddUint64(&b.counts.TotalFailures, 1)
		atomic.AddUint64(&b.counts.ConsecutiveFailures, 1)
		atomic.StoreUint64(&b.counts.ConsecutiveSuccesses, 0)
		if b.shouldOpen(ctx, args) {
			b.setStateUnsafe(ctx, StateOpen, now)
		}
	case StateHalfOpen:
		b.setStateUnsafe(ctx, StateOpen, now)
	}
}

func (b *Breaker[A, B]) evaluateStateUnsafe(ctx context.Context, now time.Time) (state State, generation int64) {
	switch b.state {
	case StateClosed:
		if b.closedFailuresExpireAt.IsZero() {
			b.closedFailuresExpireAt = now.Add(b.Config.ClosedFailureExpiryIntervalOrDefault())
		} else if b.closedFailuresExpireAt.Before(now) {
			b.counts = Counts{}
			b.closedFailuresExpireAt = now.Add(b.Config.ClosedFailureExpiryIntervalOrDefault())
		}
	case StateOpen:
		if !b.openStateExpiresAt.IsZero() && b.openStateExpiresAt.Before(now) {
			b.setStateUnsafe(ctx, StateHalfOpen, now)
		}
	}
	return b.state, b.generation
}

func (b *Breaker[A, B]) setStateUnsafe(ctx context.Context, state State, now time.Time) {
	if b.state == state {
		return
	}
	previousState := b.state
	b.state = state
	b.incrementGenerationAfterStateChange(now)
	if b.OnStateChange != nil {
		b.OnStateChange(ctx, previousState, b.state, b.generation)
	}
}

func (b *Breaker[A, B]) incrementGenerationAfterStateChange(now time.Time) {
	atomic.AddInt64(&b.generation, 1)
	b.counts = Counts{}

	var zero time.Time
	switch b.state {
	case StateClosed:
		b.openStateExpiresAt = zero
		b.closedFailuresExpireAt = now.Add(b.Config.ClosedFailureExpiryIntervalOrDefault())
	case StateOpen:
		b.openStateExpiresAt = now.Add(b.Config.OpenExpiryIntervalOrDefault())
		b.closedFailuresExpireAt = zero
	case StateHalfOpen:
		b.openStateExpiresAt = zero
		b.closedFailuresExpireAt = zero
	}
}

func (b *Breaker[A, B]) shouldOpen(ctx context.Context, args A) bool {
	if b.ShouldOpen != nil {
		return b.ShouldOpen(ctx, b.counts, args)
	}
	return b.counts.ConsecutiveFailures >= b.Config.FailureThresholdOrDefault()
}

func (b *Breaker[A, B]) now() time.Time {
	if b.Now != nil {
		return b.Now()
	}
	return time.Now()
}