github.com/wfusion/gofusion@v1.1.14/common/infra/watermill/message/router/middleware/retry.go (about)

     1  package middleware
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	"github.com/cenkalti/backoff/v4"
     8  
     9  	"github.com/wfusion/gofusion/common/infra/watermill"
    10  	"github.com/wfusion/gofusion/common/infra/watermill/message"
    11  )
    12  
    13  // Retry provides a middleware that retries the handler if errors are returned.
    14  // The retry behaviour is configurable, with exponential backoff and maximum elapsed time.
    15  type Retry struct {
    16  	// MaxRetries is maximum number of times a retry will be attempted.
    17  	MaxRetries int
    18  
    19  	// InitialInterval is the first interval between retries. Subsequent intervals will be scaled by Multiplier.
    20  	InitialInterval time.Duration
    21  	// MaxInterval sets the limit for the exponential backoff of retries. The interval will not be increased beyond MaxInterval.
    22  	MaxInterval time.Duration
    23  	// Multiplier is the factor by which the waiting interval will be multiplied between retries.
    24  	Multiplier float64
    25  	// MaxElapsedTime sets the time limit of how long retries will be attempted. Disabled if 0.
    26  	MaxElapsedTime time.Duration
    27  	// RandomizationFactor randomizes the spread of the backoff times within the interval of:
    28  	// [currentInterval * (1 - randomization_factor), currentInterval * (1 + randomization_factor)].
    29  	RandomizationFactor float64
    30  
    31  	// OnRetryHook is an optional function that will be executed on each retry attempt.
    32  	// The number of the current retry is passed as retryNum,
    33  	OnRetryHook func(retryNum int, delay time.Duration)
    34  
    35  	Logger watermill.LoggerAdapter
    36  }
    37  
    38  // Middleware returns the Retry middleware.
    39  func (r Retry) Middleware(h message.HandlerFunc) message.HandlerFunc {
    40  	return func(msg *message.Message) ([]*message.Message, error) {
    41  		producedMessages, err := h(msg)
    42  		if err == nil {
    43  			return producedMessages, nil
    44  		}
    45  
    46  		expBackoff := backoff.NewExponentialBackOff()
    47  		expBackoff.InitialInterval = r.InitialInterval
    48  		expBackoff.MaxInterval = r.MaxInterval
    49  		expBackoff.Multiplier = r.Multiplier
    50  		expBackoff.MaxElapsedTime = r.MaxElapsedTime
    51  		expBackoff.RandomizationFactor = r.RandomizationFactor
    52  
    53  		ctx := msg.Context()
    54  		if r.MaxElapsedTime > 0 {
    55  			var cancel func()
    56  			ctx, cancel = context.WithTimeout(ctx, r.MaxElapsedTime)
    57  			defer cancel()
    58  		}
    59  
    60  		retryNum := 1
    61  		expBackoff.Reset()
    62  	retryLoop:
    63  		for {
    64  			waitTime := expBackoff.NextBackOff()
    65  			select {
    66  			case <-ctx.Done():
    67  				return producedMessages, err
    68  			case <-time.After(waitTime):
    69  				// go on
    70  			}
    71  
    72  			producedMessages, err = h(msg)
    73  			if err == nil {
    74  				return producedMessages, nil
    75  			}
    76  
    77  			if r.Logger != nil {
    78  				r.Logger.Error("Error occurred, retrying", err, watermill.LogFields{
    79  					"retry_no":     retryNum,
    80  					"max_retries":  r.MaxRetries,
    81  					"wait_time":    waitTime,
    82  					"elapsed_time": expBackoff.GetElapsedTime(),
    83  				})
    84  			}
    85  			if r.OnRetryHook != nil {
    86  				r.OnRetryHook(retryNum, waitTime)
    87  			}
    88  
    89  			retryNum++
    90  			if retryNum > r.MaxRetries {
    91  				break retryLoop
    92  			}
    93  		}
    94  
    95  		return nil, err
    96  	}
    97  }