Skip to content

Commit

Permalink
Split the keys to search
Browse files Browse the repository at this point in the history 
Signed-off-by: JmPotato <[email protected]>
  • Loading branch information
@JmPotato
JmPotato committed Jan 17, 2025
1 parent 287a1a3 commit 665cbfa
Show file tree
Hide file tree
Showing 8 changed files with 47 additions and 26 deletions.
2 changes: 1 addition & 1 deletion go.mod
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions go.sum
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down
38 changes: 32 additions & 6 deletions pkg/core/region.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
Expand Down Expand Up @@ -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.
Expand Down
17 changes: 6 additions & 11 deletions server/grpc_service.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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())
}
Expand Down
2 changes: 1 addition & 1 deletion tests/integrations/go.mod
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions tests/integrations/go.sum
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down
2 changes: 1 addition & 1 deletion tools/go.mod
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
Expand Down
4 changes: 2 additions & 2 deletions tools/go.sum
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down

0 comments on commit 665cbfa

Please sign in to comment.