core: batch get region size

pkg/core/region.go

@@ -41,7 +41,10 @@ import (
 	"go.uber.org/zap"
 )
 
-const randomRegionMaxRetry = 10
+const (
+	randomRegionMaxRetry = 10
+	scanRegionLimit      = 1000
+)
 
 // errRegionIsStale is error info for region is stale.
 func errRegionIsStale(region *metapb.Region, origin *metapb.Region) error {
@@ -1610,16 +1613,31 @@ func (r *RegionsInfo) ScanRegionWithIterator(startKey []byte, iterator func(regi
 
 // GetRegionSizeByRange scans regions intersecting [start key, end key), returns the total region size of this range.
 func (r *RegionsInfo) GetRegionSizeByRange(startKey, endKey []byte) int64 {
-	r.t.RLock()
-	defer r.t.RUnlock()
 	var size int64
-	r.tree.scanRange(startKey, func(region *RegionInfo) bool {
-		if len(endKey) > 0 && bytes.Compare(region.GetStartKey(), endKey) >= 0 {
-			return false
+	for {
+		r.t.RLock()
+		var cnt int
+		r.tree.scanRange(startKey, func(region *RegionInfo) bool {
+			if len(endKey) > 0 && bytes.Compare(region.GetStartKey(), endKey) >= 0 {
+				return false
+			}
+			if cnt >= scanRegionLimit {
+				return false
+			}
+			cnt++
+			startKey = region.GetEndKey()
+			size += region.GetApproximateSize()
+			return true
+		})
+		r.t.RUnlock()
+		if cnt == 0 {
+			break
 		}
-		size += region.GetApproximateSize()
-		return true
-	})
+		if len(startKey) == 0 || (len(endKey) > 0 && bytes.Compare(startKey, endKey) >= 0) {
+			break
+		}
+	}
+
 	return size
 }
 

pkg/core/region_test.go

@@ -18,8 +18,10 @@
 	"crypto/rand"
 	"fmt"
 	"math"
+	mrand "math/rand"
 	"strconv"
 	"testing"
+	"time"
 
 	"github.com/pingcap/failpoint"
 	"github.com/pingcap/kvproto/pkg/metapb"
@@ -658,6 +660,124 @@
 	}
 }
 
+func BenchmarkRandomSetRegion(b *testing.B) {
+	regions := NewRegionsInfo()
+	var items []*RegionInfo
+	for i := 0; i < 1000000; 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)
+		items = append(items, region)
+	}
+	b.ResetTimer()
+	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)
+	}
+}
+
+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
+	for i := 0; i < 1000000; 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, SetApproximateSize(10))
+		origin, overlaps, rangeChanged := regions.SetRegion(region)
+		regions.UpdateSubTree(region, origin, overlaps, rangeChanged)
+		items = append(items, region)
+	}
+	b.ResetTimer()
+	go func() {
+		for {
+			regions.GetRegionSizeByRange([]byte(""), []byte(""))
+			time.Sleep(time.Millisecond)
+		}
+	}()
+	for i := 0; i < b.N; i++ {
+		item := items[i%len(items)]
+		item.approximateKeys = int64(200000)
+		origin, overlaps, rangeChanged := regions.SetRegion(item)
+		regions.UpdateSubTree(item, origin, overlaps, rangeChanged)
+	}
+}
+
+func BenchmarkRandomSetRegionWithGetRegionSizeByRangeParallel(b *testing.B) {
+	regions := NewRegionsInfo()
+	var items []*RegionInfo
+	for i := 0; i < 1000000; 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)
+		items = append(items, region)
+	}
+	b.ResetTimer()
+	go func() {
+		for {
+			regions.GetRegionSizeByRange([]byte(""), []byte(""))
+			time.Sleep(time.Millisecond)
+		}
+	}()
+
+	b.RunParallel(
+		func(pb *testing.PB) {
+			for pb.Next() {
+				item := items[mrand.Intn(len(items))]
+				n := item.Clone(SetApproximateSize(20))
+				origin, overlaps, rangeChanged := regions.SetRegion(n)
+				regions.UpdateSubTree(item, origin, overlaps, rangeChanged)
+			}
+		},
+	)
+}
+
 const keyLength = 100
 
 func randomBytes(n int) []byte {