synapsecns/sanguine

package retry

import (
    "context"
    "errors"
    "fmt"
    "github.com/jpillora/backoff"
    "go.uber.org/multierr"
    "time"
)

// RetryableFunc is a function that can be retried.
type RetryableFunc func(ctx context.Context) error

// retryWithBackoffConfig holds the configuration for WithBackoff.
type retryWithBackoffConfig struct {
    factor float64
    jitter bool
    min    time.Duration
    max    time.Duration
    // maxAttempts sets the maximum number of retry attempts.
    // if this is negative it is ignored
    maxAttempts    int
    maxAttemptTime time.Duration
    // maxAllAttempts sets the maximum time for all attempts.
    // if this is negative it is ignored
    maxAllAttemptsTime time.Duration
}

// returns true if the number of attempts exceeds the maximum number of attempts.
func (r *retryWithBackoffConfig) exceedsMaxAttempts(attempts int) bool {
    return r.maxAttempts > 0 && attempts > r.maxAttempts
}

// returns true if the time for all attempts exceeds the maximum time for all attempts.
func (r *retryWithBackoffConfig) exceedsMaxTime(startTime time.Time) bool {
    return r.maxAllAttemptsTime > 0 && time.Since(startTime) > r.maxAllAttemptsTime
}

// WithBackoffConfigurator configures a retryWithBackoffConfig.
type WithBackoffConfigurator func(*retryWithBackoffConfig)

// WithFactor sets the backoff factor.
func WithFactor(factor float64) WithBackoffConfigurator {
    return func(c *retryWithBackoffConfig) {
        c.factor = factor
    }
}

// WithJitter enables or disables jitter.
func WithJitter(jitter bool) WithBackoffConfigurator {
    return func(c *retryWithBackoffConfig) {
        c.jitter = jitter
    }
}

// WithMin sets the minimum backoff duration.
func WithMin(min time.Duration) WithBackoffConfigurator {
    return func(c *retryWithBackoffConfig) {
        c.min = min
    }
}

// WithMax sets the maximum backoff duration.
func WithMax(max time.Duration) WithBackoffConfigurator {
    return func(c *retryWithBackoffConfig) {
        c.max = max
    }
}

// WithMaxAttempts sets the maximum number of retry attempts.
func WithMaxAttempts(maxAttempts int) WithBackoffConfigurator {
    return func(c *retryWithBackoffConfig) {
        c.maxAttempts = maxAttempts
    }
}

// WithMaxAttemptTime sets the maximum time of all retry attempts.
func WithMaxAttemptTime(maxAttemptTime time.Duration) WithBackoffConfigurator {
    return func(c *retryWithBackoffConfig) {
        c.maxAttemptTime = maxAttemptTime
    }
}

// WithMaxTotalTime sets the maximum time of all retry attempts combined.
func WithMaxTotalTime(maxTotalTime time.Duration) WithBackoffConfigurator {
    return func(c *retryWithBackoffConfig) {
        c.maxAllAttemptsTime = maxTotalTime
    }
}

func defaultConfig() retryWithBackoffConfig {
    return retryWithBackoffConfig{
        factor:             2,
        jitter:             true,
        min:                200 * time.Millisecond,
        max:                5 * time.Second,
        maxAttempts:        -1,
        maxAllAttemptsTime: time.Second * 30,
    }
}

// WithBackoff retries the given function with exponential backoff.
func WithBackoff(ctx context.Context, doFunc RetryableFunc, configurators ...WithBackoffConfigurator) error {
    config := defaultConfig()

    for _, configurator := range configurators {
        configurator(&config)
    }

    b := &backoff.Backoff{
        Factor: config.factor,
        Jitter: config.jitter,
        Min:    config.min,
        Max:    config.max,
    }

    timeout := time.Duration(0)
    startTime := time.Now()

    var errs []error

    attempts := 0
    for !config.exceedsMaxAttempts(attempts) && !config.exceedsMaxTime(startTime) {
        select {
        case <-ctx.Done():
            return fmt.Errorf("%w while retrying", ctx.Err())
        case <-time.After(timeout):
            var funcCtx context.Context
            var cancel context.CancelFunc

            if config.maxAttemptTime > 0 {
                funcCtx, cancel = context.WithTimeout(ctx, config.maxAttemptTime)
            } else {
                funcCtx, cancel = context.WithCancel(ctx)
            }

            err := doFunc(funcCtx)
            if err != nil {
                errs = append(errs, err)
                timeout = b.Duration()
                attempts++
                cancel()
            } else {
                cancel()
                return nil
            }
        }
    }

    if config.exceedsMaxAttempts(attempts) {
        // nolint: wrapcheck
        return multierr.Append(ErrMaxAttempts, fmt.Errorf("after %d attempts (attempt errors: %w)", attempts, multierr.Combine(errs...)))
    }
    if config.exceedsMaxTime(startTime) {
        // nolint: wrapcheck
        return multierr.Append(ErrMaxTime, fmt.Errorf("after %s (max was %s) %w", time.Since(startTime).String(), config.maxAllAttemptsTime.String(), multierr.Combine(errs...)))
    }

    return ErrUnknown
}

// ErrMaxAttempts is returned when the maximum number of retry attempts is reached.
var ErrMaxAttempts = errors.New("max attempts reached")

// ErrMaxTime is returned when the maximum time for all retry attempts is reached.
var ErrMaxTime = errors.New("max time reached")

// ErrUnknown is returned when an unknown error occurs.
var ErrUnknown = errors.New("unknown error")