From 1118cec6ab7a99caf97cb24903813de3a4cd96f7 Mon Sep 17 00:00:00 2001
From: JmPotato <ghzpotato@gmail.com>
Date: Fri, 17 Jan 2025 11:11:40 +0800
Subject: [PATCH] Split the keys to search

Signed-off-by: JmPotato <ghzpotato@gmail.com>
---
 pkg/core/region.go     | 38 ++++++++++++++++++++++++++++++++------
 server/grpc_service.go | 17 ++++++-----------
 2 files changed, 38 insertions(+), 17 deletions(-)

diff --git a/pkg/core/region.go b/pkg/core/region.go
index b4ebede7722..94fc525f11b 100644
--- a/pkg/core/region.go
+++ b/pkg/core/region.go
@@ -47,6 +47,7 @@ import (
 const (
 	randomRegionMaxRetry = 10
 	scanRegionLimit      = 1000
+	batchSearchSize      = 16
 	// CollectFactor is the factor to collect the count of region.
 	CollectFactor = 0.9
 )
@@ -1512,15 +1513,40 @@ func (r *RegionsInfo) QueryRegions(
 
 // getRegionsByKeys searches RegionInfo from regionTree by keys.
 func (r *RegionsInfo) getRegionsByKeys(keys [][]byte) []*RegionInfo {
-	r.t.RLock()
-	defer r.t.RUnlock()
-	return r.tree.searchByKeys(keys)
+	regions := make([]*RegionInfo, 0, len(keys))
+	// Split the keys into multiple batches, and search each batch separately.
+	// This is to avoid the lock contention on the `regionTree`.
+	for _, batch := range splitKeysIntoBatches(keys) {
+		r.t.RLock()
+		results := r.tree.searchByKeys(batch)
+		r.t.RUnlock()
+		regions = append(regions, results...)
+	}
+	return regions
+}
+
+func splitKeysIntoBatches(keys [][]byte) [][][]byte {
+	keysLen := len(keys)
+	batches := make([][][]byte, 0, (keysLen+batchSearchSize-1)/batchSearchSize)
+	for i := 0; i < keysLen; i += batchSearchSize {
+		end := i + batchSearchSize
+		if end > keysLen {
+			end = keysLen
+		}
+		batches = append(batches, keys[i:end])
+	}
+	return batches
 }
 
 func (r *RegionsInfo) getRegionsByPrevKeys(prevKeys [][]byte) []*RegionInfo {
-	r.t.RLock()
-	defer r.t.RUnlock()
-	return r.tree.searchByPrevKeys(prevKeys)
+	regions := make([]*RegionInfo, 0, len(prevKeys))
+	for _, batch := range splitKeysIntoBatches(prevKeys) {
+		r.t.RLock()
+		results := r.tree.searchByPrevKeys(batch)
+		r.t.RUnlock()
+		regions = append(regions, results...)
+	}
+	return regions
 }
 
 // sortOutKeyIDMap will iterate the regions, convert it to a slice of regionID that corresponds to the input regions.
diff --git a/server/grpc_service.go b/server/grpc_service.go
index 77ac9ddddc6..953ff1cfbc9 100644
--- a/server/grpc_service.go
+++ b/server/grpc_service.go
@@ -1532,14 +1532,12 @@ func (s *GrpcServer) GetRegionByID(ctx context.Context, request *pdpb.GetRegionB
 
 // QueryRegion provides a stream processing of the region query.
 func (s *GrpcServer) QueryRegion(stream pdpb.PD_QueryRegionServer) error {
-	if s.GetServiceMiddlewarePersistOptions().IsGRPCRateLimitEnabled() {
-		fName := currentFunction()
-		limiter := s.GetGRPCRateLimiter()
-		if done, err := limiter.Allow(fName); err == nil {
-			defer done()
-		} else {
-			return err
-		}
+	done, err := s.rateLimitCheck()
+	if err != nil {
+		return err
+	}
+	if done != nil {
+		defer done()
 	}
 
 	for {
@@ -1553,9 +1551,6 @@ func (s *GrpcServer) QueryRegion(stream pdpb.PD_QueryRegionServer) error {
 
 		// TODO: add forwarding logic.
 
-		if s.IsClosed() {
-			return errs.ErrNotStarted
-		}
 		if clusterID := keypath.ClusterID(); request.GetHeader().GetClusterId() != clusterID {
 			return errs.ErrMismatchClusterID(clusterID, request.GetHeader().GetClusterId())
 		}