Repository: tikv/client-go
Pull request #739 (2023-03-15 06:25:49)

compatible batch copr with load-based replica read
pingcap/tidb#42322

Disable the server-is-busy retry when there are batched tasks, the application should handle it.


disksing: @you06 do you have anyone who can review this PR?
cfzjywxk: /merge

Patch:
diff --git a/internal/locate/region_cache.go b/internal/locate/region_cache.go
index 5139d5b..ecb8cea 100644
--- a/internal/locate/region_cache.go
+++ b/internal/locate/region_cache.go
@@ -603,6 +603,7 @@ type storeSelectorOp struct {
 	leaderOnly   bool
 	preferLeader bool
 	labels       []*metapb.StoreLabel
+	stores       []uint64
 }
 
 // StoreSelectorOption configures storeSelectorOp.
@@ -629,6 +630,13 @@ func WithPerferLeader() StoreSelectorOption {
 	}
 }
 
+// WithMatchStores indicates selecting stores with matched store ids.
+func WithMatchStores(stores []uint64) StoreSelectorOption {
+	return func(op *storeSelectorOp) {
+		op.stores = stores
+	}
+}
+
 // GetTiKVRPCContext returns RPCContext for a region. If it returns nil, the region
 // must be out of date and already dropped from cache.
 func (c *RegionCache) GetTiKVRPCContext(bo *retry.Backoffer, id RegionVerID, replicaRead kv.ReplicaReadType, followerStoreSeed uint32, opts ...StoreSelectorOption) (*RPCContext, error) {
@@ -2480,6 +2488,19 @@ func (s *Store) IsLabelsMatch(labels []*metapb.StoreLabel) bool {
 	return true
 }
 
+// IsStoreMatch returns whether the store's id match the target ids.
+func (s *Store) IsStoreMatch(stores []uint64) bool {
+	if len(stores) < 1 {
+		return true
+	}
+	for _, storeID := range stores {
+		if s.storeID == storeID {
+			return true
+		}
+	}
+	return false
+}
+
 func isStoreContainLabel(labels []*metapb.StoreLabel, key string, val string) (res bool) {
 	for _, label := range labels {
 		if label.GetKey() == key && label.GetValue() == val {
@@ -2636,9 +2657,9 @@ func (s *Store) GetPeerAddr() string {
 	return s.peerAddr
 }
 
-// estimatedWaitTime returns an optimistic estimation of how long a request will wait in the store.
+// EstimatedWaitTime returns an optimistic estimation of how long a request will wait in the store.
 // It's calculated by subtracting the time since the last update from the wait time returned from TiKV.
-func (s *Store) estimatedWaitTime() time.Duration {
+func (s *Store) EstimatedWaitTime() time.Duration {
 	loadStats := s.loadStats.Load()
 	if loadStats == nil {
 		return 0
diff --git a/internal/locate/region_request.go b/internal/locate/region_request.go
index 5b21f73..a10924d 100644
--- a/internal/locate/region_request.go
+++ b/internal/locate/region_request.go
@@ -346,7 +346,7 @@ func (state *accessKnownLeader) next(bo *retry.Backoffer, selector *replicaSelec
 	if selector.busyThreshold > 0 {
 		// If the leader is busy in our estimation, change to tryIdleReplica state to try other replicas.
 		// If other replicas are all busy, tryIdleReplica will try the leader again without busy threshold.
-		leaderEstimated := selector.replicas[state.leaderIdx].store.estimatedWaitTime()
+		leaderEstimated := selector.replicas[state.leaderIdx].store.EstimatedWaitTime()
 		if leaderEstimated > selector.busyThreshold {
 			selector.state = &tryIdleReplica{leaderIdx: state.leaderIdx}
 			return nil, stateChanged{}
@@ -612,14 +612,27 @@ func (state *accessFollower) onSendFailure(bo *retry.Backoffer, selector *replic
 }
 
 func (state *accessFollower) isCandidate(idx AccessIndex, replica *replica) bool {
-	return !replica.isEpochStale() && !replica.isExhausted(1) &&
-		// The request can only be sent to the leader.
-		((state.option.leaderOnly && idx == state.leaderIdx) ||
-			// Choose a replica with matched labels.
-			(!state.option.leaderOnly && (state.tryLeader || idx != state.leaderIdx) && replica.store.IsLabelsMatch(state.option.labels) && (!state.learnerOnly || replica.peer.Role == metapb.PeerRole_Learner)) &&
-				// And If the leader store is abnormal to be accessed under `ReplicaReadPreferLeader` mode, we should choose other valid followers
-				// as candidates to serve the Read request.
-				(!state.option.preferLeader || !replica.store.isSlow()))
+	// the epoch is staled or retry exhausted.
+	if replica.isEpochStale() || replica.isExhausted(1) {
+		return false
+	}
+	// The request can only be sent to the leader.
+	if state.option.leaderOnly && idx == state.leaderIdx {
+		return true
+	}
+	// Choose a replica with matched labels.
+	followerCandidate := !state.option.leaderOnly && (state.tryLeader || idx != state.leaderIdx) &&
+		replica.store.IsLabelsMatch(state.option.labels) && (!state.learnerOnly || replica.peer.Role == metapb.PeerRole_Learner)
+	if !followerCandidate {
+		return false
+	}
+	// And If the leader store is abnormal to be accessed under `ReplicaReadPreferLeader` mode, we should choose other valid followers
+	// as candidates to serve the Read request.
+	if state.option.preferLeader && replica.store.isSlow() {
+		return false
+	}
+	// If the stores are limited, check if the store is in the list.
+	return replica.store.IsStoreMatch(state.option.stores)
 }
 
 // tryIdleReplica is the state where we find the leader is busy and retry the request using replica read.
@@ -644,7 +657,7 @@ func (state *tryIdleReplica) next(bo *retry.Backoffer, selector *replicaSelector
 		if r.isExhausted(1) {
 			continue
 		}
-		estimated := r.store.estimatedWaitTime()
+		estimated := r.store.EstimatedWaitTime()
 		if estimated > selector.busyThreshold {
 			continue
 		}
@@ -950,7 +963,7 @@ func (s *replicaSelector) updateLeader(leader *metapb.Peer) {
 	s.region.invalidate(StoreNotFound)
 }
 
-func (s *replicaSelector) onServerIsBusy(bo *retry.Backoffer, ctx *RPCContext, serverIsBusy *errorpb.ServerIsBusy) (shouldRetry bool, err error) {
+func (s *replicaSelector) onServerIsBusy(bo *retry.Backoffer, ctx *RPCContext, req *tikvrpc.Request, serverIsBusy *errorpb.ServerIsBusy) (shouldRetry bool, err error) {
 	if serverIsBusy.EstimatedWaitMs != 0 && ctx != nil && ctx.Store != nil {
 		estimatedWait := time.Duration(serverIsBusy.EstimatedWaitMs) * time.Millisecond
 		// Update the estimated wait time of the store.
@@ -961,6 +974,11 @@ func (s *replicaSelector) onServerIsBusy(bo *retry.Backoffer, ctx *RPCContext, s
 		ctx.Store.loadStats.Store(loadStats)
 
 		if s.busyThreshold != 0 {
+			// do not retry with batched coprocessor requests.
+			// it'll be region misses if we send the tasks to replica.
+			if req.Type == tikvrpc.CmdCop && len(req.Cop().Tasks) > 0 {
+				return false, nil
+			}
 			switch state := s.state.(type) {
 			case *accessKnownLeader:
 				// Clear attempt history of the leader, so the leader can be accessed again.
@@ -1668,7 +1686,7 @@ func (s *RegionRequestSender) onRegionError(bo *retry.Backoffer, ctx *RPCContext
 
 	if serverIsBusy := regionErr.GetServerIsBusy(); serverIsBusy != nil {
 		if s.replicaSelector != nil {
-			return s.replicaSelector.onServerIsBusy(bo, ctx, serverIsBusy)
+			return s.replicaSelector.onServerIsBusy(bo, ctx, req, serverIsBusy)
 		}
 		logutil.BgLogger().Warn("tikv reports `ServerIsBusy` retry later",
 			zap.String("reason", regionErr.GetServerIsBusy().GetReason()),
diff --git a/tikv/region.go b/tikv/region.go
index dd7500c..6b5e487 100644
--- a/tikv/region.go
+++ b/tikv/region.go
@@ -184,6 +184,11 @@ func WithMatchLabels(labels []*metapb.StoreLabel) StoreSelectorOption {
 	return locate.WithMatchLabels(labels)
 }
 
+// WithMatchStores indicates selecting stores with matched store ids.
+func WithMatchStores(stores []uint64) StoreSelectorOption {
+	return locate.WithMatchStores(stores)
+}
+
 // NewRegionRequestRuntimeStats returns a new RegionRequestRuntimeStats.
 func NewRegionRequestRuntimeStats() RegionRequestRuntimeStats {
 	return locate.NewRegionRequestRuntimeStats()


Test patch:
diff --git a/internal/locate/region_request3_test.go b/internal/locate/region_request3_test.go
index c1950e1..a148329 100644
--- a/internal/locate/region_request3_test.go
+++ b/internal/locate/region_request3_test.go
@@ -952,7 +952,7 @@ func (s *testRegionRequestToThreeStoresSuite) TestLoadBasedReplicaRead() {
 	s.Equal(rpcCtx.Peer.Id, s.leaderPeer)
 
 	// Receive a ServerIsBusy error
-	replicaSelector.onServerIsBusy(bo, rpcCtx, &errorpb.ServerIsBusy{
+	replicaSelector.onServerIsBusy(bo, rpcCtx, req, &errorpb.ServerIsBusy{
 		EstimatedWaitMs: 500,
 	})
 
@@ -963,7 +963,7 @@ func (s *testRegionRequestToThreeStoresSuite) TestLoadBasedReplicaRead() {
 	s.True(*rpcCtx.contextPatcher.replicaRead)
 	lastPeerID := rpcCtx.Peer.Id
 
-	replicaSelector.onServerIsBusy(bo, rpcCtx, &errorpb.ServerIsBusy{
+	replicaSelector.onServerIsBusy(bo, rpcCtx, req, &errorpb.ServerIsBusy{
 		EstimatedWaitMs: 800,
 	})
 
@@ -976,7 +976,7 @@ func (s *testRegionRequestToThreeStoresSuite) TestLoadBasedReplicaRead() {
 	s.True(*rpcCtx.contextPatcher.replicaRead)
 
 	// All peers are too busy
-	replicaSelector.onServerIsBusy(bo, rpcCtx, &errorpb.ServerIsBusy{
+	replicaSelector.onServerIsBusy(bo, rpcCtx, req, &errorpb.ServerIsBusy{
 		EstimatedWaitMs: 150,
 	})
 	lessBusyPeer := rpcCtx.Peer.Id