From d952e06f1ecefd5e6d5f6aa6c650aeca3f7d26e0 Mon Sep 17 00:00:00 2001
From: JmPotato <ghzpotato@gmail.com>
Date: Thu, 30 May 2024 11:21:27 +0800
Subject: [PATCH] Refine the retryable checker

Signed-off-by: JmPotato <ghzpotato@gmail.com>
---
 client/http/client.go        |  4 ++--
 client/retry/backoff.go      | 24 ++++++++++++++----------
 client/retry/backoff_test.go | 26 ++++++++++++++++++++++----
 3 files changed, 38 insertions(+), 16 deletions(-)

diff --git a/client/http/client.go b/client/http/client.go
index 35d184e95db..0bf68d88df9 100644
--- a/client/http/client.go
+++ b/client/http/client.go
@@ -166,11 +166,11 @@ func (ci *clientInner) requestWithRetry(
 	}
 	// Copy a new backoffer for each request.
 	bo := *reqInfo.bo
-	// Set the retryable checker for the backoffer.
+	// Set the retryable checker for the backoffer if it's not set.
 	bo.SetRetryableChecker(func(err error) bool {
 		// Backoffer also needs to check the status code to determine whether to retry.
 		return err != nil && !noNeedRetry(statusCode)
-	})
+	}, false)
 	return bo.Exec(ctx, execFunc)
 }
 
diff --git a/client/retry/backoff.go b/client/retry/backoff.go
index 580e466badb..6c72b68ab9d 100644
--- a/client/retry/backoff.go
+++ b/client/retry/backoff.go
@@ -50,7 +50,7 @@ type Backoffer struct {
 	// total defines the max total time duration cost in retrying. If it's 0, it means infinite retry until success.
 	total time.Duration
 	// retryableChecker is used to check if the error is retryable.
-	// By default, all errors are retryable.
+	// If it's not set, it will use `defaultRetryableChecker` to retry on all non-nil errors.
 	retryableChecker func(err error) bool
 	// logInterval defines the log interval for retrying.
 	logInterval time.Duration
@@ -132,12 +132,9 @@ func InitialBackoffer(base, max, total time.Duration, opts ...Option) *Backoffer
 		total = base
 	}
 	bo := &Backoffer{
-		base:  base,
-		max:   max,
-		total: total,
-		retryableChecker: func(err error) bool {
-			return err != nil
-		},
+		base:         base,
+		max:          max,
+		total:        total,
 		next:         base,
 		currentTotal: 0,
 		attempt:      0,
@@ -148,18 +145,25 @@ func InitialBackoffer(base, max, total time.Duration, opts ...Option) *Backoffer
 	return bo
 }
 
-// SetRetryableChecker sets the retryable checker.
-func (bo *Backoffer) SetRetryableChecker(checker func(err error) bool) {
+// SetRetryableChecker sets the retryable checker, `overwrite` flag is used to indicate whether to overwrite the existing checker.
+func (bo *Backoffer) SetRetryableChecker(checker func(err error) bool, overwrite bool) {
+	if !overwrite && bo.retryableChecker != nil {
+		return
+	}
 	bo.retryableChecker = checker
 }
 
 func (bo *Backoffer) isRetryable(err error) bool {
 	if bo.retryableChecker == nil {
-		return true
+		return defaultRetryableChecker(err)
 	}
 	return bo.retryableChecker(err)
 }
 
+func defaultRetryableChecker(err error) bool {
+	return err != nil
+}
+
 // nextInterval for now use the `exponentialInterval`.
 func (bo *Backoffer) nextInterval() time.Duration {
 	return bo.exponentialInterval()
diff --git a/client/retry/backoff_test.go b/client/retry/backoff_test.go
index 8df06b75f94..8dd44033b55 100644
--- a/client/retry/backoff_test.go
+++ b/client/retry/backoff_test.go
@@ -95,16 +95,34 @@ func TestBackoffer(t *testing.T) {
 	// Test the retryable checker.
 	execCount = 0
 	bo = InitialBackoffer(base, max, total)
-	bo.SetRetryableChecker(func(error) bool {
+	retryableChecker := func(error) bool {
 		return execCount < 2
-	})
-	err = bo.Exec(ctx, func() error {
+	}
+	bo.SetRetryableChecker(retryableChecker, false)
+	execFunc := func() error {
 		execCount++
 		return nil
-	})
+	}
+	err = bo.Exec(ctx, execFunc)
+	re.NoError(err)
+	re.Equal(2, execCount)
+	re.True(isBackofferReset(bo))
+	// Test the retryable checker with overwrite.
+	execCount = 0
+	retryableChecker = func(error) bool {
+		return execCount < 4
+	}
+	bo.SetRetryableChecker(retryableChecker, false)
+	err = bo.Exec(ctx, execFunc)
 	re.NoError(err)
 	re.Equal(2, execCount)
 	re.True(isBackofferReset(bo))
+	execCount = 0
+	bo.SetRetryableChecker(retryableChecker, true)
+	err = bo.Exec(ctx, execFunc)
+	re.NoError(err)
+	re.Equal(4, execCount)
+	re.True(isBackofferReset(bo))
 }
 
 func isBackofferReset(bo *Backoffer) bool {