tikvrpc: avoid attaching ctx twice

Signed-off-by: zyguan <[email protected]>

internal/locate/region_request.go

@@ -857,7 +857,9 @@ func (s *RegionRequestSender) SendReqCtx(
 		if req.InputRequestSource != "" && s.replicaSelector != nil {
 			patchRequestSource(req, s.replicaSelector.replicaType())
 		}
-		if e := tikvrpc.SetContext(req, rpcCtx.Meta, rpcCtx.Peer); e != nil {
+		// Since RPCClient.SendRequest will call AttachContext (when EncodeRequest) to attach the context to the
+		// request, here we just update the req.Context to avoid attaching the ctx twice.
+		if e := tikvrpc.SetContextNoAttach(req, rpcCtx.Meta, rpcCtx.Peer); e != nil {
 			return nil, nil, retryTimes, err
 		}
 		if s.replicaSelector != nil {

tikvrpc/gen.sh

@@ -97,3 +97,25 @@ cat <<EOF >> $output
 	return true
 }
 EOF
+
+cat <<EOF >> $output
+
+func isValidReqType(cmd CmdType) bool {
+	switch cmd {
+EOF
+
+for cmd in "${cmds[@]}"; do
+cat <<EOF >> $output
+	case Cmd${cmd}:
+		return true
+EOF
+done
+
+cat <<EOF >> $output
+	case CmdCopStream, CmdMPPTask, CmdMPPConn, CmdMPPCancel, CmdMPPAlive, CmdEmpty:
+		return true
+	default:
+		return false
+	}
+}
+EOF

tikvrpc/tikvrpc.go

@@ -817,6 +817,19 @@ func SetContext(req *Request, region *metapb.Region, peer *metapb.Peer) error {
 	return nil
 }
 
+// SetContextNoAttach likes SetContext, but it doesn't attach the context to the underlying request.
+func SetContextNoAttach(req *Request, region *metapb.Region, peer *metapb.Peer) error {
+	if !isValidReqType(req.Type) {
+		return errors.Errorf("invalid request type %v", req.Type)
+	}
+	if region != nil {
+		req.Context.RegionId = region.Id
+		req.Context.RegionEpoch = region.RegionEpoch
+	}
+	req.Context.Peer = peer
+	return nil
+}
+
 // GenRegionErrorResp returns corresponding Response with specified RegionError
 // according to the given req.
 func GenRegionErrorResp(req *Request, e *errorpb.Error) (*Response, error) {