Repository: tikv/client-go (go)

Issue:
The size of BatchScanRegionsRequest might exceed the default MaxRequestBytes
The size of BatchScanRegionsRequest might exceed the default MaxRequestBytes (150MiB) when there are a lot of key ranges waiting for scan. The related messages are shown as the following.
```
grpc:received message larger than max(297800997 vs 157810688)
pdrpc backoffer.maxsleep 20000ms is exceeded,errors:batchScanRegion from PD failed,range num:1059860,limit:128
```

Pull request description:
region cache: limit key ranges sent to PD in one `BatchScanRegions` gRPC
close #1704

In some user scenario, millions of key ranges can be passed in one `BatchLocateKeyRanges` call, if we sent such number of key ranges to PD in one RPC, it can either exceeds the gRPC size limit or cause performance issues. Because we limit the number of response regions in one gRPC call, it's meaningless sending too many ranges.

This PR limit the number of key ranges to 16 times of region batch size to avoid unexpected error or network traffic/CPU waste.

Patch:
diff --git a/internal/locate/region_cache.go b/internal/locate/region_cache.go
index 93c2a09..1622980 100644
--- a/internal/locate/region_cache.go
+++ b/internal/locate/region_cache.go
@@ -1350,7 +1350,16 @@ func (c *RegionCache) BatchLocateKeyRanges(bo *retry.Backoffer, keyRanges []kv.K
 
 	// 2. load remaining regions from pd client
 	for len(uncachedRanges) > 0 {
-		regions, err := c.BatchLoadRegionsWithKeyRanges(bo, uncachedRanges, defaultRegionsPerBatch, opts...)
+		// If we send too many ranges to PD, it may exceed the size of grpc limitation or cause a timeout error.
+		// So we limit the number of ranges per batch to avoid this issue.
+		maxRangesPerBatch := 16 * defaultRegionsPerBatch
+		var toBeSentRanges []router.KeyRange
+		if len(uncachedRanges) > maxRangesPerBatch {
+			toBeSentRanges = uncachedRanges[:maxRangesPerBatch]
+		} else {
+			toBeSentRanges = uncachedRanges
+		}
+		regions, err := c.BatchLoadRegionsWithKeyRanges(bo, toBeSentRanges, defaultRegionsPerBatch, opts...)
 		if err != nil {
 			return nil, err
 		}


Test patch:
diff --git a/internal/locate/region_cache_test.go b/internal/locate/region_cache_test.go
index 10dbac9..0cf1bb5 100644
--- a/internal/locate/region_cache_test.go
+++ b/internal/locate/region_cache_test.go
@@ -3086,3 +3086,25 @@ func (s *testRegionCacheSuite) TestRegionCacheValidAfterLoading() {
 		s.True(region.isValid())
 	}
 }
+
+func (s *testRegionCacheSuite) TestBatchLoadLimitRanges() {
+	ranges := make([]kv.KeyRange, 0, 100000)
+	for i := 0; i < 100000; i++ {
+		startKey := make([]byte, 8)
+		endKey := make([]byte, 8)
+		binary.BigEndian.PutUint64(startKey, uint64(i*2))
+		binary.BigEndian.PutUint64(endKey, uint64(i*2+1))
+		ranges = append(ranges, kv.KeyRange{StartKey: startKey, EndKey: endKey})
+	}
+
+	originalBatchScanRegions := s.cache.pdClient.BatchScanRegions
+	s.cache.pdClient = &inspectedPDClient{
+		Client: s.cache.pdClient,
+		batchScanRegions: func(ctx context.Context, keyRanges []router.KeyRange, limit int, opts ...opt.GetRegionOption) ([]*router.Region, error) {
+			s.LessOrEqual(len(keyRanges), 16*defaultRegionsPerBatch)
+			return originalBatchScanRegions(ctx, keyRanges, limit, opts...)
+		},
+	}
+	_, err := s.cache.BatchLocateKeyRanges(s.bo, ranges)
+	s.Nil(err)
+}