region cache: limit key ranges sent to PD in one `BatchScanRegions` gRPC (#1712)

close tikv/client-go#1704

Signed-off-by: you06 <you1474600@gmail.com>

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
 		}

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)
+}