From bd976642f6bab41529c3a00f99a884f73920ce7f Mon Sep 17 00:00:00 2001
From: dudaodong <lanliddd.2007@163.com>
Date: Thu, 13 Jan 2022 16:18:49 +0800
Subject: [PATCH] feat: add try package for executing a function repeatedly

---
 retry/retry.go      | 89 +++++++++++++++++++++++++++++++++++++++++++++
 retry/retry_test.go | 80 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 169 insertions(+)
 create mode 100644 retry/retry.go
 create mode 100644 retry/retry_test.go

diff --git a/retry/retry.go b/retry/retry.go
new file mode 100644
index 0000000..f58c4fe
--- /dev/null
+++ b/retry/retry.go
@@ -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)
+}
diff --git a/retry/retry_test.go b/retry/retry_test.go
new file mode 100644
index 0000000..f2f0815
--- /dev/null
+++ b/retry/retry_test.go
@@ -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)
+}