bbgo_origin/pkg/util/retry.go

70 lines
1.3 KiB
Go
Raw Normal View History

2020-06-08 02:42:50 +00:00
package util
import (
"context"
"time"
"github.com/pkg/errors"
)
const (
InfiniteRetry = 0
)
type RetryPredicator func(e error) bool
2020-10-17 16:34:26 +00:00
// Retry retrys the passed function for "attempts" times, if passed function return error. Setting attempts to zero means keep retrying.
2020-06-08 02:42:50 +00:00
func Retry(ctx context.Context, attempts int, duration time.Duration, fnToRetry func() error, errHandler func(error), predicators ...RetryPredicator) (err error) {
infinite := false
if attempts == InfiniteRetry {
infinite = true
}
for attempts > 0 || infinite {
select {
case <-ctx.Done():
errMsg := "return for context done"
if err != nil {
return errors.Wrap(err, errMsg)
} else {
return errors.New(errMsg)
}
default:
if err = fnToRetry(); err == nil {
return nil
}
if !needRetry(err, predicators) {
return err
}
err = errors.Wrapf(err, "failed in retry: countdown: %v", attempts)
if errHandler != nil {
errHandler(err)
}
if !infinite {
attempts--
}
time.Sleep(duration)
}
}
return err
}
func needRetry(err error, predicators []RetryPredicator) bool {
if err == nil {
return false
}
// If no predicators specified, we will retry for all errors
if len(predicators) == 0 {
return true
}
return predicators[0](err)
}