Skip to content

Commit da1e92d

Browse files
ti-chi-botnolouchti-chi-bot[bot]
authored
core: batch get region size (#7252) (#7332)
close #7248 Signed-off-by: nolouch <nolouch@gmail.com> Co-authored-by: nolouch <nolouch@gmail.com> Co-authored-by: ti-chi-bot[bot] <108142056+ti-chi-bot[bot]@users.noreply.github.com>
1 parent 710ffcd commit da1e92d

File tree

4 files changed

+152
-13
lines changed

4 files changed

+152
-13
lines changed

‎pkg/core/region.go‎

Lines changed: 27 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,10 @@ import (
4141
"go.uber.org/zap"
4242
)
4343

44-
const randomRegionMaxRetry = 10
44+
const (
45+
randomRegionMaxRetry = 10
46+
scanRegionLimit = 1000
47+
)
4548

4649
// errRegionIsStale is error info for region is stale.
4750
func errRegionIsStale(region *metapb.Region, origin *metapb.Region) error {
@@ -1610,16 +1613,31 @@ func (r *RegionsInfo) ScanRegionWithIterator(startKey []byte, iterator func(regi
16101613

16111614
// GetRegionSizeByRange scans regions intersecting [start key, end key), returns the total region size of this range.
16121615
func (r *RegionsInfo) GetRegionSizeByRange(startKey, endKey []byte) int64 {
1613-
r.t.RLock()
1614-
defer r.t.RUnlock()
16151616
var size int64
1616-
r.tree.scanRange(startKey, func(region *RegionInfo) bool {
1617-
if len(endKey) > 0 && bytes.Compare(region.GetStartKey(), endKey) >= 0 {
1618-
return false
1617+
for {
1618+
r.t.RLock()
1619+
var cnt int
1620+
r.tree.scanRange(startKey, func(region *RegionInfo) bool {
1621+
if len(endKey) > 0 && bytes.Compare(region.GetStartKey(), endKey) >= 0 {
1622+
return false
1623+
}
1624+
if cnt >= scanRegionLimit {
1625+
return false
1626+
}
1627+
cnt++
1628+
startKey = region.GetEndKey()
1629+
size += region.GetApproximateSize()
1630+
return true
1631+
})
1632+
r.t.RUnlock()
1633+
if cnt == 0 {
1634+
break
16191635
}
1620-
size += region.GetApproximateSize()
1621-
return true
1622-
})
1636+
if len(startKey) == 0 {
1637+
break
1638+
}
1639+
}
1640+
16231641
return size
16241642
}
16251643

‎pkg/core/region_test.go‎

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,10 @@ import (
1818
"crypto/rand"
1919
"fmt"
2020
"math"
21+
mrand "math/rand"
2122
"strconv"
2223
"testing"
24+
"time"
2325

2426
"github.com/pingcap/failpoint"
2527
"github.com/pingcap/kvproto/pkg/metapb"
@@ -658,6 +660,124 @@ func BenchmarkRandomRegion(b *testing.B) {
658660
}
659661
}
660662

663+
func BenchmarkRandomSetRegion(b *testing.B) {
664+
regions := NewRegionsInfo()
665+
var items []*RegionInfo
666+
for i := 0; i < 1000000; i++ {
667+
peer := &metapb.Peer{StoreId: 1, Id: uint64(i + 1)}
668+
region := NewRegionInfo(&metapb.Region{
669+
Id: uint64(i + 1),
670+
Peers: []*metapb.Peer{peer},
671+
StartKey: []byte(fmt.Sprintf("%20d", i)),
672+
EndKey: []byte(fmt.Sprintf("%20d", i+1)),
673+
}, peer)
674+
origin, overlaps, rangeChanged := regions.SetRegion(region)
675+
regions.UpdateSubTree(region, origin, overlaps, rangeChanged)
676+
items = append(items, region)
677+
}
678+
b.ResetTimer()
679+
for i := 0; i < b.N; i++ {
680+
item := items[i%len(items)]
681+
item.approximateKeys = int64(200000)
682+
item.approximateSize = int64(20)
683+
origin, overlaps, rangeChanged := regions.SetRegion(item)
684+
regions.UpdateSubTree(item, origin, overlaps, rangeChanged)
685+
}
686+
}
687+
688+
func TestGetRegionSizeByRange(t *testing.T) {
689+
regions := NewRegionsInfo()
690+
nums := 1000010
691+
for i := 0; i < nums; i++ {
692+
peer := &metapb.Peer{StoreId: 1, Id: uint64(i + 1)}
693+
endKey := []byte(fmt.Sprintf("%20d", i+1))
694+
if i == nums-1 {
695+
endKey = []byte("")
696+
}
697+
region := NewRegionInfo(&metapb.Region{
698+
Id: uint64(i + 1),
699+
Peers: []*metapb.Peer{peer},
700+
StartKey: []byte(fmt.Sprintf("%20d", i)),
701+
EndKey: endKey,
702+
}, peer, SetApproximateSize(10))
703+
origin, overlaps, rangeChanged := regions.SetRegion(region)
704+
regions.UpdateSubTree(region, origin, overlaps, rangeChanged)
705+
}
706+
totalSize := regions.GetRegionSizeByRange([]byte(""), []byte(""))
707+
require.Equal(t, int64(nums*10), totalSize)
708+
for i := 1; i < 10; i++ {
709+
verifyNum := nums / i
710+
endKey := fmt.Sprintf("%20d", verifyNum)
711+
totalSize := regions.GetRegionSizeByRange([]byte(""), []byte(endKey))
712+
require.Equal(t, int64(verifyNum*10), totalSize)
713+
}
714+
}
715+
716+
func BenchmarkRandomSetRegionWithGetRegionSizeByRange(b *testing.B) {
717+
regions := NewRegionsInfo()
718+
var items []*RegionInfo
719+
for i := 0; i < 1000000; i++ {
720+
peer := &metapb.Peer{StoreId: 1, Id: uint64(i + 1)}
721+
region := NewRegionInfo(&metapb.Region{
722+
Id: uint64(i + 1),
723+
Peers: []*metapb.Peer{peer},
724+
StartKey: []byte(fmt.Sprintf("%20d", i)),
725+
EndKey: []byte(fmt.Sprintf("%20d", i+1)),
726+
}, peer, SetApproximateSize(10))
727+
origin, overlaps, rangeChanged := regions.SetRegion(region)
728+
regions.UpdateSubTree(region, origin, overlaps, rangeChanged)
729+
items = append(items, region)
730+
}
731+
b.ResetTimer()
732+
go func() {
733+
for {
734+
regions.GetRegionSizeByRange([]byte(""), []byte(""))
735+
time.Sleep(time.Millisecond)
736+
}
737+
}()
738+
for i := 0; i < b.N; i++ {
739+
item := items[i%len(items)]
740+
item.approximateKeys = int64(200000)
741+
origin, overlaps, rangeChanged := regions.SetRegion(item)
742+
regions.UpdateSubTree(item, origin, overlaps, rangeChanged)
743+
}
744+
}
745+
746+
func BenchmarkRandomSetRegionWithGetRegionSizeByRangeParallel(b *testing.B) {
747+
regions := NewRegionsInfo()
748+
var items []*RegionInfo
749+
for i := 0; i < 1000000; i++ {
750+
peer := &metapb.Peer{StoreId: 1, Id: uint64(i + 1)}
751+
region := NewRegionInfo(&metapb.Region{
752+
Id: uint64(i + 1),
753+
Peers: []*metapb.Peer{peer},
754+
StartKey: []byte(fmt.Sprintf("%20d", i)),
755+
EndKey: []byte(fmt.Sprintf("%20d", i+1)),
756+
}, peer)
757+
origin, overlaps, rangeChanged := regions.SetRegion(region)
758+
regions.UpdateSubTree(region, origin, overlaps, rangeChanged)
759+
items = append(items, region)
760+
}
761+
b.ResetTimer()
762+
go func() {
763+
for {
764+
regions.GetRegionSizeByRange([]byte(""), []byte(""))
765+
time.Sleep(time.Millisecond)
766+
}
767+
}()
768+
769+
b.RunParallel(
770+
func(pb *testing.PB) {
771+
for pb.Next() {
772+
item := items[mrand.Intn(len(items))]
773+
n := item.Clone(SetApproximateSize(20))
774+
origin, overlaps, rangeChanged := regions.SetRegion(n)
775+
regions.UpdateSubTree(item, origin, overlaps, rangeChanged)
776+
}
777+
},
778+
)
779+
}
780+
661781
const keyLength = 100
662782

663783
func randomBytes(n int) []byte {

‎server/cluster/cluster.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1841,12 +1841,13 @@ func (c *RaftCluster) checkStores() {
18411841
if err := c.ReadyToServe(storeID); err != nil {
18421842
log.Error("change store to serving failed",
18431843
zap.Stringer("store", store.GetMeta()),
1844+
zap.Int("region-count", c.GetTotalRegionCount()),
18441845
errs.ZapError(err))
18451846
}
18461847
} else if c.IsPrepared() {
18471848
threshold := c.getThreshold(stores, store)
1848-
log.Debug("store serving threshold", zap.Uint64("store-id", storeID), zap.Float64("threshold", threshold))
18491849
regionSize := float64(store.GetRegionSize())
1850+
log.Debug("store serving threshold", zap.Uint64("store-id", storeID), zap.Float64("threshold", threshold), zap.Float64("region-size", regionSize))
18501851
if regionSize >= threshold {
18511852
if err := c.ReadyToServe(storeID); err != nil {
18521853
log.Error("change store to serving failed",

‎tests/server/api/api_test.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -914,7 +914,7 @@ func TestPreparingProgress(t *testing.T) {
914914
tests.MustPutStore(re, cluster, store)
915915
}
916916
for i := 0; i < 100; i++ {
917-
tests.MustPutRegion(re, cluster, uint64(i+1), uint64(i)%3+1, []byte(fmt.Sprintf("p%d", i)), []byte(fmt.Sprintf("%d", i+1)), core.SetApproximateSize(10))
917+
tests.MustPutRegion(re, cluster, uint64(i+1), uint64(i)%3+1, []byte(fmt.Sprintf("%20d", i)), []byte(fmt.Sprintf("%20d", i+1)), core.SetApproximateSize(10))
918918
}
919919
// no store preparing
920920
output := sendRequest(re, leader.GetAddr()+"/pd/api/v1/stores/progress?action=preparing", http.MethodGet, http.StatusNotFound)
@@ -941,8 +941,8 @@ func TestPreparingProgress(t *testing.T) {
941941
re.Equal(math.MaxFloat64, p.LeftSeconds)
942942

943943
// update size
944-
tests.MustPutRegion(re, cluster, 1000, 4, []byte(fmt.Sprintf("%d", 1000)), []byte(fmt.Sprintf("%d", 1001)), core.SetApproximateSize(10))
945-
tests.MustPutRegion(re, cluster, 1001, 5, []byte(fmt.Sprintf("%d", 1001)), []byte(fmt.Sprintf("%d", 1002)), core.SetApproximateSize(40))
944+
tests.MustPutRegion(re, cluster, 1000, 4, []byte(fmt.Sprintf("%20d", 1000)), []byte(fmt.Sprintf("%20d", 1001)), core.SetApproximateSize(10))
945+
tests.MustPutRegion(re, cluster, 1001, 5, []byte(fmt.Sprintf("%20d", 1001)), []byte(fmt.Sprintf("%20d", 1002)), core.SetApproximateSize(40))
946946
time.Sleep(2 * time.Second)
947947
output = sendRequest(re, leader.GetAddr()+"/pd/api/v1/stores/progress?action=preparing", http.MethodGet, http.StatusOK)
948948
re.NoError(json.Unmarshal(output, &p))

0 commit comments

Comments
 (0)
⚡