fix(retry): add status to fix MixRetry race issue

Author: YangruiEmma

pkg/remote/codec/default_codec.go

@@ -194,11 +194,6 @@ func (c *defaultCodec) DecodeMeta(ctx context.Context, message remote.Message, i
 	if flagBuf, err = in.Peek(2 * Size32); err != nil {
 		return perrors.NewProtocolErrorWithErrMsg(err, fmt.Sprintf("default codec read failed: %s", err.Error()))
 	}
-
-	if err = checkRPCState(ctx, message); err != nil {
-		// there is one call has finished in retry task, it doesn't need to do decode for this call
-		return err
-	}
 	isTTHeader := IsTTHeader(flagBuf)
 	// 1. decode header
 	if isTTHeader {
@@ -254,6 +249,11 @@ func (c *defaultCodec) DecodePayload(ctx context.Context, message remote.Message
 
 // Decode implements the remote.Codec interface, it does complete message decode include header and payload.
 func (c *defaultCodec) Decode(ctx context.Context, message remote.Message, in remote.ByteBuffer) (err error) {
+	defer updateRPCDecodeState(ctx, message, err)
+	if err = checkRPCState(ctx, message); err != nil {
+		// there is one call has finished in retry task, it doesn't need to do decode for this call
+		return err
+	}
 	// 1. decode meta
 	if err = c.DecodeMeta(ctx, message, in); err != nil {
 		return err
@@ -381,6 +381,8 @@ func isThriftFramedBinary(flagBuf []byte) bool {
 	return binary.BigEndian.Uint32(flagBuf[Size32:])&MagicMask == ThriftV1Magic
 }
 
+// checkRPCState is used to avoid concurrent decoding situations
+// specific situation: request retry and timeout
 func checkRPCState(ctx context.Context, message remote.Message) error {
 	if message.RPCRole() == remote.Server {
 		return nil
@@ -398,6 +400,23 @@ func checkRPCState(ctx context.Context, message remote.Message) error {
 	return nil
 }
 
+// updateRPCDecodeState is used to update the decoding status on the client side, which is for retry scenarios.
+// If decoding is completed, update CtxRespOp to OpDone.
+// Note that the update will only occur when it is actually encoded, if it is returned by ErrRPCFinal, it cannot be updated
+func updateRPCDecodeState(ctx context.Context, message remote.Message, err error) {
+	if err == kerrors.ErrRPCFinish {
+		fmt.Println("ErrRPCFinish 不更新")
+		return
+	}
+	if message.RPCRole() == remote.Server {
+		return
+	}
+	if respOp, ok := ctx.Value(retry.CtxRespOp).(*int32); ok {
+		atomic.CompareAndSwapInt32(respOp, retry.OpDoing, retry.OpDone)
+	}
+	return
+}
+
 func checkPayload(flagBuf []byte, message remote.Message, in remote.ByteBuffer, isTTHeader bool, maxPayloadSize int) error {
 	var transProto transport.Protocol
 	var codecType serviceinfo.PayloadCodec

pkg/retry/mixed_retryer.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"runtime"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -157,10 +158,10 @@ func (r *mixedRetryer) Do(ctx context.Context, rpcCall RPCCallFunc, firstRI rpci
 			}
 		case res := <-callDone:
 			// result retry
-			if respOp, ok := ctx.Value(CtxRespOp).(*int32); ok {
-				// must set as OpNo, or the new resp cannot be decoded
-				atomic.StoreInt32(respOp, OpNo)
-			}
+
+			// must set CtxRespOp as OpNo, or the new resp cannot be decoded
+			respOp, _ := ctx.Value(CtxRespOp).(*int32)
+
 			doneCount++
 			isFinishErr := res.err != nil && errors.Is(res.err, kerrors.ErrRPCFinish)
 			if nonFinishedErrRes == nil || !isFinishErr {
@@ -175,30 +176,47 @@ func (r *mixedRetryer) Do(ctx context.Context, rpcCall RPCCallFunc, firstRI rpci
 						// But if all requests return this error, it must be a bug, preventive panic to avoid dead loop
 						panic(errUnexpectedFinish)
 					}
+					updateRespOpStat(respOp)
 					continue
 				}
 				if callCount < retryTimes+1 {
 					if msg, ok := r.ShouldRetry(ctx, nil, callCount, req, cbKey); ok {
 						if r.isRetryResult(ctx, res.ri, res.resp, res.err, &r.policy.FailurePolicy) {
 							doCall = true
 							timer.Reset(retryDelay)
-							continue
+							updateRespOpStat(respOp)
+
 						}
 					} else if msg != "" {
 						appendMsg := fmt.Sprintf("retried %d, %s", callCount-1, msg)
 						appendErrMsg(res.err, appendMsg)
 					}
 				} else if r.isRetryResult(ctx, res.ri, res.resp, res.err, &r.policy.FailurePolicy) {
+					// only the Stat is Done, can be reset
+					updateRespOpStat(respOp)
 					continue
 				}
 			}
+			atomic.StoreInt32(respOp, OpDone)
 			atomic.StoreInt32(&abort, 1)
 			recordRetryInfo(nonFinishedErrRes.ri, atomic.LoadInt32(&callTimes), callCosts.String())
 			return nonFinishedErrRes.ri, false, nonFinishedErrRes.err
 		}
 	}
 }
 
+func updateRespOpStat(respOp *int32) {
+	// only the Stat is Done, can be reset
+	if atomic.CompareAndSwapInt32(respOp, OpDone, OpNo) {
+		return
+	}
+	var i, maxTry = 0, 3
+	for i < maxTry && atomic.LoadInt32(respOp) == OpDoing {
+		runtime.Gosched()
+		i++
+	}
+}
+
 // UpdatePolicy implements the Retryer interface.
 func (r *mixedRetryer) UpdatePolicy(rp Policy) (err error) {
 	r.Lock()