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 }