region cache: batch find regions by key ranges from cache

internal/locate/region_cache.go

@@ -1240,24 +1240,50 @@ 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.
+		containsAll := false
+	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, defaultRegionsPerBatch)
+			if err != nil {
+				return nil, err
+			}
+			for _, r = range batchRegionInCache {
+				if !r.Contains(keyRange.StartKey) { // uncached hole, load the rest regions
+					break outer
+				}
+				cachedRegions = append(cachedRegions, r)
+				lastRegion = r
+				if r.ContainsByEnd(keyRange.EndKey) {
+					// the range is fully hit in the region cache.
+					containsAll = true
+					break outer
+				}
+				keyRange.StartKey = r.EndKey()
 			}
-			// 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.
+			if len(batchRegionInCache) < defaultRegionsPerBatch { // region cache miss, load the rest regions
 				break
 			}
-			keyRange.StartKey = r.EndKey()
+		}
+		if !containsAll {
+			uncachedRanges = append(uncachedRanges, pd.KeyRange{StartKey: keyRange.StartKey, EndKey: keyRange.EndKey})
 		}
 	}
 
@@ -2078,7 +2104,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.
+// It is the caller's responsibility to make sure that startKey is a node in the B-tree, otherwise, the startKey will not be included in the return regions.
 func (c *RegionCache) scanRegionsFromCache(bo *retry.Backoffer, startKey, endKey []byte, limit int) ([]*Region, error) {
 	if limit == 0 {
 		return nil, nil
@@ -2089,9 +2116,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
 }
 

internal/locate/region_cache_test.go

@@ -37,6 +37,7 @@ package locate
 import (
 	"bytes"
 	"context"
+	"encoding/binary"
 	"errors"
 	"fmt"
 	"math/rand"
@@ -2874,3 +2875,40 @@ func (s *testRegionCacheSuite) TestIssue1401() {
 	}, 3*time.Second, time.Second)
 	s.Require().True(isStoreContainLabel(newStore1.labels, "host", "0.0.0.0:20161"))
 }
+
+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()
+}

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,22 @@ func (s *SortedRegions) SearchByKey(key []byte, isEndKey bool) (r *Region) {
 }
 
 // AscendGreaterOrEqual returns all items that are greater than or equal to the key.
+// It is the caller's responsibility to make sure that startKey is a node in the B-tree, otherwise, the startKey will not be included in the return regions.
 func (s *SortedRegions) AscendGreaterOrEqual(startKey, endKey []byte, limit int) (regions []*Region) {
+	now := time.Now().Unix()
+	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(now) {
+			return false
+		}
+		if len(lastStartKey) > 0 && !region.Contains(lastStartKey) { // uncached hole
+			return false
+		}
+		lastStartKey = region.EndKey()
 		regions = append(regions, region)
 		return len(regions) < limit
 	})