retry.go raw

   1  package backoff
   2  
   3  import (
   4  	"context"
   5  	"errors"
   6  	"time"
   7  )
   8  
   9  // DefaultMaxElapsedTime sets a default limit for the total retry duration.
  10  const DefaultMaxElapsedTime = 15 * time.Minute
  11  
  12  // Operation is a function that attempts an operation and may be retried.
  13  type Operation[T any] func() (T, error)
  14  
  15  // Notify is a function called on operation error with the error and backoff duration.
  16  type Notify func(error, time.Duration)
  17  
  18  // retryOptions holds configuration settings for the retry mechanism.
  19  type retryOptions struct {
  20  	BackOff        BackOff       // Strategy for calculating backoff periods.
  21  	Timer          timer         // Timer to manage retry delays.
  22  	Notify         Notify        // Optional function to notify on each retry error.
  23  	MaxTries       uint          // Maximum number of retry attempts.
  24  	MaxElapsedTime time.Duration // Maximum total time for all retries.
  25  }
  26  
  27  type RetryOption func(*retryOptions)
  28  
  29  // WithBackOff configures a custom backoff strategy.
  30  func WithBackOff(b BackOff) RetryOption {
  31  	return func(args *retryOptions) {
  32  		args.BackOff = b
  33  	}
  34  }
  35  
  36  // withTimer sets a custom timer for managing delays between retries.
  37  func withTimer(t timer) RetryOption {
  38  	return func(args *retryOptions) {
  39  		args.Timer = t
  40  	}
  41  }
  42  
  43  // WithNotify sets a notification function to handle retry errors.
  44  func WithNotify(n Notify) RetryOption {
  45  	return func(args *retryOptions) {
  46  		args.Notify = n
  47  	}
  48  }
  49  
  50  // WithMaxTries limits the number of all attempts.
  51  func WithMaxTries(n uint) RetryOption {
  52  	return func(args *retryOptions) {
  53  		args.MaxTries = n
  54  	}
  55  }
  56  
  57  // WithMaxElapsedTime limits the total duration for retry attempts.
  58  func WithMaxElapsedTime(d time.Duration) RetryOption {
  59  	return func(args *retryOptions) {
  60  		args.MaxElapsedTime = d
  61  	}
  62  }
  63  
  64  // Retry attempts the operation until success, a permanent error, or backoff completion.
  65  // It ensures the operation is executed at least once.
  66  //
  67  // Returns the operation result or error if retries are exhausted or context is cancelled.
  68  func Retry[T any](ctx context.Context, operation Operation[T], opts ...RetryOption) (T, error) {
  69  	// Initialize default retry options.
  70  	args := &retryOptions{
  71  		BackOff:        NewExponentialBackOff(),
  72  		Timer:          &defaultTimer{},
  73  		MaxElapsedTime: DefaultMaxElapsedTime,
  74  	}
  75  
  76  	// Apply user-provided options to the default settings.
  77  	for _, opt := range opts {
  78  		opt(args)
  79  	}
  80  
  81  	defer args.Timer.Stop()
  82  
  83  	startedAt := time.Now()
  84  	args.BackOff.Reset()
  85  	for numTries := uint(1); ; numTries++ {
  86  		// Execute the operation.
  87  		res, err := operation()
  88  		if err == nil {
  89  			return res, nil
  90  		}
  91  
  92  		// Stop retrying if maximum tries exceeded.
  93  		if args.MaxTries > 0 && numTries >= args.MaxTries {
  94  			return res, err
  95  		}
  96  
  97  		// Handle permanent errors without retrying.
  98  		var permanent *PermanentError
  99  		if errors.As(err, &permanent) {
 100  			return res, permanent.Unwrap()
 101  		}
 102  
 103  		// Stop retrying if context is cancelled.
 104  		if cerr := context.Cause(ctx); cerr != nil {
 105  			return res, cerr
 106  		}
 107  
 108  		// Calculate next backoff duration.
 109  		next := args.BackOff.NextBackOff()
 110  		if next == Stop {
 111  			return res, err
 112  		}
 113  
 114  		// Reset backoff if RetryAfterError is encountered.
 115  		var retryAfter *RetryAfterError
 116  		if errors.As(err, &retryAfter) {
 117  			next = retryAfter.Duration
 118  			args.BackOff.Reset()
 119  		}
 120  
 121  		// Stop retrying if maximum elapsed time exceeded.
 122  		if args.MaxElapsedTime > 0 && time.Since(startedAt)+next > args.MaxElapsedTime {
 123  			return res, err
 124  		}
 125  
 126  		// Notify on error if a notifier function is provided.
 127  		if args.Notify != nil {
 128  			args.Notify(err, next)
 129  		}
 130  
 131  		// Wait for the next backoff period or context cancellation.
 132  		args.Timer.Start(next)
 133  		select {
 134  		case <-args.Timer.C():
 135  		case <-ctx.Done():
 136  			return res, context.Cause(ctx)
 137  		}
 138  	}
 139  }
 140