From a97cdde2afc0add88a0584102e7dfcec724bb1c3 Mon Sep 17 00:00:00 2001
From: YangruiEmma <yangrui.emma@bytedance.com>
Date: Wed, 28 Aug 2024 05:41:04 +0800
Subject: [PATCH] fix: mixed retry

---
 pkg/retry/mixed_retryer.go |  6 +++---
 pkg/retry/mixed_test.go    | 28 ++++++++++++++++++++++++++++
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/pkg/retry/mixed_retryer.go b/pkg/retry/mixed_retryer.go
index 4d6e450b6c..4669287c3c 100644
--- a/pkg/retry/mixed_retryer.go
+++ b/pkg/retry/mixed_retryer.go
@@ -165,6 +165,9 @@ func (r *mixedRetryer) Do(ctx context.Context, rpcCall RPCCallFunc, firstRI rpci
 			}
 			doneCount++
 			isFinishErr := res.err != nil && errors.Is(res.err, kerrors.ErrRPCFinish)
+			if nonFinishedErrRes == nil || !isFinishErr {
+				nonFinishedErrRes = res
+			}
 			if doneCount < retryTimes+1 {
 				if isFinishErr {
 					// There will be only one request (goroutine) pass the `checkRPCState`, others will skip decoding
@@ -191,9 +194,6 @@ func (r *mixedRetryer) Do(ctx context.Context, rpcCall RPCCallFunc, firstRI rpci
 					continue
 				}
 			}
-			if nonFinishedErrRes == nil || !isFinishErr {
-				nonFinishedErrRes = res
-			}
 			atomic.StoreInt32(&abort, 1)
 			recordRetryInfo(nonFinishedErrRes.ri, atomic.LoadInt32(&callTimes), callCosts.String())
 			return nonFinishedErrRes.ri, false, nonFinishedErrRes.err
diff --git a/pkg/retry/mixed_test.go b/pkg/retry/mixed_test.go
index 46cd2fa7d9..c8312e107f 100644
--- a/pkg/retry/mixed_test.go
+++ b/pkg/retry/mixed_test.go
@@ -235,6 +235,34 @@ func TestMixedRetry(t *testing.T) {
 		test.Assert(t, ok)
 		test.Assert(t, v == remoteTagValue)
 	})
+
+	// case4: RPCFinishErr
+	t.Run("RPCFinishErr", func(t *testing.T) {
+		mockErr := errors.New("mock")
+		retryWithRPCFinishErr := func(callCount int32) RPCCallFunc {
+			// fails for the first call if callTimes is initialized to 0
+			return func(ctx context.Context, r Retryer) (rpcinfo.RPCInfo, interface{}, error) {
+				time.Sleep(50 * time.Millisecond)
+				ct := atomic.AddInt32(&callCount, 1)
+				if ct == 1 || ct == 2 {
+					// first call retry TransErr with specified errCode
+					return genRPCInfo(), nil, mockErr
+				} else {
+					return genRPCInfo(), nil, kerrors.ErrRPCFinish
+				}
+			}
+		}
+
+		rc := NewRetryContainer()
+		mp := NewMixedPolicyWithResultRetry(10, AllErrorRetry())
+		mp.WithMaxRetryTimes(3)
+		p := BuildMixedPolicy(mp)
+		ri = genRPCInfo()
+		_, ok, err := rc.WithRetryIfNeeded(ctx, &p, retryWithRPCFinishErr(0), ri, nil)
+		test.Assert(t, err != nil, err)
+		test.Assert(t, err == mockErr, err)
+		test.Assert(t, !ok)
+	})
 }
 
 // Assuming the first request returns at 300ms, the second request costs 150ms