fix

Signed-off-by: nolouch <[email protected]>
tikv · Nov 7, 2023 · 2a8ae9c · 2a8ae9c
1 parent 8073cb9
commit 2a8ae9c
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 6 deletions.
diff --git a/pkg/core/region.go b/pkg/core/region.go
@@ -1619,23 +1619,21 @@ 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
 		})
 		r.t.RUnlock()
 		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
@@ -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)
 	}