-
Notifications
You must be signed in to change notification settings - Fork 24
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
2 changed files
with
197 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,107 @@ | ||
package retryexecutor | ||
|
||
import ( | ||
"context" | ||
"errors" | ||
"fmt" | ||
"github.com/jfrog/gofrog/log" | ||
"time" | ||
) | ||
|
||
type ExecutionHandlerFunc func() (shouldRetry bool, err error) | ||
|
||
type RetryExecutor struct { | ||
// The context | ||
Context context.Context | ||
|
||
// The amount of retries to perform. | ||
MaxRetries int | ||
|
||
// Number of milliseconds to sleep between retries. | ||
RetriesIntervalMilliSecs int | ||
|
||
// Message to display when retrying. | ||
ErrorMessage string | ||
|
||
// Prefix to print at the beginning of each log. | ||
LogMsgPrefix string | ||
|
||
// ExecutionHandler is the operation to run with retries. | ||
ExecutionHandler ExecutionHandlerFunc | ||
} | ||
|
||
func (runner *RetryExecutor) Execute() error { | ||
var err error | ||
var shouldRetry bool | ||
for i := 0; i <= runner.MaxRetries; i++ { | ||
// Run ExecutionHandler | ||
shouldRetry, err = runner.ExecutionHandler() | ||
|
||
// If we should not retry, return. | ||
if !shouldRetry { | ||
return err | ||
} | ||
if cancelledErr := runner.checkCancelled(); cancelledErr != nil { | ||
return cancelledErr | ||
} | ||
|
||
// Print retry log message | ||
runner.LogRetry(i, err) | ||
|
||
// Going to sleep for RetryInterval milliseconds | ||
if runner.RetriesIntervalMilliSecs > 0 && i < runner.MaxRetries { | ||
time.Sleep(time.Millisecond * time.Duration(runner.RetriesIntervalMilliSecs)) | ||
} | ||
} | ||
// If the error is not nil, return it and log the timeout message. Otherwise, generate new error. | ||
if err != nil { | ||
log.Info(runner.getTimeoutErrorMsg()) | ||
return err | ||
} | ||
return TimeoutError{runner.getTimeoutErrorMsg()} | ||
} | ||
|
||
// Error of this type will be returned if the executor reaches timeout and no other error is returned by the execution handler. | ||
type TimeoutError struct { | ||
errMsg string | ||
} | ||
|
||
func (retryErr TimeoutError) Error() string { | ||
return retryErr.errMsg | ||
} | ||
|
||
func (runner *RetryExecutor) getTimeoutErrorMsg() string { | ||
prefix := "" | ||
if runner.LogMsgPrefix != "" { | ||
prefix = runner.LogMsgPrefix + " " | ||
} | ||
return fmt.Sprintf("%sexecutor timeout after %v attempts with %v milliseconds wait intervals", prefix, runner.MaxRetries, runner.RetriesIntervalMilliSecs) | ||
} | ||
|
||
func (runner *RetryExecutor) LogRetry(attemptNumber int, err error) { | ||
message := fmt.Sprintf("%s(Attempt %v)", runner.LogMsgPrefix, attemptNumber+1) | ||
if runner.ErrorMessage != "" { | ||
message = fmt.Sprintf("%s - %s", message, runner.ErrorMessage) | ||
} | ||
if err != nil { | ||
message = fmt.Sprintf("%s: %s", message, err.Error()) | ||
} | ||
|
||
if err != nil || runner.ErrorMessage != "" { | ||
log.Warn(message) | ||
} else { | ||
log.Debug(message) | ||
} | ||
} | ||
|
||
func (runner *RetryExecutor) checkCancelled() error { | ||
if runner.Context == nil { | ||
return nil | ||
} | ||
contextErr := runner.Context.Err() | ||
if errors.Is(contextErr, context.Canceled) { | ||
log.Info("Retry executor was cancelled") | ||
return contextErr | ||
} | ||
return nil | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,90 @@ | ||
package retryexecutor | ||
|
||
import ( | ||
"context" | ||
"errors" | ||
"github.com/jfrog/gofrog/log" | ||
"github.com/stretchr/testify/assert" | ||
"testing" | ||
) | ||
|
||
func TestRetryExecutorSuccess(t *testing.T) { | ||
retriesToPerform := 10 | ||
breakRetriesAt := 4 | ||
runCount := 0 | ||
executor := RetryExecutor{ | ||
MaxRetries: retriesToPerform, | ||
RetriesIntervalMilliSecs: 0, | ||
ErrorMessage: "Testing RetryExecutor", | ||
ExecutionHandler: func() (bool, error) { | ||
runCount++ | ||
if runCount == breakRetriesAt { | ||
log.Warn("Breaking after", runCount-1, "retries") | ||
return false, nil | ||
} | ||
return true, nil | ||
}, | ||
} | ||
|
||
assert.NoError(t, executor.Execute()) | ||
assert.Equal(t, breakRetriesAt, runCount) | ||
} | ||
|
||
func TestRetryExecutorTimeoutWithDefaultError(t *testing.T) { | ||
retriesToPerform := 5 | ||
runCount := 0 | ||
|
||
executor := RetryExecutor{ | ||
MaxRetries: retriesToPerform, | ||
RetriesIntervalMilliSecs: 0, | ||
ErrorMessage: "Testing RetryExecutor", | ||
ExecutionHandler: func() (bool, error) { | ||
runCount++ | ||
return true, nil | ||
}, | ||
} | ||
|
||
assert.Equal(t, executor.Execute(), TimeoutError{executor.getTimeoutErrorMsg()}) | ||
assert.Equal(t, retriesToPerform+1, runCount) | ||
} | ||
|
||
func TestRetryExecutorTimeoutWithCustomError(t *testing.T) { | ||
retriesToPerform := 5 | ||
runCount := 0 | ||
|
||
executionHandler := errors.New("retry failed due to reason") | ||
|
||
executor := RetryExecutor{ | ||
MaxRetries: retriesToPerform, | ||
RetriesIntervalMilliSecs: 0, | ||
ErrorMessage: "Testing RetryExecutor", | ||
ExecutionHandler: func() (bool, error) { | ||
runCount++ | ||
return true, executionHandler | ||
}, | ||
} | ||
|
||
assert.Equal(t, executor.Execute(), executionHandler) | ||
assert.Equal(t, retriesToPerform+1, runCount) | ||
} | ||
|
||
func TestRetryExecutorCancel(t *testing.T) { | ||
retriesToPerform := 5 | ||
runCount := 0 | ||
|
||
retryContext, cancelFunc := context.WithCancel(context.Background()) | ||
executor := RetryExecutor{ | ||
Context: retryContext, | ||
MaxRetries: retriesToPerform, | ||
RetriesIntervalMilliSecs: 0, | ||
ErrorMessage: "Testing RetryExecutor", | ||
ExecutionHandler: func() (bool, error) { | ||
runCount++ | ||
return true, nil | ||
}, | ||
} | ||
|
||
cancelFunc() | ||
assert.EqualError(t, executor.Execute(), context.Canceled.Error()) | ||
assert.Equal(t, 1, runCount) | ||
} |