From fd070fec9113d1b0919d8eccc3c443925a8bae27 Mon Sep 17 00:00:00 2001
From: Kevin Wan <wanjunfeng@gmail.com>
Date: Sun, 15 Oct 2023 21:39:44 +0800
Subject: [PATCH] feat: retry with ctx deadline (#3626)

---
 .codecov.yml          |  4 +++
 core/fx/retry.go      | 14 ++++-------
 core/fx/retry_test.go | 58 +++++++++++++++++++++++++++++++++----------
 3 files changed, 54 insertions(+), 22 deletions(-)

diff --git a/.codecov.yml b/.codecov.yml
index 4a78255b..ea0a4218 100644
--- a/.codecov.yml
+++ b/.codecov.yml
@@ -1,3 +1,7 @@
+coverage:
+  status:
+    patch: true
+    project: false # disabled because project coverage is not stable
 comment:
   layout: "flags, files"
   behavior: once
diff --git a/core/fx/retry.go b/core/fx/retry.go
index 4c0c0376..e79c9002 100644
--- a/core/fx/retry.go
+++ b/core/fx/retry.go
@@ -2,7 +2,6 @@ package fx
 
 import (
 	"context"
-	"errors"
 	"time"
 
 	"github.com/zeromicro/go-zero/core/errorx"
@@ -10,8 +9,6 @@ import (
 
 const defaultRetryTimes = 3
 
-var errTimeout = errors.New("retry timeout")
-
 type (
 	// RetryOption defines the method to customize DoWithRetry.
 	RetryOption func(*retryOptions)
@@ -28,7 +25,7 @@ type (
 // 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(func(errChan chan error, retryCount int) {
+	return retry(context.Background(), func(errChan chan error, retryCount int) {
 		errChan <- fn()
 	}, opts...)
 }
@@ -40,12 +37,12 @@ func DoWithRetry(fn func() error, opts ...RetryOption) error {
 // otherwise there may be data race issues
 func DoWithRetryCtx(ctx context.Context, fn func(ctx context.Context, retryCount int) error,
 	opts ...RetryOption) error {
-	return retry(func(errChan chan error, retryCount int) {
+	return retry(ctx, func(errChan chan error, retryCount int) {
 		errChan <- fn(ctx, retryCount)
 	}, opts...)
 }
 
-func retry(fn func(errChan chan error, retryCount int), opts ...RetryOption) error {
+func retry(ctx context.Context, fn func(errChan chan error, retryCount int), opts ...RetryOption) error {
 	options := newRetryOptions()
 	for _, opt := range opts {
 		opt(options)
@@ -53,7 +50,6 @@ func retry(fn func(errChan chan error, retryCount int), opts ...RetryOption) err
 
 	var berr errorx.BatchError
 	var cancelFunc context.CancelFunc
-	ctx := context.Background()
 	if options.timeout > 0 {
 		ctx, cancelFunc = context.WithTimeout(ctx, options.timeout)
 		defer cancelFunc()
@@ -71,14 +67,14 @@ func retry(fn func(errChan chan error, retryCount int), opts ...RetryOption) err
 				return nil
 			}
 		case <-ctx.Done():
-			berr.Add(errTimeout)
+			berr.Add(ctx.Err())
 			return berr.Err()
 		}
 
 		if options.interval > 0 {
 			select {
 			case <-ctx.Done():
-				berr.Add(errTimeout)
+				berr.Add(ctx.Err())
 				return berr.Err()
 			case <-time.After(options.interval):
 			}
diff --git a/core/fx/retry_test.go b/core/fx/retry_test.go
index d4569dc4..045d782a 100644
--- a/core/fx/retry_test.go
+++ b/core/fx/retry_test.go
@@ -98,19 +98,51 @@ func TestRetryWithInterval(t *testing.T) {
 }
 
 func TestRetryCtx(t *testing.T) {
-	assert.NotNil(t, DoWithRetryCtx(context.Background(), func(ctx context.Context, retryCount int) error {
-		if retryCount == 0 {
+	t.Run("with timeout", func(t *testing.T) {
+		assert.NotNil(t, DoWithRetryCtx(context.Background(), func(ctx context.Context, retryCount int) error {
+			if retryCount == 0 {
+				return errors.New("any")
+			}
+			time.Sleep(time.Millisecond * 150)
+			return nil
+		}, WithTimeout(time.Millisecond*250), WithInterval(time.Millisecond*150)))
+
+		assert.NotNil(t, DoWithRetryCtx(context.Background(), func(ctx context.Context, retryCount int) error {
+			if retryCount == 1 {
+				return nil
+			}
+			time.Sleep(time.Millisecond * 150)
+			return errors.New("any ")
+		}, WithTimeout(time.Millisecond*250), WithInterval(time.Millisecond*150)))
+	})
+
+	t.Run("with deadline exceeded", func(t *testing.T) {
+		ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(time.Millisecond*250))
+		defer cancel()
+
+		var times int
+		assert.Error(t, DoWithRetryCtx(ctx, func(ctx context.Context, retryCount int) error {
+			times++
+			time.Sleep(time.Millisecond * 150)
 			return errors.New("any")
-		}
-		time.Sleep(time.Millisecond * 150)
-		return nil
-	}, WithTimeout(time.Millisecond*250), WithInterval(time.Millisecond*150)))
+		}, WithInterval(time.Millisecond*150)))
+		assert.Equal(t, 1, times)
+	})
 
-	assert.NotNil(t, DoWithRetryCtx(context.Background(), func(ctx context.Context, retryCount int) error {
-		if retryCount == 1 {
-			return nil
-		}
-		time.Sleep(time.Millisecond * 150)
-		return errors.New("any ")
-	}, WithTimeout(time.Millisecond*250), WithInterval(time.Millisecond*150)))
+	t.Run("with deadline not exceeded", func(t *testing.T) {
+		ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(time.Millisecond*250))
+		defer cancel()
+
+		var times int
+		assert.NoError(t, DoWithRetryCtx(ctx, func(ctx context.Context, retryCount int) error {
+			times++
+			if times == defaultRetryTimes {
+				return nil
+			}
+
+			time.Sleep(time.Millisecond * 50)
+			return errors.New("any")
+		}))
+		assert.Equal(t, defaultRetryTimes, times)
+	})
 }