Repository: tikv/client-go
Pull request #959 (2023-09-04 05:27:34)

add retry info to request source (#953)
Cherry pick #953 to 6.5.




Patch:
diff --git a/internal/locate/region_request.go b/internal/locate/region_request.go
index e47f0d8..c840029 100644
--- a/internal/locate/region_request.go
+++ b/internal/locate/region_request.go
@@ -1117,7 +1117,6 @@ func (s *RegionRequestSender) SendReqCtx(
 	totalErrors := make(map[string]int)
 	for {
 		if tryTimes > 0 {
-			req.IsRetryRequest = true
 			if tryTimes%100 == 0 {
 				logutil.Logger(bo.GetCtx()).Warn("retry", zap.Uint64("region", regionID.GetID()), zap.Int("times", tryTimes))
 			}
@@ -1166,8 +1165,17 @@ func (s *RegionRequestSender) SendReqCtx(
 			}
 		}
 
+		if e := tikvrpc.SetContext(req, rpcCtx.Meta, rpcCtx.Peer); e != nil {
+			return nil, nil, err
+		}
+		rpcCtx.contextPatcher.applyTo(&req.Context)
+		if req.InputRequestSource != "" && s.replicaSelector != nil {
+			s.replicaSelector.patchRequestSource(req, rpcCtx)
+		}
+
 		var retry bool
 		resp, retry, err = s.sendReqToRegion(bo, rpcCtx, req, timeout)
+		req.IsRetryRequest = true
 		if err != nil {
 			msg := fmt.Sprintf("send request failed, err: %v", err.Error())
 			s.logSendReqError(bo, msg, regionID, tryTimes, req, totalErrors)
@@ -1356,10 +1364,6 @@ func fetchRespInfo(resp *tikvrpc.Response) string {
 func (s *RegionRequestSender) sendReqToRegion(bo *retry.Backoffer, rpcCtx *RPCContext, req *tikvrpc.Request, timeout time.Duration) (resp *tikvrpc.Response, retry bool, err error) {
 	req.ApiVersion = s.apiVersion
 
-	if e := tikvrpc.SetContext(req, rpcCtx.Meta, rpcCtx.Peer); e != nil {
-		return nil, false, err
-	}
-	rpcCtx.contextPatcher.applyTo(&req.Context)
 	// judge the store limit switch.
 	if limit := kv.StoreLimit.Load(); limit > 0 {
 		if err := s.getStoreToken(rpcCtx.Store, limit); err != nil {
@@ -1921,3 +1925,51 @@ func (s *staleReadMetricsCollector) onResp(tp tikvrpc.CmdType, resp *tikvrpc.Res
 		metrics.StaleReadRemoteInBytes.Add(float64(size))
 	}
 }
+
+func (s *replicaSelector) replicaType(rpcCtx *RPCContext) string {
+	leaderIdx := -1
+	switch v := s.state.(type) {
+	case *accessKnownLeader:
+		return "leader"
+	case *tryFollower:
+		return "follower"
+	case *accessFollower:
+		leaderIdx = int(v.leaderIdx)
+	}
+	if leaderIdx > -1 && rpcCtx != nil && rpcCtx.Peer != nil {
+		for idx, replica := range s.replicas {
+			if replica.peer.Id == rpcCtx.Peer.Id {
+				if idx == leaderIdx {
+					return "leader"
+				}
+				return "follower"
+			}
+		}
+	}
+	return "unknown"
+}
+
+func (s *replicaSelector) patchRequestSource(req *tikvrpc.Request, rpcCtx *RPCContext) {
+	var sb strings.Builder
+	sb.WriteString(req.InputRequestSource)
+	sb.WriteByte('-')
+	defer func() {
+		req.RequestSource = sb.String()
+	}()
+
+	replicaType := s.replicaType(rpcCtx)
+
+	if req.IsRetryRequest {
+		sb.WriteString("retry_")
+		sb.WriteString(req.ReadType)
+		sb.WriteByte('_')
+		sb.WriteString(replicaType)
+		return
+	}
+	if req.StaleRead {
+		req.ReadType = "stale_" + replicaType
+	} else {
+		req.ReadType = replicaType
+	}
+	sb.WriteString(req.ReadType)
+}
diff --git a/tikvrpc/tikvrpc.go b/tikvrpc/tikvrpc.go
index 33a4e05..0119e55 100644
--- a/tikvrpc/tikvrpc.go
+++ b/tikvrpc/tikvrpc.go
@@ -229,6 +229,10 @@ type Request struct {
 	// If it's not empty, the store which receive the request will forward it to
 	// the forwarded host. It's useful when network partition occurs.
 	ForwardedHost string
+	// The initial read type, note this will be assigned in the first try, no need to set it outside the client.
+	ReadType string
+	// InputRequestSource is the input source of the request, if it's not empty, the final RequestSource sent to store will be attached with the retry info.
+	InputRequestSource string
 }
 
 // NewRequest returns new kv rpc request.
diff --git a/txnkv/txnsnapshot/scan.go b/txnkv/txnsnapshot/scan.go
index 2e24d52..25579de 100644
--- a/txnkv/txnsnapshot/scan.go
+++ b/txnkv/txnsnapshot/scan.go
@@ -202,6 +202,8 @@ func (s *Scanner) getData(bo *retry.Backoffer) error {
 	var loc *locate.KeyLocation
 	var resolvingRecordToken *int
 	var err error
+	// the states in request need to keep when retry request.
+	var readType string
 	for {
 		if !s.reverse {
 			loc, err = s.snapshot.store.GetRegionCache().LocateKey(bo, s.nextStartKey)
@@ -251,8 +253,12 @@ func (s *Scanner) getData(bo *retry.Backoffer) error {
 			TaskId:           s.snapshot.mu.taskID,
 			ResourceGroupTag: s.snapshot.mu.resourceGroupTag,
 			IsolationLevel:   s.snapshot.isolationLevel.ToPB(),
-			RequestSource:    s.snapshot.GetRequestSource(),
 		})
+		if readType != "" {
+			req.ReadType = readType
+			req.IsRetryRequest = true
+		}
+		req.InputRequestSource = s.snapshot.GetRequestSource()
 		if s.snapshot.mu.resourceGroupTag == nil && s.snapshot.mu.resourceGroupTagger != nil {
 			s.snapshot.mu.resourceGroupTagger(req)
 		}
@@ -265,6 +271,7 @@ func (s *Scanner) getData(bo *retry.Backoffer) error {
 		if err != nil {
 			return err
 		}
+		readType = req.ReadType
 		if regionErr != nil {
 			logutil.BgLogger().Debug("scanner getData failed",
 				zap.Stringer("regionErr", regionErr))
diff --git a/txnkv/txnsnapshot/snapshot.go b/txnkv/txnsnapshot/snapshot.go
index 76b5aed..e79f0fc 100644
--- a/txnkv/txnsnapshot/snapshot.go
+++ b/txnkv/txnsnapshot/snapshot.go
@@ -370,6 +370,8 @@ func (s *KVSnapshot) batchGetSingleRegion(bo *retry.Backoffer, batch batchKeys,
 
 	pending := batch.keys
 	var resolvingRecordToken *int
+	// the states in request need to keep when retry request.
+	var readType string
 	for {
 		s.mu.RLock()
 		req := tikvrpc.NewReplicaReadRequest(tikvrpc.CmdBatchGet, &kvrpcpb.BatchGetRequest{
@@ -381,8 +383,12 @@ func (s *KVSnapshot) batchGetSingleRegion(bo *retry.Backoffer, batch batchKeys,
 			TaskId:           s.mu.taskID,
 			ResourceGroupTag: s.mu.resourceGroupTag,
 			IsolationLevel:   s.isolationLevel.ToPB(),
-			RequestSource:    s.GetRequestSource(),
 		})
+		req.InputRequestSource = s.GetRequestSource()
+		if readType != "" {
+			req.ReadType = readType
+			req.IsRetryRequest = true
+		}
 		if s.mu.resourceGroupTag == nil && s.mu.resourceGroupTagger != nil {
 			s.mu.resourceGroupTagger(req)
 		}
@@ -414,6 +420,7 @@ func (s *KVSnapshot) batchGetSingleRegion(bo *retry.Backoffer, batch batchKeys,
 		if err != nil {
 			return err
 		}
+		readType = req.ReadType
 		if regionErr != nil {
 			// For other region error and the fake region error, backoff because
 			// there's something wrong.
@@ -587,8 +594,8 @@ func (s *KVSnapshot) get(ctx context.Context, bo *retry.Backoffer, k []byte) ([]
 			TaskId:           s.mu.taskID,
 			ResourceGroupTag: s.mu.resourceGroupTag,
 			IsolationLevel:   s.isolationLevel.ToPB(),
-			RequestSource:    s.GetRequestSource(),
 		})
+	req.InputRequestSource = s.GetRequestSource()
 	if s.mu.resourceGroupTag == nil && s.mu.resourceGroupTagger != nil {
 		s.mu.resourceGroupTagger(req)
 	}


Test patch:
diff --git a/internal/client/client_test.go b/internal/client/client_test.go
index fd8916f..52bf305 100644
--- a/internal/client/client_test.go
+++ b/internal/client/client_test.go
@@ -700,7 +700,7 @@ func TestBatchClientRecoverAfterServerRestart(t *testing.T) {
 			cli.unlockForSend()
 			break
 		}
-		if time.Since(start) > time.Second*5 {
+		if time.Since(start) > time.Second*10 {
 			// It shouldn't take too long for batch_client to reconnect.
 			require.Fail(t, "wait batch client reconnect timeout")
 		}
diff --git a/internal/locate/region_request3_test.go b/internal/locate/region_request3_test.go
index c56c521..feb60fb 100644
--- a/internal/locate/region_request3_test.go
+++ b/internal/locate/region_request3_test.go
@@ -1229,3 +1229,80 @@ func (s *testRegionRequestToThreeStoresSuite) TestDoNotTryUnreachableLeader() {
 	// `tryFollower` always try the local peer firstly
 	s.Equal(follower.addr, string(resp.Resp.(*kvrpcpb.GetResponse).Value))
 }
+
+func (s *testRegionRequestToThreeStoresSuite) TestRetryRequestSource() {
+	leaderStore, _ := s.loadAndGetLeaderStore()
+	regionLoc, err := s.cache.LocateRegionByID(s.bo, s.regionID)
+	s.Nil(err)
+	req := tikvrpc.NewRequest(tikvrpc.CmdGet, &kvrpcpb.GetRequest{
+		Key: []byte("key"),
+	})
+	req.InputRequestSource = "test"
+
+	setReadType := func(req *tikvrpc.Request, readType string) {
+		req.StaleRead = false
+		req.ReplicaRead = false
+		switch readType {
+		case "leader":
+			return
+		case "follower":
+			req.ReplicaRead = true
+			req.ReplicaReadType = kv.ReplicaReadFollower
+		case "stale_follower", "stale_leader":
+			req.EnableStaleRead()
+		default:
+			panic("unreachable")
+		}
+	}
+
+	setTargetReplica := func(selector *replicaSelector, readType string) {
+		var leader bool
+		switch readType {
+		case "leader", "stale_leader":
+			leader = true
+		case "follower", "stale_follower":
+			leader = false
+		default:
+			panic("unreachable")
+		}
+		for idx, replica := range selector.replicas {
+			if replica.store.storeID == leaderStore.storeID && leader {
+				selector.targetIdx = AccessIndex(idx)
+				return
+			}
+			if replica.store.storeID != leaderStore.storeID && !leader {
+				selector.targetIdx = AccessIndex(idx)
+				return
+			}
+		}
+		panic("unreachable")
+	}
+
+	firstReadReplicas := []string{"leader", "follower", "stale_follower", "stale_leader"}
+	retryReadReplicas := []string{"leader", "follower"}
+	for _, firstReplica := range firstReadReplicas {
+		for _, retryReplica := range retryReadReplicas {
+			bo := retry.NewBackoffer(context.Background(), -1)
+			req.IsRetryRequest = false
+			setReadType(req, firstReplica)
+			replicaSelector, err := newReplicaSelector(s.cache, regionLoc.Region, req)
+			s.Nil(err)
+			setTargetReplica(replicaSelector, firstReplica)
+			rpcCtx, err := replicaSelector.buildRPCContext(bo)
+			s.Nil(err)
+			replicaSelector.patchRequestSource(req, rpcCtx)
+			s.Equal("test-"+firstReplica, req.RequestSource)
+
+			// retry
+			setReadType(req, retryReplica)
+			replicaSelector, err = newReplicaSelector(s.cache, regionLoc.Region, req)
+			s.Nil(err)
+			setTargetReplica(replicaSelector, retryReplica)
+			rpcCtx, err = replicaSelector.buildRPCContext(bo)
+			s.Nil(err)
+			req.IsRetryRequest = true
+			replicaSelector.patchRequestSource(req, rpcCtx)
+			s.Equal("test-retry_"+firstReplica+"_"+retryReplica, req.RequestSource)
+		}
+	}
+}