region_cache: support fall back to leader handle when using ScanRegions/BatchScanRegions (#1794)

close #1793

Signed-off-by: okjiang <819421878@qq.com>

Author: okJiang

internal/locate/region_cache.go

@@ -2370,6 +2370,7 @@ func (c *RegionCache) scanRegions(bo *retry.Backoffer, startKey, endKey []byte,
 		ctx = opentracing.ContextWithSpan(ctx, span1)
 	}
 
+	pdOpts := []opt.GetRegionOption{opt.WithAllowFollowerHandle()}
 	var backoffErr error
 	for {
 		if backoffErr != nil {
@@ -2380,7 +2381,7 @@ func (c *RegionCache) scanRegions(bo *retry.Backoffer, startKey, endKey []byte,
 		}
 		start := time.Now()
 		// TODO: ScanRegions has been deprecated in favor of BatchScanRegions.
-		regionsInfo, err := c.pdClient.ScanRegions(withPDCircuitBreaker(ctx), startKey, endKey, limit, opt.WithAllowFollowerHandle())
+		regionsInfo, err := c.pdClient.ScanRegions(withPDCircuitBreaker(ctx), startKey, endKey, limit, pdOpts...)
 		metrics.LoadRegionCacheHistogramWithRegions.Observe(time.Since(start).Seconds())
 		if err != nil {
 			if apicodec.IsDecodeError(err) {
@@ -2399,10 +2400,18 @@ func (c *RegionCache) scanRegions(bo *retry.Backoffer, startKey, endKey []byte,
 
 		if len(regionsInfo) == 0 {
 			backoffErr = errors.Errorf("PD returned no region, limit: %d", limit)
+			metrics.TiKVStaleRegionFromPDCounter.Inc()
+			if len(pdOpts) > 0 {
+				pdOpts = nil
+			}
 			continue
 		}
 		if regionsHaveGapInRanges([]router.KeyRange{{StartKey: startKey, EndKey: endKey}}, regionsInfo, limit) {
 			backoffErr = errors.Errorf("PD returned regions have gaps, limit: %d", limit)
+			metrics.TiKVStaleRegionFromPDCounter.Inc()
+			if len(pdOpts) > 0 {
+				pdOpts = nil
+			}
 			continue
 		}
 		validRegions, err := c.handleRegionInfos(bo, regionsInfo, true)
@@ -2414,6 +2423,10 @@ func (c *RegionCache) scanRegions(bo *retry.Backoffer, startKey, endKey []byte,
 		// Retry if there is no valid regions with leaders.
 		if len(validRegions) == 0 {
 			backoffErr = errors.Errorf("All returned regions have no leaders, limit: %d", limit)
+			metrics.TiKVStaleRegionFromPDCounter.Inc()
+			if len(pdOpts) > 0 {
+				pdOpts = nil
+			}
 			continue
 		}
 		return validRegions, nil
@@ -2435,23 +2448,31 @@ func (c *RegionCache) batchScanRegions(bo *retry.Backoffer, keyRanges []router.K
 	for _, op := range opts {
 		op(&batchOpt)
 	}
+	needFollowerHandle := true
+	initPdOpts := func() []opt.GetRegionOption {
+		pdOpts := []opt.GetRegionOption{
+			opt.WithOutputMustContainAllKeyRange(),
+		}
+		if needFollowerHandle {
+			pdOpts = append(pdOpts, opt.WithAllowFollowerHandle())
+		}
+		if batchOpt.needBuckets {
+			pdOpts = append(pdOpts, opt.WithBuckets())
+		}
+		return pdOpts
+	}
 	// TODO: return start key and end key after redact is introduced.
 	var backoffErr error
 	for {
+		pdOpts := initPdOpts()
+
 		if backoffErr != nil {
 			err := bo.Backoff(retry.BoPDRPC, backoffErr)
 			if err != nil {
 				return nil, errors.WithStack(err)
 			}
 		}
 		start := time.Now()
-		pdOpts := []opt.GetRegionOption{
-			opt.WithAllowFollowerHandle(),
-			opt.WithOutputMustContainAllKeyRange(),
-		}
-		if batchOpt.needBuckets {
-			pdOpts = append(pdOpts, opt.WithBuckets())
-		}
 		regionsInfo, err := c.pdClient.BatchScanRegions(withPDCircuitBreaker(ctx), keyRanges, limit, pdOpts...)
 		metrics.LoadRegionCacheHistogramWithBatchScanRegions.Observe(time.Since(start).Seconds())
 		if err != nil {
@@ -2477,13 +2498,21 @@ func (c *RegionCache) batchScanRegions(bo *retry.Backoffer, keyRanges []router.K
 				"PD returned no region, range num: %d, limit: %d",
 				len(keyRanges), limit,
 			)
+			metrics.TiKVStaleRegionFromPDCounter.Inc()
+			if needFollowerHandle {
+				needFollowerHandle = false
+			}
 			continue
 		}
 		if regionsHaveGapInRanges(keyRanges, regionsInfo, limit) {
 			backoffErr = errors.Errorf(
 				"PD returned regions have gaps, range num: %d, limit: %d",
 				len(keyRanges), limit,
 			)
+			metrics.TiKVStaleRegionFromPDCounter.Inc()
+			if needFollowerHandle {
+				needFollowerHandle = false
+			}
 			continue
 		}
 		validRegions, err := c.handleRegionInfos(bo, regionsInfo, batchOpt.needRegionHasLeaderPeer)
@@ -2495,6 +2524,10 @@ func (c *RegionCache) batchScanRegions(bo *retry.Backoffer, keyRanges []router.K
 		// Retry if there is no valid regions with leaders.
 		if len(validRegions) == 0 {
 			backoffErr = errors.Errorf("All returned regions have no leaders, limit: %d", limit)
+			metrics.TiKVStaleRegionFromPDCounter.Inc()
+			if needFollowerHandle {
+				needFollowerHandle = false
+			}
 			continue
 		}
 		return validRegions, nil