From 0206a3c142bd094da05f8bf6381a8963a5e78a5f Mon Sep 17 00:00:00 2001 From: you06 Date: Wed, 3 Jul 2024 16:38:53 +0900 Subject: [PATCH] region cache: fallback to ScanRegions when BatchScanRegions receive unimplement error (#1378) * add fallback Signed-off-by: you06 * fallback to ScanRegions when BatchScanRegions receive unimplement error Signed-off-by: you06 * add comment & test Signed-off-by: you06 --------- Signed-off-by: you06 --- internal/locate/region_cache.go | 37 ++++++++++++++++++++++++++++ internal/locate/region_cache_test.go | 10 ++++++++ internal/mockstore/mocktikv/pd.go | 6 +++++ 3 files changed, 53 insertions(+) diff --git a/internal/locate/region_cache.go b/internal/locate/region_cache.go index 80565b0d..b15d502e 100644 --- a/internal/locate/region_cache.go +++ b/internal/locate/region_cache.go @@ -2192,6 +2192,9 @@ func (c *RegionCache) batchScanRegions(bo *retry.Backoffer, keyRanges []pd.KeyRa regionsInfo, err := c.pdClient.BatchScanRegions(ctx, keyRanges, limit, pdOpts...) metrics.LoadRegionCacheHistogramWithBatchScanRegions.Observe(time.Since(start).Seconds()) if err != nil { + if st, ok := status.FromError(err); ok && st.Code() == codes.Unimplemented { + return c.batchScanRegionsFallback(bo, keyRanges, limit, opts...) + } if apicodec.IsDecodeError(err) { return nil, errors.Errorf("failed to decode region range key, range num: %d, limit: %d, err: %v", len(keyRanges), limit, err) @@ -2216,6 +2219,40 @@ func (c *RegionCache) batchScanRegions(bo *retry.Backoffer, keyRanges []pd.KeyRa } } +func (c *RegionCache) batchScanRegionsFallback(bo *retry.Backoffer, keyRanges []pd.KeyRange, limit int, opts ...BatchLocateKeyRangesOpt) ([]*Region, error) { + logutil.BgLogger().Warn("batch scan regions fallback to scan regions", zap.Int("range-num", len(keyRanges))) + res := make([]*Region, 0, len(keyRanges)) + var lastRegion *Region + for _, keyRange := range keyRanges { + if lastRegion != nil { + endKey := lastRegion.EndKey() + if len(endKey) == 0 { + // end_key is empty means the last region is the last region of the store, which certainly contains all the rest ranges. + break + } + if bytes.Compare(endKey, keyRange.EndKey) >= 0 { + continue + } + if bytes.Compare(endKey, keyRange.StartKey) > 0 { + keyRange.StartKey = endKey + } + } + regions, err := c.scanRegions(bo, keyRange.StartKey, keyRange.EndKey, limit) + if err != nil { + return nil, err + } + if len(regions) > 0 { + lastRegion = regions[len(regions)-1] + } + res = append(res, regions...) + if len(regions) >= limit { + return res, nil + } + limit -= len(regions) + } + return res, nil +} + func (c *RegionCache) handleRegionInfos(bo *retry.Backoffer, regionsInfo []*pd.Region, needLeader bool) ([]*Region, error) { regions := make([]*Region, 0, len(regionsInfo)) for _, r := range regionsInfo { diff --git a/internal/locate/region_cache_test.go b/internal/locate/region_cache_test.go index 68965405..e53e587e 100644 --- a/internal/locate/region_cache_test.go +++ b/internal/locate/region_cache_test.go @@ -2594,6 +2594,16 @@ func (s *testRegionCacheSuite) TestSplitKeyRanges() { } func (s *testRegionCacheSuite) TestBatchScanRegions() { + s.testBatchScanRegions() +} + +func (s *testRegionCacheSuite) TestBatchScanRegionsFallback() { + s.Nil(failpoint.Enable("tikvclient/mockBatchScanRegionsUnimplemented", `return`)) + s.testBatchScanRegions() + s.Nil(failpoint.Disable("tikvclient/mockBatchScanRegionsUnimplemented")) +} + +func (s *testRegionCacheSuite) testBatchScanRegions() { // Split at "a", "b", "c", "d", "e", "f", "g" // nil --- 'a' --- 'b' --- 'c' --- 'd' --- 'e' --- 'f' --- 'g' --- nil // <- 0 -> <- 1 -> <- 2 -> <- 3 -> <- 4 -> <- 5 -> <- 6 -> <- 7 -> diff --git a/internal/mockstore/mocktikv/pd.go b/internal/mockstore/mocktikv/pd.go index f7690c98..0a1998d0 100644 --- a/internal/mockstore/mocktikv/pd.go +++ b/internal/mockstore/mocktikv/pd.go @@ -49,8 +49,11 @@ import ( rmpb "github.com/pingcap/kvproto/pkg/resource_manager" "github.com/pkg/errors" "github.com/tikv/client-go/v2/oracle" + "github.com/tikv/client-go/v2/util" pd "github.com/tikv/pd/client" "go.uber.org/atomic" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" ) // Use global variables to prevent pdClients from creating duplicate timestamps. @@ -254,6 +257,9 @@ func (c *pdClient) ScanRegions(ctx context.Context, startKey []byte, endKey []by } func (c *pdClient) BatchScanRegions(ctx context.Context, keyRanges []pd.KeyRange, limit int, opts ...pd.GetRegionOption) ([]*pd.Region, error) { + if _, err := util.EvalFailpoint("mockBatchScanRegionsUnimplemented"); err == nil { + return nil, status.Errorf(codes.Unimplemented, "mock BatchScanRegions is not implemented") + } regions := make([]*pd.Region, 0, len(keyRanges)) var lastRegion *pd.Region for _, keyRange := range keyRanges {