retry.go

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")