Repository: tikv/client-go
Pull request #1016 (2023-10-12 03:03:34)

request_source: add retry info to request source(#953)
Cherry pick #953 and #960 to 7.1.


Patch:
diff --git a/go.mod b/go.mod
index 1962a77..360bbba 100644
--- a/go.mod
+++ b/go.mod
@@ -14,7 +14,7 @@ require (
 	github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c
 	github.com/pingcap/failpoint v0.0.0-20220801062533-2eaa32854a6c
 	github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989
-	github.com/pingcap/kvproto v0.0.0-20230523065550-8b641fa69bf3
+	github.com/pingcap/kvproto v0.0.0-20231011074246-fa00d2b03372
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3
 	github.com/pkg/errors v0.9.1
 	github.com/prometheus/client_golang v1.15.1
diff --git a/go.sum b/go.sum
index cc14fb8..cefa298 100644
--- a/go.sum
+++ b/go.sum
@@ -147,8 +147,8 @@ github.com/pingcap/failpoint v0.0.0-20220801062533-2eaa32854a6c h1:CgbKAHto5CQgW
 github.com/pingcap/failpoint v0.0.0-20220801062533-2eaa32854a6c/go.mod h1:4qGtCB0QK0wBzKtFEGDhxXnSnbQApw1gc9siScUl8ew=
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989 h1:surzm05a8C9dN8dIUmo4Be2+pMRb6f55i+UIYrluu2E=
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989/go.mod h1:O17XtbryoCJhkKGbT62+L2OlrniwqiGLSqrmdHCMzZw=
-github.com/pingcap/kvproto v0.0.0-20230523065550-8b641fa69bf3 h1:iAIjJVgrPuXjpAiMDcJvz4Y4Qf4KypiCsqy3UVzU6FQ=
-github.com/pingcap/kvproto v0.0.0-20230523065550-8b641fa69bf3/go.mod h1:guCyM5N+o+ru0TsoZ1hi9lDjUMs2sIBjW3ARTEpVbnk=
+github.com/pingcap/kvproto v0.0.0-20231011074246-fa00d2b03372 h1:so6ffMPEx9HYGOOi9T0ezTZlI7LvpYBUQyWWyz8kND8=
+github.com/pingcap/kvproto v0.0.0-20231011074246-fa00d2b03372/go.mod h1:guCyM5N+o+ru0TsoZ1hi9lDjUMs2sIBjW3ARTEpVbnk=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 h1:HR/ylkkLmGdSSDaD8IDP+SZrdhV1Kibl9KrHxJ9eciw=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=
 github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
diff --git a/internal/locate/region_cache.go b/internal/locate/region_cache.go
index e60efe6..1669fb8 100644
--- a/internal/locate/region_cache.go
+++ b/internal/locate/region_cache.go
@@ -620,16 +620,17 @@ func (c *RegionCache) SetPDClient(client pd.Client) {
 
 // RPCContext contains data that is needed to send RPC to a region.
 type RPCContext struct {
-	Region     RegionVerID
-	Meta       *metapb.Region
-	Peer       *metapb.Peer
-	AccessIdx  AccessIndex
-	Store      *Store
-	Addr       string
-	AccessMode accessMode
-	ProxyStore *Store // nil means proxy is not used
-	ProxyAddr  string // valid when ProxyStore is not nil
-	TiKVNum    int    // Number of TiKV nodes among the region's peers. Assuming non-TiKV peers are all TiFlash peers.
+	Region        RegionVerID
+	Meta          *metapb.Region
+	Peer          *metapb.Peer
+	AccessIdx     AccessIndex
+	Store         *Store
+	Addr          string
+	AccessMode    accessMode
+	ProxyStore    *Store // nil means proxy is not used
+	ProxyAddr     string // valid when ProxyStore is not nil
+	TiKVNum       int    // Number of TiKV nodes among the region's peers. Assuming non-TiKV peers are all TiFlash peers.
+	BucketVersion uint64
 
 	contextPatcher contextPatcher // kvrpcpb.Context fields that need to be overridden
 }
@@ -1929,6 +1930,26 @@ func (c *RegionCache) getStoresByLabels(labels []*metapb.StoreLabel) []*Store {
 	return s
 }
 
+// OnBucketVersionNotMatch removes the old buckets meta if the version is stale.
+func (c *RegionCache) OnBucketVersionNotMatch(ctx *RPCContext, version uint64, keys [][]byte) {
+	r := c.GetCachedRegionWithRLock(ctx.Region)
+	if r == nil {
+		return
+	}
+
+	buckets := r.getStore().buckets
+	if buckets == nil || buckets.GetVersion() < version {
+		oldStore := r.getStore()
+		store := oldStore.clone()
+		store.buckets = &metapb.Buckets{
+			Version:  version,
+			Keys:     keys,
+			RegionId: r.meta.GetId(),
+		}
+		r.compareAndSwapStore(oldStore, store)
+	}
+}
+
 // OnRegionEpochNotMatch removes the old region and inserts new regions into the cache.
 // It returns whether retries the request because it's possible the region epoch is ahead of TiKV's due to slow appling.
 func (c *RegionCache) OnRegionEpochNotMatch(bo *retry.Backoffer, ctx *RPCContext, currentRegions []*metapb.Region) (bool, error) {
diff --git a/internal/locate/region_request.go b/internal/locate/region_request.go
index 176850c..54f335b 100644
--- a/internal/locate/region_request.go
+++ b/internal/locate/region_request.go
@@ -1289,7 +1289,6 @@ func (s *RegionRequestSender) SendReqCtx(
 	totalErrors := make(map[string]int)
 	for {
 		if retryTimes > 0 {
-			req.IsRetryRequest = true
 			if retryTimes%100 == 0 {
 				logutil.Logger(bo.GetCtx()).Warn(
 					"retry",
@@ -1342,8 +1341,17 @@ func (s *RegionRequestSender) SendReqCtx(
 			}
 		}
 
+		if e := tikvrpc.SetContext(req, rpcCtx.Meta, rpcCtx.Peer); e != nil {
+			return nil, nil, retryTimes, 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, retryTimes, req, totalErrors)
@@ -1539,10 +1547,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) {
-	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 {
@@ -1849,6 +1853,8 @@ func regionErrorToLabel(e *errorpb.Error) string {
 		return "flashback_not_prepared"
 	} else if e.GetIsWitness() != nil {
 		return "peer_is_witness"
+	} else if e.GetBucketVersionNotMatch() != nil {
+		return "bucket_version_not_match"
 	} else if strings.Contains(e.Message, "mismatch peer id") {
 		// the error message is like "[components/raftstore/src/store/util.rs:428]: mismatch peer id ? != ?"
 		// the `mismatch peer id` error does not has a specific error type, so we have to match the error message.
@@ -2006,6 +2012,18 @@ func (s *RegionRequestSender) onRegionError(
 		return retry, err
 	}
 
+	if bucketVersionNotMatch := regionErr.GetBucketVersionNotMatch(); bucketVersionNotMatch != nil {
+		logutil.Logger(bo.GetCtx()).Debug(
+			"tikv reports `BucketVersionNotMatch` retry later",
+			zap.Uint64("latest-bucket-version", bucketVersionNotMatch.Version),
+			zap.Uint64("request-bucket-version", ctx.BucketVersion),
+			zap.Stringer("ctx", ctx),
+		)
+		// bucket version is not match, we should split this cop request again.
+		s.regionCache.OnBucketVersionNotMatch(ctx, bucketVersionNotMatch.Version, bucketVersionNotMatch.Keys)
+		return false, nil
+	}
+
 	if serverIsBusy := regionErr.GetServerIsBusy(); serverIsBusy != nil {
 		if s.replicaSelector != nil {
 			return s.replicaSelector.onServerIsBusy(bo, ctx, req, serverIsBusy)
@@ -2209,3 +2227,54 @@ 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)
+	case *tryIdleReplica:
+		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
+	defer func() {
+		// TiKV does the limit control by the last part of the request source.
+		sb.WriteByte('_')
+		sb.WriteString(req.InputRequestSource)
+		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 4d81e14..76f4e61 100644
--- a/tikvrpc/tikvrpc.go
+++ b/tikvrpc/tikvrpc.go
@@ -234,6 +234,10 @@ type Request struct {
 	ForwardedHost string
 	// ReplicaNumber is the number of current replicas, which is used to calculate the RU cost.
 	ReplicaNumber int64
+	// 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 221cf59..7b07a92 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)
@@ -245,12 +247,16 @@ 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(),
 			ResourceControlContext: &kvrpcpb.ResourceControlContext{
 				ResourceGroupName: s.snapshot.mu.resourceGroupName,
 			},
 			BusyThresholdMs: uint32(s.snapshot.mu.busyThreshold.Milliseconds()),
 		})
+		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)
 		}
@@ -263,6 +269,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 b4671d6..00f2b5d 100644
--- a/txnkv/txnsnapshot/snapshot.go
+++ b/txnkv/txnsnapshot/snapshot.go
@@ -387,6 +387,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{
@@ -398,12 +400,16 @@ func (s *KVSnapshot) batchGetSingleRegion(bo *retry.Backoffer, batch batchKeys,
 			TaskId:           s.mu.taskID,
 			ResourceGroupTag: s.mu.resourceGroupTag,
 			IsolationLevel:   s.isolationLevel.ToPB(),
-			RequestSource:    s.GetRequestSource(),
 			ResourceControlContext: &kvrpcpb.ResourceControlContext{
 				ResourceGroupName: s.mu.resourceGroupName,
 			},
 			BusyThresholdMs: uint32(busyThresholdMs),
 		})
+		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)
 		}
@@ -435,6 +441,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.
@@ -618,12 +625,12 @@ 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(),
 			ResourceControlContext: &kvrpcpb.ResourceControlContext{
 				ResourceGroupName: s.mu.resourceGroupName,
 			},
 			BusyThresholdMs: uint32(s.mu.busyThreshold.Milliseconds()),
 		})
+	req.InputRequestSource = s.GetRequestSource()
 	if s.mu.resourceGroupTag == nil && s.mu.resourceGroupTagger != nil {
 		s.mu.resourceGroupTagger(req)
 	}


Test patch:
diff --git a/integration_tests/go.mod b/integration_tests/go.mod
index 090cf56..5f4229b 100644
--- a/integration_tests/go.mod
+++ b/integration_tests/go.mod
@@ -6,7 +6,7 @@ require (
 	github.com/ninedraft/israce v0.0.3
 	github.com/pingcap/errors v0.11.5-0.20221009092201-b66cddb77c32
 	github.com/pingcap/failpoint v0.0.0-20220801062533-2eaa32854a6c
-	github.com/pingcap/kvproto v0.0.0-20230523065550-8b641fa69bf3
+	github.com/pingcap/kvproto v0.0.0-20231011074246-fa00d2b03372
 	github.com/pingcap/tidb v1.1.0-beta.0.20230428082552-991d99c1e0b9
 	github.com/pkg/errors v0.9.1
 	github.com/stretchr/testify v1.8.2
diff --git a/integration_tests/go.sum b/integration_tests/go.sum
index 009827c..88bbff1 100644
--- a/integration_tests/go.sum
+++ b/integration_tests/go.sum
@@ -342,8 +342,8 @@ github.com/pingcap/fn v0.0.0-20200306044125-d5540d389059 h1:Pe2LbxRmbTfAoKJ65bZL
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989 h1:surzm05a8C9dN8dIUmo4Be2+pMRb6f55i+UIYrluu2E=
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989/go.mod h1:O17XtbryoCJhkKGbT62+L2OlrniwqiGLSqrmdHCMzZw=
 github.com/pingcap/kvproto v0.0.0-20191211054548-3c6b38ea5107/go.mod h1:WWLmULLO7l8IOcQG+t+ItJ3fEcrL5FxF0Wu+HrMy26w=
-github.com/pingcap/kvproto v0.0.0-20230523065550-8b641fa69bf3 h1:iAIjJVgrPuXjpAiMDcJvz4Y4Qf4KypiCsqy3UVzU6FQ=
-github.com/pingcap/kvproto v0.0.0-20230523065550-8b641fa69bf3/go.mod h1:guCyM5N+o+ru0TsoZ1hi9lDjUMs2sIBjW3ARTEpVbnk=
+github.com/pingcap/kvproto v0.0.0-20231011074246-fa00d2b03372 h1:so6ffMPEx9HYGOOi9T0ezTZlI7LvpYBUQyWWyz8kND8=
+github.com/pingcap/kvproto v0.0.0-20231011074246-fa00d2b03372/go.mod h1:guCyM5N+o+ru0TsoZ1hi9lDjUMs2sIBjW3ARTEpVbnk=
 github.com/pingcap/log v0.0.0-20210625125904-98ed8e2eb1c7/go.mod h1:8AanEdAHATuRurdGxZXBz0At+9avep+ub7U1AGYLIMM=
 github.com/pingcap/log v1.1.0/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=
 github.com/pingcap/log v1.1.1-0.20230317032135-a0d097d16e22 h1:2SOzvGvE8beiC1Y4g9Onkvu6UmuBBOeWRGQEjJaT/JY=
diff --git a/internal/locate/region_cache_test.go b/internal/locate/region_cache_test.go
index 619da2d..6226a1c 100644
--- a/internal/locate/region_cache_test.go
+++ b/internal/locate/region_cache_test.go
@@ -1646,6 +1646,15 @@ func (s *testRegionCacheSuite) TestShouldNotRetryFlashback() {
 	shouldRetry, err = reqSend.onRegionError(s.bo, ctx, nil, &errorpb.Error{FlashbackNotPrepared: &errorpb.FlashbackNotPrepared{}})
 	s.Error(err)
 	s.False(shouldRetry)
+
+	shouldRetry, err = reqSend.onRegionError(s.bo, ctx, nil, &errorpb.Error{BucketVersionNotMatch: &errorpb.BucketVersionNotMatch{Keys: [][]byte{[]byte("a")}, Version: 1}})
+	s.Nil(err)
+	s.False(shouldRetry)
+	ctx.Region.GetID()
+	key, err := s.cache.LocateKey(s.bo, []byte("a"))
+	s.Nil(err)
+	s.Equal(key.Buckets.Keys, [][]byte{[]byte("a")})
+	s.Equal(key.Buckets.Version, uint64(1))
 }
 
 func (s *testRegionCacheSuite) TestBackgroundCacheGC() {
diff --git a/internal/locate/region_request3_test.go b/internal/locate/region_request3_test.go
index 4757a60..cd9a8e3 100644
--- a/internal/locate/region_request3_test.go
+++ b/internal/locate/region_request3_test.go
@@ -1338,3 +1338,80 @@ func (s *testRegionRequestToThreeStoresSuite) TestStaleReadFallback() {
 	s.True(ok)
 	s.Equal(getResp.Value, 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(firstReplica+"_test", 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("retry_"+firstReplica+"_"+retryReplica+"_test", req.RequestSource)
+		}
+	}
+}