github.com/swisspost/terratest@v0.0.0-20230214120104-7ec6de2e1ae0/modules/retry/retry.go (about)

     1  // Package retry contains logic to retry actions with certain conditions.
     2  package retry
     3  
     4  import (
     5  	"fmt"
     6  	"regexp"
     7  	"time"
     8  
     9  	"github.com/stretchr/testify/require"
    10  
    11  	"github.com/gruntwork-io/terratest/modules/logger"
    12  	"github.com/gruntwork-io/terratest/modules/testing"
    13  	"golang.org/x/net/context"
    14  )
    15  
    16  // Either contains a result and potentially an error.
    17  type Either struct {
    18  	Result string
    19  	Error  error
    20  }
    21  
    22  // DoWithTimeout runs the specified action and waits up to the specified timeout for it to complete. Return the output of the action if
    23  // it completes on time or fail the test otherwise.
    24  func DoWithTimeout(t testing.TestingT, actionDescription string, timeout time.Duration, action func() (string, error)) string {
    25  	out, err := DoWithTimeoutE(t, actionDescription, timeout, action)
    26  	if err != nil {
    27  		t.Fatal(err)
    28  	}
    29  	return out
    30  }
    31  
    32  // DoWithTimeoutE runs the specified action and waits up to the specified timeout for it to complete. Return the output of the action if
    33  // it completes on time or an error otherwise.
    34  func DoWithTimeoutE(t testing.TestingT, actionDescription string, timeout time.Duration, action func() (string, error)) (string, error) {
    35  	ctx, cancel := context.WithTimeout(context.Background(), timeout)
    36  	defer cancel()
    37  
    38  	resultChannel := make(chan Either, 1)
    39  
    40  	go func() {
    41  		out, err := action()
    42  		resultChannel <- Either{Result: out, Error: err}
    43  	}()
    44  
    45  	select {
    46  	case either := <-resultChannel:
    47  		return either.Result, either.Error
    48  	case <-ctx.Done():
    49  		return "", TimeoutExceeded{Description: actionDescription, Timeout: timeout}
    50  	}
    51  }
    52  
    53  // DoWithRetry runs the specified action. If it returns a string, return that string. If it returns a FatalError, return that error
    54  // immediately. If it returns any other type of error, sleep for sleepBetweenRetries and try again, up to a maximum of
    55  // maxRetries retries. If maxRetries is exceeded, fail the test.
    56  func DoWithRetry(t testing.TestingT, actionDescription string, maxRetries int, sleepBetweenRetries time.Duration, action func() (string, error)) string {
    57  	out, err := DoWithRetryE(t, actionDescription, maxRetries, sleepBetweenRetries, action)
    58  	if err != nil {
    59  		t.Fatal(err)
    60  	}
    61  	return out
    62  }
    63  
    64  // DoWithRetryE runs the specified action. If it returns a string, return that string. If it returns a FatalError, return that error
    65  // immediately. If it returns any other type of error, sleep for sleepBetweenRetries and try again, up to a maximum of
    66  // maxRetries retries. If maxRetries is exceeded, return a MaxRetriesExceeded error.
    67  func DoWithRetryE(t testing.TestingT, actionDescription string, maxRetries int, sleepBetweenRetries time.Duration, action func() (string, error)) (string, error) {
    68  	out, err := DoWithRetryInterfaceE(t, actionDescription, maxRetries, sleepBetweenRetries, func() (interface{}, error) { return action() })
    69  	return out.(string), err
    70  }
    71  
    72  // DoWithRetryInterface runs the specified action. If it returns a value, return that value. If it returns a FatalError, return that error
    73  // immediately. If it returns any other type of error, sleep for sleepBetweenRetries and try again, up to a maximum of
    74  // maxRetries retries. If maxRetries is exceeded, fail the test.
    75  func DoWithRetryInterface(t testing.TestingT, actionDescription string, maxRetries int, sleepBetweenRetries time.Duration, action func() (interface{}, error)) interface{} {
    76  	out, err := DoWithRetryInterfaceE(t, actionDescription, maxRetries, sleepBetweenRetries, action)
    77  	if err != nil {
    78  		t.Fatal(err)
    79  	}
    80  	return out
    81  }
    82  
    83  // DoWithRetryInterfaceE runs the specified action. If it returns a value, return that value. If it returns a FatalError, return that error
    84  // immediately. If it returns any other type of error, sleep for sleepBetweenRetries and try again, up to a maximum of
    85  // maxRetries retries. If maxRetries is exceeded, return a MaxRetriesExceeded error.
    86  func DoWithRetryInterfaceE(t testing.TestingT, actionDescription string, maxRetries int, sleepBetweenRetries time.Duration, action func() (interface{}, error)) (interface{}, error) {
    87  	var output interface{}
    88  	var err error
    89  
    90  	for i := 0; i <= maxRetries; i++ {
    91  		logger.Log(t, actionDescription)
    92  
    93  		output, err = action()
    94  		if err == nil {
    95  			return output, nil
    96  		}
    97  
    98  		if _, isFatalErr := err.(FatalError); isFatalErr {
    99  			logger.Logf(t, "Returning due to fatal error: %v", err)
   100  			return output, err
   101  		}
   102  
   103  		logger.Logf(t, "%s returned an error: %s. Sleeping for %s and will try again.", actionDescription, err.Error(), sleepBetweenRetries)
   104  		time.Sleep(sleepBetweenRetries)
   105  	}
   106  
   107  	return output, MaxRetriesExceeded{Description: actionDescription, MaxRetries: maxRetries}
   108  }
   109  
   110  // DoWithRetryableErrors runs the specified action. If it returns a value, return that value. If it returns an error,
   111  // check if error message or the string output from the action (which is often stdout/stderr from running some command)
   112  // matches any of the regular expressions in the specified retryableErrors map. If there is a match, sleep for
   113  // sleepBetweenRetries, and retry the specified action, up to a maximum of maxRetries retries. If there is no match,
   114  // return that error immediately, wrapped in a FatalError. If maxRetries is exceeded, return a MaxRetriesExceeded error.
   115  func DoWithRetryableErrors(t testing.TestingT, actionDescription string, retryableErrors map[string]string, maxRetries int, sleepBetweenRetries time.Duration, action func() (string, error)) string {
   116  	out, err := DoWithRetryableErrorsE(t, actionDescription, retryableErrors, maxRetries, sleepBetweenRetries, action)
   117  	require.NoError(t, err)
   118  	return out
   119  }
   120  
   121  // DoWithRetryableErrorsE runs the specified action. If it returns a value, return that value. If it returns an error,
   122  // check if error message or the string output from the action (which is often stdout/stderr from running some command)
   123  // matches any of the regular expressions in the specified retryableErrors map. If there is a match, sleep for
   124  // sleepBetweenRetries, and retry the specified action, up to a maximum of maxRetries retries. If there is no match,
   125  // return that error immediately, wrapped in a FatalError. If maxRetries is exceeded, return a MaxRetriesExceeded error.
   126  func DoWithRetryableErrorsE(t testing.TestingT, actionDescription string, retryableErrors map[string]string, maxRetries int, sleepBetweenRetries time.Duration, action func() (string, error)) (string, error) {
   127  	retryableErrorsRegexp := map[*regexp.Regexp]string{}
   128  	for errorStr, errorMessage := range retryableErrors {
   129  		errorRegex, err := regexp.Compile(errorStr)
   130  		if err != nil {
   131  			return "", FatalError{Underlying: err}
   132  		}
   133  		retryableErrorsRegexp[errorRegex] = errorMessage
   134  	}
   135  
   136  	return DoWithRetryE(t, actionDescription, maxRetries, sleepBetweenRetries, func() (string, error) {
   137  		output, err := action()
   138  		if err == nil {
   139  			return output, nil
   140  		}
   141  
   142  		for errorRegexp, errorMessage := range retryableErrorsRegexp {
   143  			if errorRegexp.MatchString(output) || errorRegexp.MatchString(err.Error()) {
   144  				logger.Logf(t, "'%s' failed with the error '%s' but this error was expected and warrants a retry. Further details: %s\n", actionDescription, err.Error(), errorMessage)
   145  				return output, err
   146  			}
   147  		}
   148  
   149  		return output, FatalError{Underlying: err}
   150  	})
   151  }
   152  
   153  // Done can be stopped.
   154  type Done struct {
   155  	stop chan bool
   156  }
   157  
   158  // Done stops the execution.
   159  func (done Done) Done() {
   160  	done.stop <- true
   161  }
   162  
   163  // DoInBackgroundUntilStopped runs the specified action in the background (in a goroutine) repeatedly, waiting the specified amount of time between
   164  // repetitions. To stop this action, call the Done() function on the returned value.
   165  func DoInBackgroundUntilStopped(t testing.TestingT, actionDescription string, sleepBetweenRepeats time.Duration, action func()) Done {
   166  	stop := make(chan bool)
   167  
   168  	go func() {
   169  		for {
   170  			logger.Logf(t, "Executing action '%s'", actionDescription)
   171  
   172  			action()
   173  
   174  			logger.Logf(t, "Sleeping for %s before repeating action '%s'", sleepBetweenRepeats, actionDescription)
   175  
   176  			select {
   177  			case <-time.After(sleepBetweenRepeats):
   178  				// Nothing to do, just allow the loop to continue
   179  			case <-stop:
   180  				logger.Logf(t, "Received stop signal for action '%s'.", actionDescription)
   181  				return
   182  			}
   183  		}
   184  	}()
   185  
   186  	return Done{stop: stop}
   187  }
   188  
   189  // Custom error types
   190  
   191  // TimeoutExceeded is an error that occurs when a timeout is exceeded.
   192  type TimeoutExceeded struct {
   193  	Description string
   194  	Timeout     time.Duration
   195  }
   196  
   197  func (err TimeoutExceeded) Error() string {
   198  	return fmt.Sprintf("'%s' did not complete before timeout of %s", err.Description, err.Timeout)
   199  }
   200  
   201  // MaxRetriesExceeded is an error that occurs when the maximum amount of retries is exceeded.
   202  type MaxRetriesExceeded struct {
   203  	Description string
   204  	MaxRetries  int
   205  }
   206  
   207  func (err MaxRetriesExceeded) Error() string {
   208  	return fmt.Sprintf("'%s' unsuccessful after %d retries", err.Description, err.MaxRetries)
   209  }
   210  
   211  // FatalError is a marker interface for errors that should not be retried.
   212  type FatalError struct {
   213  	Underlying error
   214  }
   215  
   216  func (err FatalError) Error() string {
   217  	return fmt.Sprintf("FatalError{Underlying: %v}", err.Underlying)
   218  }