src/pkg/utils/retry.go

package utils

import (
	"context"
	"fmt"
	"time"

	"github.com/golang/glog"
)

type RetryOptions struct {
	// Number of times to try a function before failing.
	// If unset, the function will be attempted once.
	Attempts int
	// The time to wait between function attempts.
	Delay time.Duration
	// Function which determines whether or not the function should be retried after a given error.
	// If this is unset, every error is considered retryable.
	ShouldRetry func(error) bool
}

// Retry calls a function repeatedly until it succeeds or another stopping condition is met.
//
// The stopping conditions are determined by the provided RetryOptions.
func Retry(ctx context.Context, fn func(ctx context.Context) error, opts RetryOptions) error {
	attempts := opts.Attempts
	if attempts <= 0 {
		attempts = 1
	}

	for i := range attempts {
		err := fn(ctx)
		if err == nil {
			return nil
		}

		if i + 1 == attempts {
			break
		}

		if opts.ShouldRetry == nil || opts.ShouldRetry(err) {
			glog.Infof("error %w; retrying", err)
		} else {
			return fmt.Errorf("fatal error, not retrying: %w", err)
		}

		if opts.Delay > 0 {
			// Block until the delay interval or the context deadline, whichever happens first.
			select {
				case <-time.After(opts.Delay): {}
				case <-ctx.Done():
					return fmt.Errorf("context deadline exceeded")
			}
		}
	}

	return fmt.Errorf("no successes after %d attempts", opts.Attempts)
}