mirror of
https://github.com/duke-git/lancet.git
synced 2026-02-04 12:52:28 +08:00
feat: add try package for executing a function repeatedly
This commit is contained in:
89
retry/retry.go
Normal file
89
retry/retry.go
Normal file
@@ -0,0 +1,89 @@
|
||||
// Copyright 2021 dudaodong@gmail.com. All rights reserved.
|
||||
// Use of this source code is governed by MIT license
|
||||
|
||||
// Package retry is for executing a function repeatedly until it was successful or canceled by the context.
|
||||
package retry
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"reflect"
|
||||
"runtime"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
DefaultRetryTimes = 5
|
||||
DefaultRetryDuration = time.Second * 3
|
||||
)
|
||||
|
||||
// RetryConfig is config for retry
|
||||
type RetryConfig struct {
|
||||
context context.Context
|
||||
retryTimes uint
|
||||
retryDuration time.Duration
|
||||
}
|
||||
|
||||
// RetryFn is function that retry executes
|
||||
type RetryFunc func() error
|
||||
|
||||
// Option is for adding retry config
|
||||
type Option func(*RetryConfig)
|
||||
|
||||
// RetryTimes set times of retry
|
||||
func RetryTimes(n uint) Option {
|
||||
return func(rc *RetryConfig) {
|
||||
rc.retryTimes = n
|
||||
}
|
||||
}
|
||||
|
||||
// RetryDuration set duration of retries
|
||||
func RetryDuration(d time.Duration) Option {
|
||||
return func(rc *RetryConfig) {
|
||||
rc.retryDuration = d
|
||||
}
|
||||
}
|
||||
|
||||
// Context set retry context config
|
||||
func Context(ctx context.Context) Option {
|
||||
return func(rc *RetryConfig) {
|
||||
rc.context = ctx
|
||||
}
|
||||
}
|
||||
|
||||
// Retry executes the retryFunc repeatedly until it was successful or canceled by the context
|
||||
// The default times of retries is 5 and the default duration between retries is 3 seconds
|
||||
func Retry(retryFunc RetryFunc, opts ...Option) error {
|
||||
config := &RetryConfig{
|
||||
retryTimes: DefaultRetryTimes,
|
||||
retryDuration: DefaultRetryDuration,
|
||||
context: context.TODO(),
|
||||
}
|
||||
|
||||
for _, opt := range opts {
|
||||
opt(config)
|
||||
}
|
||||
|
||||
var i uint
|
||||
for i < config.retryTimes {
|
||||
err := retryFunc()
|
||||
if err != nil {
|
||||
select {
|
||||
case <-time.After(config.retryDuration):
|
||||
case <-config.context.Done():
|
||||
return errors.New("retry is cancelled")
|
||||
}
|
||||
} else {
|
||||
return nil
|
||||
}
|
||||
i++
|
||||
}
|
||||
|
||||
funcPath := runtime.FuncForPC(reflect.ValueOf(retryFunc).Pointer()).Name()
|
||||
lastSlash := strings.LastIndex(funcPath, "/")
|
||||
funcName := funcPath[lastSlash+1:]
|
||||
|
||||
return fmt.Errorf("function %s run failed after %d times retry", funcName, i)
|
||||
}
|
||||
80
retry/retry_test.go
Normal file
80
retry/retry_test.go
Normal file
@@ -0,0 +1,80 @@
|
||||
package retry
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/duke-git/lancet/internal"
|
||||
)
|
||||
|
||||
func TestRetryFailed(t *testing.T) {
|
||||
assert := internal.NewAssert(t, "TestRetryFailed")
|
||||
|
||||
var number int
|
||||
increaseNumber := func() error {
|
||||
number++
|
||||
return errors.New("error occurs")
|
||||
}
|
||||
|
||||
err := Retry(increaseNumber, RetryDuration(time.Microsecond*50))
|
||||
|
||||
assert.IsNotNil(err)
|
||||
assert.Equal(DefaultRetryTimes, number)
|
||||
}
|
||||
|
||||
func TestRetrySucceeded(t *testing.T) {
|
||||
assert := internal.NewAssert(t, "TestRetrySucceeded")
|
||||
|
||||
var number int
|
||||
increaseNumber := func() error {
|
||||
number++
|
||||
if number == DefaultRetryTimes {
|
||||
return nil
|
||||
}
|
||||
return errors.New("error occurs")
|
||||
}
|
||||
|
||||
err := Retry(increaseNumber, RetryDuration(time.Microsecond*50))
|
||||
|
||||
assert.IsNil(err)
|
||||
assert.Equal(DefaultRetryTimes, number)
|
||||
}
|
||||
|
||||
func TestSetRetryTimes(t *testing.T) {
|
||||
assert := internal.NewAssert(t, "TestSetRetryTimes")
|
||||
|
||||
var number int
|
||||
increaseNumber := func() error {
|
||||
number++
|
||||
return errors.New("error occurs")
|
||||
}
|
||||
|
||||
err := Retry(increaseNumber, RetryDuration(time.Microsecond*50), RetryTimes(3))
|
||||
|
||||
assert.IsNotNil(err)
|
||||
assert.Equal(3, number)
|
||||
}
|
||||
|
||||
func TestCancelRetry(t *testing.T) {
|
||||
assert := internal.NewAssert(t, "TestCancelRetry")
|
||||
|
||||
ctx, cancel := context.WithCancel(context.TODO())
|
||||
var number int
|
||||
increaseNumber := func() error {
|
||||
number++
|
||||
if number > 3 {
|
||||
cancel()
|
||||
}
|
||||
return errors.New("error occurs")
|
||||
}
|
||||
|
||||
err := Retry(increaseNumber,
|
||||
RetryDuration(time.Microsecond*50),
|
||||
Context(ctx),
|
||||
)
|
||||
|
||||
assert.IsNotNil(err)
|
||||
assert.Equal(4, number)
|
||||
}
|
||||
Reference in New Issue
Block a user