From f83fc7b263b3bad4c2a192269e63f84676be3874 Mon Sep 17 00:00:00 2001 From: JmPotato Date: Tue, 21 May 2024 15:59:25 +0800 Subject: [PATCH 1/5] Optimize the efficiency of random regions selecting Signed-off-by: JmPotato --- pkg/core/region_test.go | 16 ++++-- pkg/core/region_tree.go | 106 ++++++++++++++++++++--------------- pkg/core/region_tree_test.go | 2 + 3 files changed, 74 insertions(+), 50 deletions(-) diff --git a/pkg/core/region_test.go b/pkg/core/region_test.go index b09c1dfd601..f86144a9297 100644 --- a/pkg/core/region_test.go +++ b/pkg/core/region_test.go @@ -654,10 +654,18 @@ func BenchmarkRandomRegion(b *testing.B) { origin, overlaps, rangeChanged := regions.SetRegion(region) regions.UpdateSubTree(region, origin, overlaps, rangeChanged) } - b.ResetTimer() - for i := 0; i < b.N; i++ { - regions.RandLeaderRegion(1, nil) - } + b.Run("random region", func(b *testing.B) { + b.ResetTimer() + for i := 0; i < b.N; i++ { + regions.RandLeaderRegion(1, nil) + } + }) + b.Run("random regions", func(b *testing.B) { + b.ResetTimer() + for i := 0; i < b.N; i++ { + regions.RandLeaderRegions(1, nil) + } + }) } func BenchmarkRandomSetRegion(b *testing.B) { diff --git a/pkg/core/region_tree.go b/pkg/core/region_tree.go index 6c3c71c5158..28eaf70f548 100644 --- a/pkg/core/region_tree.go +++ b/pkg/core/region_tree.go @@ -328,62 +328,76 @@ func (t *regionTree) getAdjacentItem(item *regionItem) (prev *regionItem, next * return prev, next } -// RandomRegion is used to get a random region within ranges. +// RandomRegion returns a random region within the given ranges. func (t *regionTree) RandomRegion(ranges []KeyRange) *RegionInfo { - if t.length() == 0 { + regions := t.RandomRegions(1, ranges) + if len(regions) == 0 { return nil } + return regions[0] +} +// RandomRegions get n random regions within the given ranges. +func (t *regionTree) RandomRegions(n int, ranges []KeyRange) []*RegionInfo { + if t.length() == 0 || n < 1 { + return nil + } if len(ranges) == 0 { ranges = []KeyRange{NewKeyRange("", "")} } + var ( + startKey, endKey []byte + startIndex, endIndex, randIdx int + startItem *regionItem + pivotItem = ®ionItem{&RegionInfo{meta: &metapb.Region{}}} + region *RegionInfo + regions = make([]*RegionInfo, 0, n) + rangeNum, curLen = len(ranges), len(regions) + ) + // Keep retrying until we get enough regions. + for curLen < n { + // Shuffle the ranges to increase the randomness. + rand.Shuffle(rangeNum, func(i, j int) { + ranges[i], ranges[j] = ranges[j], ranges[i] + }) + for _, r := range ranges { + startKey, endKey = r.StartKey, r.EndKey + pivotItem.meta.StartKey = startKey + startItem, startIndex = t.tree.GetWithIndex(pivotItem) + if len(endKey) != 0 { + pivotItem.meta.StartKey = endKey + _, endIndex = t.tree.GetWithIndex(pivotItem) + } else { + endIndex = t.tree.Len() + } + // Consider that the item in the tree may not be continuous, + // we need to check if the previous item contains the key. + if startIndex != 0 && startItem == nil && t.tree.GetAt(startIndex-1).Contains(startKey) { + startIndex-- + } + if endIndex <= startIndex { + if len(endKey) > 0 && bytes.Compare(startKey, endKey) > 0 { + log.Error("wrong range keys", + logutil.ZapRedactString("start-key", string(HexRegionKey(startKey))), + logutil.ZapRedactString("end-key", string(HexRegionKey(endKey))), + errs.ZapError(errs.ErrWrongRangeKeys)) + } + continue + } - for _, i := range rand.Perm(len(ranges)) { - var endIndex int - startKey, endKey := ranges[i].StartKey, ranges[i].EndKey - startRegion, startIndex := t.tree.GetWithIndex(®ionItem{RegionInfo: &RegionInfo{meta: &metapb.Region{StartKey: startKey}}}) - - if len(endKey) != 0 { - _, endIndex = t.tree.GetWithIndex(®ionItem{RegionInfo: &RegionInfo{meta: &metapb.Region{StartKey: endKey}}}) - } else { - endIndex = t.tree.Len() - } - - // Consider that the item in the tree may not be continuous, - // we need to check if the previous item contains the key. - if startIndex != 0 && startRegion == nil && t.tree.GetAt(startIndex-1).Contains(startKey) { - startIndex-- - } - - if endIndex <= startIndex { - if len(endKey) > 0 && bytes.Compare(startKey, endKey) > 0 { - log.Error("wrong range keys", - logutil.ZapRedactString("start-key", string(HexRegionKey(startKey))), - logutil.ZapRedactString("end-key", string(HexRegionKey(endKey))), - errs.ZapError(errs.ErrWrongRangeKeys)) + randIdx = rand.Intn(endIndex-startIndex) + startIndex + region = t.tree.GetAt(randIdx).RegionInfo + if curLen < n && region.isInvolved(startKey, endKey) { + regions = append(regions, region) + curLen++ + } else if curLen == n { + return regions } - continue - } - index := rand.Intn(endIndex-startIndex) + startIndex - region := t.tree.GetAt(index).RegionInfo - if region.isInvolved(startKey, endKey) { - return region } - } - - return nil -} - -func (t *regionTree) RandomRegions(n int, ranges []KeyRange) []*RegionInfo { - if t.length() == 0 { - return nil - } - - regions := make([]*RegionInfo, 0, n) - - for i := 0; i < n; i++ { - if region := t.RandomRegion(ranges); region != nil { - regions = append(regions, region) + // No region found in the given ranges, directly break. + // This is to avoid infinite loop. + if len(regions) == 0 { + break } } return regions diff --git a/pkg/core/region_tree_test.go b/pkg/core/region_tree_test.go index 3f2ca0c1fb8..0148b7d5317 100644 --- a/pkg/core/region_tree_test.go +++ b/pkg/core/region_tree_test.go @@ -281,6 +281,8 @@ func TestRandomRegion(t *testing.T) { updateNewItem(tree, regionA) ra := tree.RandomRegion([]KeyRange{NewKeyRange("", "")}) re.Equal(regionA, ra) + ra2 := tree.RandomRegions(2, []KeyRange{NewKeyRange("", "")}) + re.Equal([]*RegionInfo{regionA, regionA}, ra2) regionB := NewTestRegionInfo(2, 2, []byte("g"), []byte("n")) regionC := NewTestRegionInfo(3, 3, []byte("n"), []byte("t")) From aea156a4d60f50f66c9682445fca10d484b810e2 Mon Sep 17 00:00:00 2001 From: JmPotato Date: Tue, 21 May 2024 18:03:46 +0800 Subject: [PATCH 2/5] Remove the unused code Signed-off-by: JmPotato --- pkg/core/region.go | 33 ++------------------------ pkg/core/region_test.go | 2 +- pkg/core/region_tree.go | 3 +-- pkg/core/region_tree_test.go | 46 ++++++++++++++++++------------------ 4 files changed, 27 insertions(+), 57 deletions(-) diff --git a/pkg/core/region.go b/pkg/core/region.go index a1a61d505a9..f0e9fec1cdd 100644 --- a/pkg/core/region.go +++ b/pkg/core/region.go @@ -1673,13 +1673,6 @@ func (r *RegionsInfo) GetStoreWitnessCount(storeID uint64) int { return r.witnesses[storeID].length() } -// RandPendingRegion randomly gets a store's region with a pending peer. -func (r *RegionsInfo) RandPendingRegion(storeID uint64, ranges []KeyRange) *RegionInfo { - r.st.RLock() - defer r.st.RUnlock() - return r.pendingPeers[storeID].RandomRegion(ranges) -} - // RandPendingRegions randomly gets a store's n regions with a pending peer. func (r *RegionsInfo) RandPendingRegions(storeID uint64, ranges []KeyRange) []*RegionInfo { r.st.RLock() @@ -1687,11 +1680,10 @@ func (r *RegionsInfo) RandPendingRegions(storeID uint64, ranges []KeyRange) []*R return r.pendingPeers[storeID].RandomRegions(randomRegionMaxRetry, ranges) } -// RandLeaderRegion randomly gets a store's leader region. -func (r *RegionsInfo) RandLeaderRegion(storeID uint64, ranges []KeyRange) *RegionInfo { +func (r *RegionsInfo) randLeaderRegion(storeID uint64, ranges []KeyRange) *RegionInfo { r.st.RLock() defer r.st.RUnlock() - return r.leaders[storeID].RandomRegion(ranges) + return r.leaders[storeID].randomRegion(ranges) } // RandLeaderRegions randomly gets a store's n leader regions. @@ -1701,13 +1693,6 @@ func (r *RegionsInfo) RandLeaderRegions(storeID uint64, ranges []KeyRange) []*Re return r.leaders[storeID].RandomRegions(randomRegionMaxRetry, ranges) } -// RandFollowerRegion randomly gets a store's follower region. -func (r *RegionsInfo) RandFollowerRegion(storeID uint64, ranges []KeyRange) *RegionInfo { - r.st.RLock() - defer r.st.RUnlock() - return r.followers[storeID].RandomRegion(ranges) -} - // RandFollowerRegions randomly gets a store's n follower regions. func (r *RegionsInfo) RandFollowerRegions(storeID uint64, ranges []KeyRange) []*RegionInfo { r.st.RLock() @@ -1715,13 +1700,6 @@ func (r *RegionsInfo) RandFollowerRegions(storeID uint64, ranges []KeyRange) []* return r.followers[storeID].RandomRegions(randomRegionMaxRetry, ranges) } -// RandLearnerRegion randomly gets a store's learner region. -func (r *RegionsInfo) RandLearnerRegion(storeID uint64, ranges []KeyRange) *RegionInfo { - r.st.RLock() - defer r.st.RUnlock() - return r.learners[storeID].RandomRegion(ranges) -} - // RandLearnerRegions randomly gets a store's n learner regions. func (r *RegionsInfo) RandLearnerRegions(storeID uint64, ranges []KeyRange) []*RegionInfo { r.st.RLock() @@ -1729,13 +1707,6 @@ func (r *RegionsInfo) RandLearnerRegions(storeID uint64, ranges []KeyRange) []*R return r.learners[storeID].RandomRegions(randomRegionMaxRetry, ranges) } -// RandWitnessRegion randomly gets a store's witness region. -func (r *RegionsInfo) RandWitnessRegion(storeID uint64, ranges []KeyRange) *RegionInfo { - r.st.RLock() - defer r.st.RUnlock() - return r.witnesses[storeID].RandomRegion(ranges) -} - // RandWitnessRegions randomly gets a store's n witness regions. func (r *RegionsInfo) RandWitnessRegions(storeID uint64, ranges []KeyRange) []*RegionInfo { r.st.RLock() diff --git a/pkg/core/region_test.go b/pkg/core/region_test.go index f86144a9297..804bc052cc8 100644 --- a/pkg/core/region_test.go +++ b/pkg/core/region_test.go @@ -657,7 +657,7 @@ func BenchmarkRandomRegion(b *testing.B) { b.Run("random region", func(b *testing.B) { b.ResetTimer() for i := 0; i < b.N; i++ { - regions.RandLeaderRegion(1, nil) + regions.randLeaderRegion(1, nil) } }) b.Run("random regions", func(b *testing.B) { diff --git a/pkg/core/region_tree.go b/pkg/core/region_tree.go index 28eaf70f548..40e0988b9c3 100644 --- a/pkg/core/region_tree.go +++ b/pkg/core/region_tree.go @@ -328,8 +328,7 @@ func (t *regionTree) getAdjacentItem(item *regionItem) (prev *regionItem, next * return prev, next } -// RandomRegion returns a random region within the given ranges. -func (t *regionTree) RandomRegion(ranges []KeyRange) *RegionInfo { +func (t *regionTree) randomRegion(ranges []KeyRange) *RegionInfo { regions := t.RandomRegions(1, ranges) if len(regions) == 0 { return nil diff --git a/pkg/core/region_tree_test.go b/pkg/core/region_tree_test.go index 0148b7d5317..a86f2f52c47 100644 --- a/pkg/core/region_tree_test.go +++ b/pkg/core/region_tree_test.go @@ -274,12 +274,12 @@ func TestRegionTreeSplitAndMerge(t *testing.T) { func TestRandomRegion(t *testing.T) { re := require.New(t) tree := newRegionTree() - r := tree.RandomRegion(nil) + r := tree.randomRegion(nil) re.Nil(r) regionA := NewTestRegionInfo(1, 1, []byte(""), []byte("g")) updateNewItem(tree, regionA) - ra := tree.RandomRegion([]KeyRange{NewKeyRange("", "")}) + ra := tree.randomRegion([]KeyRange{NewKeyRange("", "")}) re.Equal(regionA, ra) ra2 := tree.RandomRegions(2, []KeyRange{NewKeyRange("", "")}) re.Equal([]*RegionInfo{regionA, regionA}, ra2) @@ -291,20 +291,20 @@ func TestRandomRegion(t *testing.T) { updateNewItem(tree, regionC) updateNewItem(tree, regionD) - rb := tree.RandomRegion([]KeyRange{NewKeyRange("g", "n")}) + rb := tree.randomRegion([]KeyRange{NewKeyRange("g", "n")}) re.Equal(regionB, rb) - rc := tree.RandomRegion([]KeyRange{NewKeyRange("n", "t")}) + rc := tree.randomRegion([]KeyRange{NewKeyRange("n", "t")}) re.Equal(regionC, rc) - rd := tree.RandomRegion([]KeyRange{NewKeyRange("t", "")}) + rd := tree.randomRegion([]KeyRange{NewKeyRange("t", "")}) re.Equal(regionD, rd) - rf := tree.RandomRegion([]KeyRange{NewKeyRange("", "a")}) + rf := tree.randomRegion([]KeyRange{NewKeyRange("", "a")}) re.Nil(rf) - rf = tree.RandomRegion([]KeyRange{NewKeyRange("o", "s")}) + rf = tree.randomRegion([]KeyRange{NewKeyRange("o", "s")}) re.Nil(rf) - rf = tree.RandomRegion([]KeyRange{NewKeyRange("", "a")}) + rf = tree.randomRegion([]KeyRange{NewKeyRange("", "a")}) re.Nil(rf) - rf = tree.RandomRegion([]KeyRange{NewKeyRange("z", "")}) + rf = tree.randomRegion([]KeyRange{NewKeyRange("z", "")}) re.Nil(rf) checkRandomRegion(re, tree, []*RegionInfo{regionA, regionB, regionC, regionD}, []KeyRange{NewKeyRange("", "")}) @@ -317,43 +317,43 @@ func TestRandomRegion(t *testing.T) { func TestRandomRegionDiscontinuous(t *testing.T) { re := require.New(t) tree := newRegionTree() - r := tree.RandomRegion([]KeyRange{NewKeyRange("c", "f")}) + r := tree.randomRegion([]KeyRange{NewKeyRange("c", "f")}) re.Nil(r) // test for single region regionA := NewTestRegionInfo(1, 1, []byte("c"), []byte("f")) updateNewItem(tree, regionA) - ra := tree.RandomRegion([]KeyRange{NewKeyRange("c", "e")}) + ra := tree.randomRegion([]KeyRange{NewKeyRange("c", "e")}) re.Nil(ra) - ra = tree.RandomRegion([]KeyRange{NewKeyRange("c", "f")}) + ra = tree.randomRegion([]KeyRange{NewKeyRange("c", "f")}) re.Equal(regionA, ra) - ra = tree.RandomRegion([]KeyRange{NewKeyRange("c", "g")}) + ra = tree.randomRegion([]KeyRange{NewKeyRange("c", "g")}) re.Equal(regionA, ra) - ra = tree.RandomRegion([]KeyRange{NewKeyRange("a", "e")}) + ra = tree.randomRegion([]KeyRange{NewKeyRange("a", "e")}) re.Nil(ra) - ra = tree.RandomRegion([]KeyRange{NewKeyRange("a", "f")}) + ra = tree.randomRegion([]KeyRange{NewKeyRange("a", "f")}) re.Equal(regionA, ra) - ra = tree.RandomRegion([]KeyRange{NewKeyRange("a", "g")}) + ra = tree.randomRegion([]KeyRange{NewKeyRange("a", "g")}) re.Equal(regionA, ra) regionB := NewTestRegionInfo(2, 2, []byte("n"), []byte("x")) updateNewItem(tree, regionB) - rb := tree.RandomRegion([]KeyRange{NewKeyRange("g", "x")}) + rb := tree.randomRegion([]KeyRange{NewKeyRange("g", "x")}) re.Equal(regionB, rb) - rb = tree.RandomRegion([]KeyRange{NewKeyRange("g", "y")}) + rb = tree.randomRegion([]KeyRange{NewKeyRange("g", "y")}) re.Equal(regionB, rb) - rb = tree.RandomRegion([]KeyRange{NewKeyRange("n", "y")}) + rb = tree.randomRegion([]KeyRange{NewKeyRange("n", "y")}) re.Equal(regionB, rb) - rb = tree.RandomRegion([]KeyRange{NewKeyRange("o", "y")}) + rb = tree.randomRegion([]KeyRange{NewKeyRange("o", "y")}) re.Nil(rb) regionC := NewTestRegionInfo(3, 3, []byte("z"), []byte("")) updateNewItem(tree, regionC) - rc := tree.RandomRegion([]KeyRange{NewKeyRange("y", "")}) + rc := tree.randomRegion([]KeyRange{NewKeyRange("y", "")}) re.Equal(regionC, rc) regionD := NewTestRegionInfo(4, 4, []byte(""), []byte("a")) updateNewItem(tree, regionD) - rd := tree.RandomRegion([]KeyRange{NewKeyRange("", "b")}) + rd := tree.randomRegion([]KeyRange{NewKeyRange("", "b")}) re.Equal(regionD, rd) checkRandomRegion(re, tree, []*RegionInfo{regionA, regionB, regionC, regionD}, []KeyRange{NewKeyRange("", "")}) @@ -367,7 +367,7 @@ func updateNewItem(tree *regionTree, region *RegionInfo) { func checkRandomRegion(re *require.Assertions, tree *regionTree, regions []*RegionInfo, ranges []KeyRange) { keys := make(map[string]struct{}) for i := 0; i < 10000 && len(keys) < len(regions); i++ { - re := tree.RandomRegion(ranges) + re := tree.randomRegion(ranges) if re == nil { continue } From 7df701c8c80223d39132986398b9e1b7c9ddba13 Mon Sep 17 00:00:00 2001 From: JmPotato Date: Wed, 22 May 2024 11:00:22 +0800 Subject: [PATCH 3/5] Refine the code Signed-off-by: JmPotato --- pkg/core/region_test.go | 64 +++++++++++++++++++++++++++-------------- pkg/core/region_tree.go | 61 +++++++++++++++++++++++++++++---------- 2 files changed, 88 insertions(+), 37 deletions(-) diff --git a/pkg/core/region_test.go b/pkg/core/region_test.go index 804bc052cc8..8956bd8a357 100644 --- a/pkg/core/region_test.go +++ b/pkg/core/region_test.go @@ -642,30 +642,50 @@ func BenchmarkUpdateBuckets(b *testing.B) { } func BenchmarkRandomRegion(b *testing.B) { - regions := NewRegionsInfo() - for i := 0; i < 5000000; i++ { - peer := &metapb.Peer{StoreId: 1, Id: uint64(i + 1)} - region := NewRegionInfo(&metapb.Region{ - Id: uint64(i + 1), - Peers: []*metapb.Peer{peer}, - StartKey: []byte(fmt.Sprintf("%20d", i)), - EndKey: []byte(fmt.Sprintf("%20d", i+1)), - }, peer) - origin, overlaps, rangeChanged := regions.SetRegion(region) - regions.UpdateSubTree(region, origin, overlaps, rangeChanged) - } - b.Run("random region", func(b *testing.B) { - b.ResetTimer() - for i := 0; i < b.N; i++ { - regions.randLeaderRegion(1, nil) + for _, size := range []int{10, 100, 1000, 10000, 100000, 1000000, 10000000} { + regions := NewRegionsInfo() + for i := 0; i < size; i++ { + peer := &metapb.Peer{StoreId: 1, Id: uint64(i + 1)} + region := NewRegionInfo(&metapb.Region{ + Id: uint64(i + 1), + Peers: []*metapb.Peer{peer}, + StartKey: []byte(fmt.Sprintf("%20d", i)), + EndKey: []byte(fmt.Sprintf("%20d", i+1)), + }, peer) + origin, overlaps, rangeChanged := regions.SetRegion(region) + regions.UpdateSubTree(region, origin, overlaps, rangeChanged) } - }) - b.Run("random regions", func(b *testing.B) { - b.ResetTimer() - for i := 0; i < b.N; i++ { - regions.RandLeaderRegions(1, nil) + b.Run(fmt.Sprintf("random region whole range with size %d", size), func(b *testing.B) { + b.ResetTimer() + for i := 0; i < b.N; i++ { + regions.randLeaderRegion(1, nil) + } + }) + b.Run(fmt.Sprintf("random regions whole range with size %d", size), func(b *testing.B) { + b.ResetTimer() + for i := 0; i < b.N; i++ { + regions.RandLeaderRegions(1, nil) + } + }) + ranges := []KeyRange{ + NewKeyRange(fmt.Sprintf("%20d", 0), fmt.Sprintf("%20d", size/4)), + NewKeyRange(fmt.Sprintf("%20d", size/4), fmt.Sprintf("%20d", size/2)), + NewKeyRange(fmt.Sprintf("%20d", size/2), fmt.Sprintf("%20d", size*3/4)), + NewKeyRange(fmt.Sprintf("%20d", size*3/4), fmt.Sprintf("%20d", size)), } - }) + b.Run(fmt.Sprintf("random region given ranges with size %d", size), func(b *testing.B) { + b.ResetTimer() + for i := 0; i < b.N; i++ { + regions.randLeaderRegion(1, ranges) + } + }) + b.Run(fmt.Sprintf("random regions given ranges with size %d", size), func(b *testing.B) { + b.ResetTimer() + for i := 0; i < b.N; i++ { + regions.RandLeaderRegions(1, ranges) + } + }) + } } func BenchmarkRandomSetRegion(b *testing.B) { diff --git a/pkg/core/region_tree.go b/pkg/core/region_tree.go index 40e0988b9c3..9cc19e29c24 100644 --- a/pkg/core/region_tree.go +++ b/pkg/core/region_tree.go @@ -341,9 +341,7 @@ func (t *regionTree) RandomRegions(n int, ranges []KeyRange) []*RegionInfo { if t.length() == 0 || n < 1 { return nil } - if len(ranges) == 0 { - ranges = []KeyRange{NewKeyRange("", "")} - } + // Pre-allocate the variables to reduce the temporary memory allocations. var ( startKey, endKey []byte startIndex, endIndex, randIdx int @@ -352,15 +350,10 @@ func (t *regionTree) RandomRegions(n int, ranges []KeyRange) []*RegionInfo { region *RegionInfo regions = make([]*RegionInfo, 0, n) rangeNum, curLen = len(ranges), len(regions) - ) - // Keep retrying until we get enough regions. - for curLen < n { - // Shuffle the ranges to increase the randomness. - rand.Shuffle(rangeNum, func(i, j int) { - ranges[i], ranges[j] = ranges[j], ranges[i] - }) - for _, r := range ranges { - startKey, endKey = r.StartKey, r.EndKey + // setStartEndIndices is a helper function to set `startIndex` and `endIndex` + // according to the `startKey` and `endKey`. + // TODO: maybe we could cache the `startIndex` and `endIndex` for each range. + setStartEndIndices = func() { pivotItem.meta.StartKey = startKey startItem, startIndex = t.tree.GetWithIndex(pivotItem) if len(endKey) != 0 { @@ -374,6 +367,45 @@ func (t *regionTree) RandomRegions(n int, ranges []KeyRange) []*RegionInfo { if startIndex != 0 && startItem == nil && t.tree.GetAt(startIndex-1).Contains(startKey) { startIndex-- } + } + ) + // If no ranges specified, select randomly from the whole tree. + // This is a fast path to reduce the unnecessary iterations. + if rangeNum == 0 { + startKey, endKey = []byte(""), []byte("") + setStartEndIndices() + if endIndex <= startIndex { + if len(endKey) > 0 && bytes.Compare(startKey, endKey) > 0 { + log.Error("wrong range keys", + logutil.ZapRedactString("start-key", string(HexRegionKey(startKey))), + logutil.ZapRedactString("end-key", string(HexRegionKey(endKey))), + errs.ZapError(errs.ErrWrongRangeKeys)) + } + return regions + } + for curLen < n { + randIdx = rand.Intn(endIndex-startIndex) + startIndex + region = t.tree.GetAt(randIdx).RegionInfo + if curLen < n && region.isInvolved(startKey, endKey) { + regions = append(regions, region) + curLen++ + } else if curLen == n { + break + } + // No region found, directly break to avoid infinite loop. + if curLen == 0 { + break + } + } + return regions + } + // When there are multiple ranges provided, + // keep retrying until we get enough regions. + for curLen < n { + // Shuffle the ranges to increase the randomness. + for _, i := range rand.Perm(rangeNum) { + startKey, endKey = ranges[i].StartKey, ranges[i].EndKey + setStartEndIndices() if endIndex <= startIndex { if len(endKey) > 0 && bytes.Compare(startKey, endKey) > 0 { log.Error("wrong range keys", @@ -393,9 +425,8 @@ func (t *regionTree) RandomRegions(n int, ranges []KeyRange) []*RegionInfo { return regions } } - // No region found in the given ranges, directly break. - // This is to avoid infinite loop. - if len(regions) == 0 { + // No region found, directly break to avoid infinite loop. + if curLen == 0 { break } } From 228db1dfdfb6a53a0ee650820e7353076fff24e5 Mon Sep 17 00:00:00 2001 From: JmPotato Date: Wed, 22 May 2024 11:36:58 +0800 Subject: [PATCH 4/5] Fix the check Signed-off-by: JmPotato --- pkg/core/region.go | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/pkg/core/region.go b/pkg/core/region.go index f0e9fec1cdd..a5e2d0f9d1e 100644 --- a/pkg/core/region.go +++ b/pkg/core/region.go @@ -1680,10 +1680,11 @@ func (r *RegionsInfo) RandPendingRegions(storeID uint64, ranges []KeyRange) []*R return r.pendingPeers[storeID].RandomRegions(randomRegionMaxRetry, ranges) } -func (r *RegionsInfo) randLeaderRegion(storeID uint64, ranges []KeyRange) *RegionInfo { +// This function is used for test only. +func (r *RegionsInfo) randLeaderRegion(storeID uint64, ranges []KeyRange) { r.st.RLock() defer r.st.RUnlock() - return r.leaders[storeID].randomRegion(ranges) + _ = r.leaders[storeID].randomRegion(ranges) } // RandLeaderRegions randomly gets a store's n leader regions. From 12142edefd36048ba17771d3bed644233d076cb7 Mon Sep 17 00:00:00 2001 From: JmPotato Date: Wed, 22 May 2024 15:04:05 +0800 Subject: [PATCH 5/5] Address the comments Signed-off-by: JmPotato --- pkg/core/region.go | 5 ++++ pkg/core/region_tree.go | 63 +++++++++++++++++------------------------ 2 files changed, 31 insertions(+), 37 deletions(-) diff --git a/pkg/core/region.go b/pkg/core/region.go index a5e2d0f9d1e..19c1d0d4794 100644 --- a/pkg/core/region.go +++ b/pkg/core/region.go @@ -727,6 +727,11 @@ func (r *RegionInfo) isRegionRecreated() bool { return r.GetRegionEpoch().GetVersion() == 1 && r.GetRegionEpoch().GetConfVer() == 1 && (len(r.GetStartKey()) != 0 || len(r.GetEndKey()) != 0) } +func (r *RegionInfo) Contains(key []byte) bool { + start, end := r.GetStartKey(), r.GetEndKey() + return bytes.Compare(key, start) >= 0 && (len(end) == 0 || bytes.Compare(key, end) < 0) +} + // RegionGuideFunc is a function that determines which follow-up operations need to be performed based on the origin // and new region information. type RegionGuideFunc func(ctx *MetaProcessContext, region, origin *RegionInfo) (saveKV, saveCache, needSync, retained bool) diff --git a/pkg/core/region_tree.go b/pkg/core/region_tree.go index 9cc19e29c24..e6d05d443a1 100644 --- a/pkg/core/region_tree.go +++ b/pkg/core/region_tree.go @@ -52,11 +52,6 @@ func (r *regionItem) Less(other *regionItem) bool { return bytes.Compare(left, right) < 0 } -func (r *regionItem) Contains(key []byte) bool { - start, end := r.GetStartKey(), r.GetEndKey() - return bytes.Compare(key, start) >= 0 && (len(end) == 0 || bytes.Compare(key, end) < 0) -} - const ( defaultBTreeDegree = 64 ) @@ -338,18 +333,19 @@ func (t *regionTree) randomRegion(ranges []KeyRange) *RegionInfo { // RandomRegions get n random regions within the given ranges. func (t *regionTree) RandomRegions(n int, ranges []KeyRange) []*RegionInfo { - if t.length() == 0 || n < 1 { + treeLen := t.length() + if treeLen == 0 || n < 1 { return nil } // Pre-allocate the variables to reduce the temporary memory allocations. var ( - startKey, endKey []byte - startIndex, endIndex, randIdx int - startItem *regionItem - pivotItem = ®ionItem{&RegionInfo{meta: &metapb.Region{}}} - region *RegionInfo - regions = make([]*RegionInfo, 0, n) - rangeNum, curLen = len(ranges), len(regions) + startKey, endKey []byte + startIndex, endIndex, randIndex int + startItem *regionItem + pivotItem = ®ionItem{&RegionInfo{meta: &metapb.Region{}}} + region *RegionInfo + regions = make([]*RegionInfo, 0, n) + rangeLen, curLen = len(ranges), len(regions) // setStartEndIndices is a helper function to set `startIndex` and `endIndex` // according to the `startKey` and `endKey`. // TODO: maybe we could cache the `startIndex` and `endIndex` for each range. @@ -360,37 +356,29 @@ func (t *regionTree) RandomRegions(n int, ranges []KeyRange) []*RegionInfo { pivotItem.meta.StartKey = endKey _, endIndex = t.tree.GetWithIndex(pivotItem) } else { - endIndex = t.tree.Len() + endIndex = treeLen } // Consider that the item in the tree may not be continuous, // we need to check if the previous item contains the key. - if startIndex != 0 && startItem == nil && t.tree.GetAt(startIndex-1).Contains(startKey) { - startIndex-- + if startIndex != 0 && startItem == nil { + region = t.tree.GetAt(startIndex - 1).RegionInfo + if region.Contains(startKey) { + startIndex-- + } } } ) // If no ranges specified, select randomly from the whole tree. // This is a fast path to reduce the unnecessary iterations. - if rangeNum == 0 { + if rangeLen == 0 { startKey, endKey = []byte(""), []byte("") setStartEndIndices() - if endIndex <= startIndex { - if len(endKey) > 0 && bytes.Compare(startKey, endKey) > 0 { - log.Error("wrong range keys", - logutil.ZapRedactString("start-key", string(HexRegionKey(startKey))), - logutil.ZapRedactString("end-key", string(HexRegionKey(endKey))), - errs.ZapError(errs.ErrWrongRangeKeys)) - } - return regions - } for curLen < n { - randIdx = rand.Intn(endIndex-startIndex) + startIndex - region = t.tree.GetAt(randIdx).RegionInfo - if curLen < n && region.isInvolved(startKey, endKey) { + randIndex = rand.Intn(endIndex-startIndex) + startIndex + region = t.tree.GetAt(randIndex).RegionInfo + if region.isInvolved(startKey, endKey) { regions = append(regions, region) curLen++ - } else if curLen == n { - break } // No region found, directly break to avoid infinite loop. if curLen == 0 { @@ -403,7 +391,7 @@ func (t *regionTree) RandomRegions(n int, ranges []KeyRange) []*RegionInfo { // keep retrying until we get enough regions. for curLen < n { // Shuffle the ranges to increase the randomness. - for _, i := range rand.Perm(rangeNum) { + for _, i := range rand.Perm(rangeLen) { startKey, endKey = ranges[i].StartKey, ranges[i].EndKey setStartEndIndices() if endIndex <= startIndex { @@ -416,13 +404,14 @@ func (t *regionTree) RandomRegions(n int, ranges []KeyRange) []*RegionInfo { continue } - randIdx = rand.Intn(endIndex-startIndex) + startIndex - region = t.tree.GetAt(randIdx).RegionInfo - if curLen < n && region.isInvolved(startKey, endKey) { + randIndex = rand.Intn(endIndex-startIndex) + startIndex + region = t.tree.GetAt(randIndex).RegionInfo + if region.isInvolved(startKey, endKey) { regions = append(regions, region) curLen++ - } else if curLen == n { - return regions + if curLen == n { + return regions + } } } // No region found, directly break to avoid infinite loop.