update: expand the retry method to support timeout and interval control (#3283)

core/fx/retry.go

@@ -1,31 +1,87 @@
 package fx
 
-import "github.com/zeromicro/go-zero/core/errorx"
+import (
+	"context"
+	"errors"
+	"time"
+
+	"github.com/zeromicro/go-zero/core/errorx"
+)
 
 const defaultRetryTimes = 3
 
+var errTimeout = errors.New("retry timeout")
+
 type (
 	// RetryOption defines the method to customize DoWithRetry.
 	RetryOption func(*retryOptions)
 
 	retryOptions struct {
-		times int
+		times    int
+		interval time.Duration
+		timeout  time.Duration
 	}
 )
 
 // DoWithRetry runs fn, and retries if failed. Default to retry 3 times.
+// Note that if the fn function accesses global variables outside the function and performs modification operations,
+// it is best to lock them, otherwise there may be data race issues
 func DoWithRetry(fn func() error, opts ...RetryOption) error {
+	return retry(fn, opts...)
+}
+
+// DoWithRetryCtx runs fn, and retries if failed. Default to retry 3 times.
+// fn retryCount indicates the current number of retries,starting from 0
+// Note that if the fn function accesses global variables outside the function and performs modification operations,
+// it is best to lock them, otherwise there may be data race issues
+func DoWithRetryCtx(fn func(ctx context.Context, retryCount int) error, opts ...RetryOption) error {
+	return retry(fn, opts...)
+}
+
+func retry(fn interface{}, opts ...RetryOption) error {
 	options := newRetryOptions()
 	for _, opt := range opts {
 		opt(options)
 	}
 
+	sign := make(chan error, 1)
 	var berr errorx.BatchError
+	var cancelFunc context.CancelFunc
+	ctx := context.Background()
+	if options.timeout > 0 {
+		ctx, cancelFunc = context.WithTimeout(ctx, options.timeout)
+		defer cancelFunc()
+	}
+
 	for i := 0; i < options.times; i++ {
-		if err := fn(); err != nil {
-			berr.Add(err)
-		} else {
-			return nil
+		go func(retryCount int) {
+			switch f := fn.(type) {
+			case func() error:
+				sign <- f()
+			case func(ctx context.Context, retryCount int) error:
+				sign <- f(ctx, retryCount)
+			}
+		}(i)
+
+		select {
+		case err := <-sign:
+			if err != nil {
+				berr.Add(err)
+			} else {
+				return nil
+			}
+		case <-ctx.Done():
+			berr.Add(errTimeout)
+			return berr.Err()
+		}
+
+		if options.interval > 0 {
+			select {
+			case <-ctx.Done():
+				berr.Add(errTimeout)
+				return berr.Err()
+			case <-time.After(options.interval):
+			}
 		}
 	}
 
@@ -39,8 +95,22 @@ func WithRetry(times int) RetryOption {
 	}
 }
 
-func newRetryOptions() *retryOptions {
-	return &retryOptions{
-		times: defaultRetryTimes,
+func WithInterval(interval time.Duration) RetryOption {
+	return func(options *retryOptions) {
+		options.interval = interval
+	}
+}
+
+func WithTimeout(timeout time.Duration) RetryOption {
+	return func(options *retryOptions) {
+		options.timeout = timeout
+	}
+}
+
+func newRetryOptions() *retryOptions {
+	return &retryOptions{
+		times:    defaultRetryTimes,
+		interval: 0,
+		timeout:  0,
 	}
 }