diff --git a/pkg/core/region.go b/pkg/core/region.go index 2f2d2ad48c3..1a0a74abc0d 100644 --- a/pkg/core/region.go +++ b/pkg/core/region.go @@ -1619,14 +1619,13 @@ func (r *RegionsInfo) GetRegionSizeByRange(startKey, endKey []byte) int64 { var cnt int r.tree.scanRange(startKey, func(region *RegionInfo) bool { if len(endKey) > 0 && bytes.Compare(region.GetStartKey(), endKey) >= 0 { - startKey = region.GetEndKey() return false } if cnt >= scanRegionLimit { - startKey = region.GetEndKey() return false } cnt++ + startKey = region.GetEndKey() size += region.GetApproximateSize() return true }) @@ -1634,8 +1633,7 @@ func (r *RegionsInfo) GetRegionSizeByRange(startKey, endKey []byte) int64 { if cnt == 0 { break } - - if len(startKey) == 0 { + if len(startKey) == 0 || (len(endKey) > 0 && bytes.Compare(startKey, endKey) >= 0) { break } } diff --git a/pkg/core/region_test.go b/pkg/core/region_test.go index 3928c774fbd..0663c3144ca 100644 --- a/pkg/core/region_test.go +++ b/pkg/core/region_test.go @@ -685,6 +685,34 @@ func BenchmarkRandomSetRegion(b *testing.B) { } } +func TestGetRegionSizeByRange(t *testing.T) { + regions := NewRegionsInfo() + nums := 1000010 + for i := 0; i < nums; i++ { + peer := &metapb.Peer{StoreId: 1, Id: uint64(i + 1)} + endKey := []byte(fmt.Sprintf("%20d", i+1)) + if i == nums-1 { + endKey = []byte("") + } + region := NewRegionInfo(&metapb.Region{ + Id: uint64(i + 1), + Peers: []*metapb.Peer{peer}, + StartKey: []byte(fmt.Sprintf("%20d", i)), + EndKey: endKey, + }, peer, SetApproximateSize(10)) + origin, overlaps, rangeChanged := regions.SetRegion(region) + regions.UpdateSubTree(region, origin, overlaps, rangeChanged) + } + totalSize := regions.GetRegionSizeByRange([]byte(""), []byte("")) + require.Equal(t, int64(nums*10), totalSize) + for i := 1; i < 10; i++ { + verifyNum := int(nums / i) + endKey := fmt.Sprintf("%20d", verifyNum) + totalSize := regions.GetRegionSizeByRange([]byte(""), []byte(endKey)) + require.Equal(t, int64(verifyNum*10), totalSize) + } +} + func BenchmarkRandomSetRegionWithGetRegionSizeByRange(b *testing.B) { regions := NewRegionsInfo() var items []*RegionInfo @@ -695,7 +723,7 @@ func BenchmarkRandomSetRegionWithGetRegionSizeByRange(b *testing.B) { Peers: []*metapb.Peer{peer}, StartKey: []byte(fmt.Sprintf("%20d", i)), EndKey: []byte(fmt.Sprintf("%20d", i+1)), - }, peer) + }, peer, SetApproximateSize(10)) origin, overlaps, rangeChanged := regions.SetRegion(region) regions.UpdateSubTree(region, origin, overlaps, rangeChanged) items = append(items, region) @@ -710,7 +738,6 @@ func BenchmarkRandomSetRegionWithGetRegionSizeByRange(b *testing.B) { for i := 0; i < b.N; i++ { item := items[i%len(items)] item.approximateKeys = int64(200000) - item.approximateSize = int64(20) origin, overlaps, rangeChanged := regions.SetRegion(item) regions.UpdateSubTree(item, origin, overlaps, rangeChanged) }