From 665cbfa5e0d81dfa3a293c9b45ad471e2e7ed653 Mon Sep 17 00:00:00 2001 From: JmPotato Date: Fri, 17 Jan 2025 11:11:40 +0800 Subject: [PATCH] Split the keys to search Signed-off-by: JmPotato --- go.mod | 2 +- go.sum | 4 ++-- pkg/core/region.go | 38 ++++++++++++++++++++++++++++++++------ server/grpc_service.go | 17 ++++++----------- tests/integrations/go.mod | 2 +- tests/integrations/go.sum | 4 ++-- tools/go.mod | 2 +- tools/go.sum | 4 ++-- 8 files changed, 47 insertions(+), 26 deletions(-) diff --git a/go.mod b/go.mod index 44a34dd3cc3..0d151362a0c 100644 --- a/go.mod +++ b/go.mod @@ -5,7 +5,7 @@ go 1.23 // When you modify PD cooperatively with kvproto, this will be useful to submit the PR to PD and the PR to // kvproto at the same time. You can run `go mod tidy` to make it replaced with go-mod style specification. // After the PR to kvproto is merged, remember to comment this out and run `go mod tidy`. -replace github.com/pingcap/kvproto => github.com/JmPotato/kvproto v0.0.0-20250107061039-28ba77c8a2bc +replace github.com/pingcap/kvproto => github.com/JmPotato/kvproto v0.0.0-20250117041351-579bbd411428 require ( github.com/AlekSi/gocov-xml v1.0.0 diff --git a/go.sum b/go.sum index 730949974c2..b8d0f5670d5 100644 --- a/go.sum +++ b/go.sum @@ -818,8 +818,8 @@ github.com/AlekSi/gocov-xml v1.0.0/go.mod h1:J0qYeZ6tDg4oZubW9mAAgxlqw39PDfoEkzB github.com/BurntSushi/toml v0.3.1 h1:WXkYYl6Yr3qBf1K79EBnL4mak0OimBfB0XUf9Vl28OQ= github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU= github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym/WlBOVXweHU+Q+/VP0lqqI8lqeDx9IjBqo= -github.com/JmPotato/kvproto v0.0.0-20250107061039-28ba77c8a2bc h1:0PSQWfSdl+PEbEtTH9drz+LgfJw2PimLrIqeANjLtCg= -github.com/JmPotato/kvproto v0.0.0-20250107061039-28ba77c8a2bc/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8= +github.com/JmPotato/kvproto v0.0.0-20250117041351-579bbd411428 h1:DB94vRBY1YKZPIrRyyoH1IcgpkTdUF/i7BoZ4k7/6nM= +github.com/JmPotato/kvproto v0.0.0-20250117041351-579bbd411428/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8= github.com/JohnCGriffin/overflow v0.0.0-20211019200055-46fa312c352c/go.mod h1:X0CRv0ky0k6m906ixxpzmDRLvX58TFUKS2eePweuyxk= github.com/KyleBanks/depth v1.2.1 h1:5h8fQADFrWtarTdtDudMmGsC7GPbOAu6RVB3ffsVFHc= github.com/KyleBanks/depth v1.2.1/go.mod h1:jzSb9d0L43HxTQfT+oSA1EEp2q+ne2uh6XgeJcm8brE= 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()) } diff --git a/tests/integrations/go.mod b/tests/integrations/go.mod index cfa9e48de38..4898e2761ec 100644 --- a/tests/integrations/go.mod +++ b/tests/integrations/go.mod @@ -3,7 +3,7 @@ module github.com/tikv/pd/tests/integrations go 1.23 replace ( - github.com/pingcap/kvproto => github.com/JmPotato/kvproto v0.0.0-20250107061039-28ba77c8a2bc + github.com/pingcap/kvproto => github.com/JmPotato/kvproto v0.0.0-20250117041351-579bbd411428 github.com/tikv/pd => ../../ github.com/tikv/pd/client => ../../client github.com/tikv/pd/tests/integrations/mcs => ./mcs diff --git a/tests/integrations/go.sum b/tests/integrations/go.sum index 4c581ee7473..5ca38267e3c 100644 --- a/tests/integrations/go.sum +++ b/tests/integrations/go.sum @@ -818,8 +818,8 @@ github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03 github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym/WlBOVXweHU+Q+/VP0lqqI8lqeDx9IjBqo= github.com/DATA-DOG/go-sqlmock v1.5.0 h1:Shsta01QNfFxHCfpW6YH2STWB0MudeXXEWMr20OEh60= github.com/DATA-DOG/go-sqlmock v1.5.0/go.mod h1:f/Ixk793poVmq4qj/V1dPUg2JEAKC73Q5eFN3EC/SaM= -github.com/JmPotato/kvproto v0.0.0-20250107061039-28ba77c8a2bc h1:0PSQWfSdl+PEbEtTH9drz+LgfJw2PimLrIqeANjLtCg= -github.com/JmPotato/kvproto v0.0.0-20250107061039-28ba77c8a2bc/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8= +github.com/JmPotato/kvproto v0.0.0-20250117041351-579bbd411428 h1:DB94vRBY1YKZPIrRyyoH1IcgpkTdUF/i7BoZ4k7/6nM= +github.com/JmPotato/kvproto v0.0.0-20250117041351-579bbd411428/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8= github.com/JohnCGriffin/overflow v0.0.0-20211019200055-46fa312c352c/go.mod h1:X0CRv0ky0k6m906ixxpzmDRLvX58TFUKS2eePweuyxk= github.com/KyleBanks/depth v1.2.1 h1:5h8fQADFrWtarTdtDudMmGsC7GPbOAu6RVB3ffsVFHc= github.com/KyleBanks/depth v1.2.1/go.mod h1:jzSb9d0L43HxTQfT+oSA1EEp2q+ne2uh6XgeJcm8brE= diff --git a/tools/go.mod b/tools/go.mod index ec8489ee356..b25e5e9f50d 100644 --- a/tools/go.mod +++ b/tools/go.mod @@ -3,7 +3,7 @@ module github.com/tikv/pd/tools go 1.23 replace ( - github.com/pingcap/kvproto => github.com/JmPotato/kvproto v0.0.0-20250107061039-28ba77c8a2bc + github.com/pingcap/kvproto => github.com/JmPotato/kvproto v0.0.0-20250117041351-579bbd411428 github.com/tikv/pd => ../ github.com/tikv/pd/client => ../client ) diff --git a/tools/go.sum b/tools/go.sum index d08c667b1d1..e3cf1f3417d 100644 --- a/tools/go.sum +++ b/tools/go.sum @@ -816,8 +816,8 @@ git.sr.ht/~sbinet/gg v0.3.1/go.mod h1:KGYtlADtqsqANL9ueOFkWymvzUvLMQllU5Ixo+8v3p github.com/BurntSushi/toml v0.3.1 h1:WXkYYl6Yr3qBf1K79EBnL4mak0OimBfB0XUf9Vl28OQ= github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU= github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym/WlBOVXweHU+Q+/VP0lqqI8lqeDx9IjBqo= -github.com/JmPotato/kvproto v0.0.0-20250107061039-28ba77c8a2bc h1:0PSQWfSdl+PEbEtTH9drz+LgfJw2PimLrIqeANjLtCg= -github.com/JmPotato/kvproto v0.0.0-20250107061039-28ba77c8a2bc/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8= +github.com/JmPotato/kvproto v0.0.0-20250117041351-579bbd411428 h1:DB94vRBY1YKZPIrRyyoH1IcgpkTdUF/i7BoZ4k7/6nM= +github.com/JmPotato/kvproto v0.0.0-20250117041351-579bbd411428/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8= github.com/JohnCGriffin/overflow v0.0.0-20211019200055-46fa312c352c/go.mod h1:X0CRv0ky0k6m906ixxpzmDRLvX58TFUKS2eePweuyxk= github.com/KyleBanks/depth v1.2.1 h1:5h8fQADFrWtarTdtDudMmGsC7GPbOAu6RVB3ffsVFHc= github.com/KyleBanks/depth v1.2.1/go.mod h1:jzSb9d0L43HxTQfT+oSA1EEp2q+ne2uh6XgeJcm8brE=