Count only one tiflash replica in disaggregated mode. (tikv#805)

Signed-off-by: iosmanthus <myosmanthustree@gmail.com>
iosmanthus · Dec 21, 2023 · 091227e · 091227e
1 parent 4be5abd
commit 091227e
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 7 deletions.
diff --git a/internal/locate/region_request.go b/internal/locate/region_request.go
@@ -1608,6 +1608,32 @@ func fetchRespInfo(resp *tikvrpc.Response) string {
 	return extraInfo
 }
 
+func (s *RegionRequestSender) countReplicaNumber(peers []*metapb.Peer) int {
+	isTiFlashWriteNode := func(storeId uint64) bool {
+		store := s.regionCache.getStoreByStoreID(storeId)
+		engine, _ := store.GetLabelValue("engine")
+		engineRole, _ := store.GetLabelValue("engine_role")
+		return engine == "tiflash" && engineRole == "write"
+	}
+	// In disaggregated-mode(tiflash write-node), only count 1 replica for tiflash, no matter how many tiflash write-nodes there are.
+	replicaNumber := 0
+	hasTiFlashWriteNode := false
+	for _, peer := range peers {
+		role := peer.GetRole()
+		if role == metapb.PeerRole_Voter {
+			replicaNumber++
+		} else if role == metapb.PeerRole_Learner {
+			if !isTiFlashWriteNode(peer.StoreId) {
+				replicaNumber++
+			} else if !hasTiFlashWriteNode {
+				hasTiFlashWriteNode = true
+				replicaNumber++
+			}
+		}
+	}
+	return replicaNumber
+}
+
 func (s *RegionRequestSender) sendReqToRegion(
 	bo *retry.Backoffer, rpcCtx *RPCContext, req *tikvrpc.Request, timeout time.Duration,
 ) (resp *tikvrpc.Response, retry bool, err error) {
@@ -1642,13 +1668,7 @@ func (s *RegionRequestSender) sendReqToRegion(
 	// Count the replica number as the RU cost factor.
 	req.ReplicaNumber = 1
 	if rpcCtx.Meta != nil && len(rpcCtx.Meta.GetPeers()) > 0 {
-		req.ReplicaNumber = 0
-		for _, peer := range rpcCtx.Meta.GetPeers() {
-			role := peer.GetRole()
-			if role == metapb.PeerRole_Voter || role == metapb.PeerRole_Learner {
-				req.ReplicaNumber++
-			}
-		}
+		req.ReplicaNumber = int64(s.countReplicaNumber(rpcCtx.Meta.GetPeers()))
 	}
 
 	var sessionID uint64

diff --git a/internal/locate/region_request_test.go b/internal/locate/region_request_test.go
@@ -784,3 +784,43 @@ func (s *testRegionRequestToSingleStoreSuite) TestBatchClientSendLoopPanic() {
 	// batchSendLoop should not panic.
 	s.Equal(atomic.LoadInt64(&client.BatchSendLoopPanicCounter), int64(0))
 }
+
+func (s *testRegionRequestToSingleStoreSuite) TestCountReplicaNumber() {
+	fmt.Println("TestCountReplicaNumber")
+	fmt.Println(s.cache.storeMu.stores)
+	tikvLabels := []*metapb.StoreLabel{{Key: "engine", Value: "tikv"}}
+	tiflashLabels := []*metapb.StoreLabel{{Key: "engine", Value: "tiflash"}}
+	tiflashWNLabels := []*metapb.StoreLabel{{Key: "engine", Value: "tiflash"}, {Key: "engine_role", Value: "write"}}
+
+	s.cache.SetRegionCacheStore(1, "", "", tikvrpc.TiKV, 0, tikvLabels)
+	s.cache.SetRegionCacheStore(2, "", "", tikvrpc.TiKV, 0, tikvLabels)
+	s.cache.SetRegionCacheStore(3, "", "", tikvrpc.TiKV, 0, tikvLabels)
+	s.cache.SetRegionCacheStore(4, "", "", tikvrpc.TiFlash, 0, tiflashLabels)
+	s.cache.SetRegionCacheStore(5, "", "", tikvrpc.TiFlash, 0, tiflashLabels)
+	{
+		peers := []*metapb.Peer{{StoreId: 1, Role: metapb.PeerRole_Voter}}
+		s.Equal(1, s.regionRequestSender.countReplicaNumber(peers))
+		peers = append(peers, &metapb.Peer{StoreId: 2, Role: metapb.PeerRole_Voter})
+		s.Equal(2, s.regionRequestSender.countReplicaNumber(peers))
+		peers = append(peers, &metapb.Peer{StoreId: 3, Role: metapb.PeerRole_Voter})
+		s.Equal(3, s.regionRequestSender.countReplicaNumber(peers))
+		peers = append(peers, &metapb.Peer{StoreId: 4, Role: metapb.PeerRole_Learner})
+		s.Equal(4, s.regionRequestSender.countReplicaNumber(peers))
+		peers = append(peers, &metapb.Peer{StoreId: 5, Role: metapb.PeerRole_Learner})
+		s.Equal(5, s.regionRequestSender.countReplicaNumber(peers))
+	}
+	s.cache.SetRegionCacheStore(4, "", "", tikvrpc.TiFlash, 0, tiflashWNLabels)
+	s.cache.SetRegionCacheStore(5, "", "", tikvrpc.TiFlash, 0, tiflashWNLabels)
+	{
+		peers := []*metapb.Peer{{StoreId: 1, Role: metapb.PeerRole_Voter}}
+		s.Equal(1, s.regionRequestSender.countReplicaNumber(peers))
+		peers = append(peers, &metapb.Peer{StoreId: 2, Role: metapb.PeerRole_Voter})
+		s.Equal(2, s.regionRequestSender.countReplicaNumber(peers))
+		peers = append(peers, &metapb.Peer{StoreId: 3, Role: metapb.PeerRole_Voter})
+		s.Equal(3, s.regionRequestSender.countReplicaNumber(peers))
+		peers = append(peers, &metapb.Peer{StoreId: 4, Role: metapb.PeerRole_Learner})
+		s.Equal(4, s.regionRequestSender.countReplicaNumber(peers))
+		peers = append(peers, &metapb.Peer{StoreId: 5, Role: metapb.PeerRole_Learner})
+		s.Equal(4, s.regionRequestSender.countReplicaNumber(peers)) // Only count 1 tiflash replica for tiflash write-nodes.
+	}
+}