From 8eb74226625312216252a8a75bf5ef446f9e62d6 Mon Sep 17 00:00:00 2001 From: wshwsh12 <793703860@qq.com> Date: Wed, 31 Jul 2024 14:43:03 +0800 Subject: [PATCH] try add benchmark Signed-off-by: wshwsh12 <793703860@qq.com> --- internal/locate/region_cache.go | 60 +++++++++++++++++++--------- internal/locate/region_cache_test.go | 38 ++++++++++++++++++ internal/locate/sorted_btree.go | 10 +++++ 3 files changed, 89 insertions(+), 19 deletions(-) diff --git a/internal/locate/region_cache.go b/internal/locate/region_cache.go index 151de21cf..92a705cb4 100644 --- a/internal/locate/region_cache.go +++ b/internal/locate/region_cache.go @@ -1240,24 +1240,48 @@ func (c *RegionCache) BatchLocateKeyRanges(bo *retry.Backoffer, keyRanges []kv.K keyRange.StartKey = lastRegion.EndKey() } } + // TODO: find all the cached regions in the range. + // now we only check if the region is cached from the lower bound, if there is a uncached hole in the middle, + // we will load the rest regions even they are cached. + r := c.tryFindRegionByKey(keyRange.StartKey, false) + lastRegion = r + if r == nil { + // region cache miss, add the cut range to uncachedRanges, load from PD later. + uncachedRanges = append(uncachedRanges, pd.KeyRange{StartKey: keyRange.StartKey, EndKey: keyRange.EndKey}) + continue + } + // region cache hit, add the region to cachedRegions. + cachedRegions = append(cachedRegions, r) + if r.ContainsByEnd(keyRange.EndKey) { + // the range is fully hit in the region cache. + continue + } + keyRange.StartKey = r.EndKey() + // Batch load rest regions from Cache. + batchSize := 100 + outer: for { - // TODO: find all the cached regions in the range. - // now we only check if the region is cached from the lower bound, if there is a uncached hole in the middle, - // we will load the rest regions even they are cached. - r := c.tryFindRegionByKey(keyRange.StartKey, false) - lastRegion = r - if r == nil { - // region cache miss, add the cut range to uncachedRanges, load from PD later. - uncachedRanges = append(uncachedRanges, pd.KeyRange{StartKey: keyRange.StartKey, EndKey: keyRange.EndKey}) - break + batchRegionInCache, err := c.scanRegionsFromCache(bo, keyRange.StartKey, keyRange.EndKey, batchSize) + if err != nil { + return nil, err } - // region cache hit, add the region to cachedRegions. - cachedRegions = append(cachedRegions, r) - if r.ContainsByEnd(keyRange.EndKey) { - // the range is fully hit in the region cache. - break + for _, r = range batchRegionInCache { + if !r.Contains(keyRange.StartKey) { // uncached hole, load the rest regions + uncachedRanges = append(uncachedRanges, pd.KeyRange{StartKey: keyRange.StartKey, EndKey: keyRange.EndKey}) + break outer + } + cachedRegions = append(cachedRegions, r) + lastRegion = r + if r.ContainsByEnd(keyRange.EndKey) { + // the range is fully hit in the region cache. + break outer + } + keyRange.StartKey = r.EndKey() + } + if len(batchRegionInCache) < batchSize { // region cache miss, load the rest regions + uncachedRanges = append(uncachedRanges, pd.KeyRange{StartKey: keyRange.StartKey, EndKey: keyRange.EndKey}) + break outer } - keyRange.StartKey = r.EndKey() } } @@ -2078,7 +2102,8 @@ func (c *RegionCache) loadRegionByID(bo *retry.Backoffer, regionID uint64) (*Reg } } -// TODO(youjiali1995): for optimizing BatchLoadRegionsWithKeyRange, not used now. +// For optimizing BatchLocateKeyRanges, scanRegionsFromCache scans at most `limit` regions from cache. +// The first region's startKey is `startKey`, all regions returned are continuous and have no hole. func (c *RegionCache) scanRegionsFromCache(bo *retry.Backoffer, startKey, endKey []byte, limit int) ([]*Region, error) { if limit == 0 { return nil, nil @@ -2089,9 +2114,6 @@ func (c *RegionCache) scanRegionsFromCache(bo *retry.Backoffer, startKey, endKey defer c.mu.RUnlock() regions = c.mu.sorted.AscendGreaterOrEqual(startKey, endKey, limit) - if len(regions) == 0 { - return nil, errors.New("no regions in the cache") - } return regions, nil } diff --git a/internal/locate/region_cache_test.go b/internal/locate/region_cache_test.go index 888b96d15..14c5fd93b 100644 --- a/internal/locate/region_cache_test.go +++ b/internal/locate/region_cache_test.go @@ -37,6 +37,7 @@ package locate import ( "bytes" "context" + "encoding/binary" "errors" "fmt" "math/rand" @@ -2833,3 +2834,40 @@ func (s *testRegionCacheSuite) TestScanRegionsWithGaps() { }) s.Equal(batchScanRegionRes, regions) } + +func BenchmarkBatchLocateKeyRangesFromCache(t *testing.B) { + t.StopTimer() + s := new(testRegionCacheSuite) + s.SetT(&testing.T{}) + s.SetupTest() + + regionNum := 10000 + regions := s.cluster.AllocIDs(regionNum) + regions = append([]uint64{s.region1}, regions...) + + peers := [][]uint64{{s.peer1, s.peer2}} + for i := 0; i < regionNum-1; i++ { + peers = append(peers, s.cluster.AllocIDs(2)) + } + + for i := 0; i < regionNum-1; i++ { + b := make([]byte, 8) + binary.BigEndian.PutUint64(b, uint64(i*2)) + s.cluster.Split(regions[i], regions[i+1], b, peers[i+1], peers[i+1][0]) + } + + // cache all regions + keyLocation, err := s.cache.BatchLocateKeyRanges(s.bo, []kv.KeyRange{{StartKey: []byte(""), EndKey: []byte("")}}) + if err != nil || len(keyLocation) != regionNum { + t.FailNow() + } + + t.StartTimer() + for i := 0; i < t.N; i++ { + keyLocation, err := s.cache.BatchLocateKeyRanges(s.bo, []kv.KeyRange{{StartKey: []byte(""), EndKey: []byte("")}}) + if err != nil || len(keyLocation) != regionNum { + t.FailNow() + } + } + s.TearDownTest() +} diff --git a/internal/locate/sorted_btree.go b/internal/locate/sorted_btree.go index 5e2b88536..57bdead8e 100644 --- a/internal/locate/sorted_btree.go +++ b/internal/locate/sorted_btree.go @@ -36,6 +36,7 @@ package locate import ( "bytes" + "time" "github.com/google/btree" "github.com/tikv/client-go/v2/internal/logutil" @@ -79,12 +80,21 @@ func (s *SortedRegions) SearchByKey(key []byte, isEndKey bool) (r *Region) { } // AscendGreaterOrEqual returns all items that are greater than or equal to the key. +// The first region's startKey is `startKey`, all regions returned are continuous and have no hole. func (s *SortedRegions) AscendGreaterOrEqual(startKey, endKey []byte, limit int) (regions []*Region) { + lastStartKey := startKey s.b.AscendGreaterOrEqual(newBtreeSearchItem(startKey), func(item *btreeItem) bool { region := item.cachedRegion if len(endKey) > 0 && bytes.Compare(region.StartKey(), endKey) >= 0 { return false } + if !region.checkRegionCacheTTL(time.Now().Unix()) { + return false + } + if !region.Contains(lastStartKey) { // uncached hole + return false + } + lastStartKey = region.EndKey() regions = append(regions, region) return len(regions) < limit })