fix bug

Signed-off-by: you06 <[email protected]> fix bug Signed-off-by: you06 <[email protected]> remove with leader option Signed-off-by: you06 <[email protected]> update go.mod Signed-off-by: you06 <[email protected]> update client-go Signed-off-by: you06 <[email protected]> bazel Signed-off-by: you06 <[email protected]> comment public func Signed-off-by: you06 <[email protected]> update client-go Signed-off-by: you06 <[email protected]> update bazel Signed-off-by: you06 <[email protected]> deprecate ScanRegions Signed-off-by: you06 <[email protected]> bazel Signed-off-by: you06 <[email protected]> sort imports Signed-off-by: you06 <[email protected]> impl BatchScanRegions for MockPDClientForSplit Signed-off-by: you06 <[email protected]> fix injected error Signed-off-by: you06 <[email protected]> impl BatchScanRegions for FakePDClient Signed-off-by: you06 <[email protected]>
pingcap · Jun 24, 2024 · 006cbc9 · 006cbc9
1 parent ea565fc
commit 006cbc9
Show file tree

Hide file tree

Showing 16 changed files with 178 additions and 81 deletions.
diff --git a/DEPS.bzl b/DEPS.bzl
@@ -5984,13 +5984,13 @@ def go_deps():
         name = "com_github_pingcap_kvproto",
         build_file_proto_mode = "disable_global",
         importpath = "github.com/pingcap/kvproto",
-        sha256 = "ab79db554bdb20f7ebb44cd88cac36f3bf56a9d087912eda252fe40c7c828a55",
-        strip_prefix = "github.com/pingcap/[email protected]20240513094934-d9297553c900",
+        sha256 = "9c11f4c5354534ac08526d61d180bf873786fef9089b1f018fc7365350839207",
+        strip_prefix = "github.com/pingcap/[email protected]20240620063548-118a4cab53e4",
         urls = [
-            "http://bazel-cache.pingcap.net:8080/gomod/github.com/pingcap/kvproto/com_github_pingcap_kvproto-v0.0.0-20240513094934-d9297553c900.zip",
-            "http://ats.apps.svc/gomod/github.com/pingcap/kvproto/com_github_pingcap_kvproto-v0.0.0-20240513094934-d9297553c900.zip",
-            "https://cache.hawkingrei.com/gomod/github.com/pingcap/kvproto/com_github_pingcap_kvproto-v0.0.0-20240513094934-d9297553c900.zip",
-            "https://storage.googleapis.com/pingcapmirror/gomod/github.com/pingcap/kvproto/com_github_pingcap_kvproto-v0.0.0-20240513094934-d9297553c900.zip",
+            "http://bazel-cache.pingcap.net:8080/gomod/github.com/pingcap/kvproto/com_github_pingcap_kvproto-v0.0.0-20240620063548-118a4cab53e4.zip",
+            "http://ats.apps.svc/gomod/github.com/pingcap/kvproto/com_github_pingcap_kvproto-v0.0.0-20240620063548-118a4cab53e4.zip",
+            "https://cache.hawkingrei.com/gomod/github.com/pingcap/kvproto/com_github_pingcap_kvproto-v0.0.0-20240620063548-118a4cab53e4.zip",
+            "https://storage.googleapis.com/pingcapmirror/gomod/github.com/pingcap/kvproto/com_github_pingcap_kvproto-v0.0.0-20240620063548-118a4cab53e4.zip",
         ],
     )
     go_repository(
@@ -7180,26 +7180,26 @@ def go_deps():
         name = "com_github_tikv_client_go_v2",
         build_file_proto_mode = "disable_global",
         importpath = "github.com/tikv/client-go/v2",
-        sha256 = "80d99c005eaf307c2389690c1d8b6c3aadf982770a449d35e18733ee09159bcd",
-        strip_prefix = "github.com/tikv/client-go/[email protected].8-0.20240621090319-d4f0f4cf12a9",
+        sha256 = "f426f7eb6f28ca30da9240db5f103d630e960d8f185abf2afa14e11d0b2d1c65",
+        strip_prefix = "github.com/you06/client-go/[email protected].0-alpha.0.20240624045708-f42c9d878488",
         urls = [
-            "http://bazel-cache.pingcap.net:8080/gomod/github.com/tikv/client-go/v2/com_github_tikv_client_go_v2-v2.0.8-0.20240621090319-d4f0f4cf12a9.zip",
-            "http://ats.apps.svc/gomod/github.com/tikv/client-go/v2/com_github_tikv_client_go_v2-v2.0.8-0.20240621090319-d4f0f4cf12a9.zip",
-            "https://cache.hawkingrei.com/gomod/github.com/tikv/client-go/v2/com_github_tikv_client_go_v2-v2.0.8-0.20240621090319-d4f0f4cf12a9.zip",
-            "https://storage.googleapis.com/pingcapmirror/gomod/github.com/tikv/client-go/v2/com_github_tikv_client_go_v2-v2.0.8-0.20240621090319-d4f0f4cf12a9.zip",
+            "http://bazel-cache.pingcap.net:8080/gomod/github.com/you06/client-go/v2/com_github_you06_client_go_v2-v2.0.0-alpha.0.20240624045708-f42c9d878488.zip",
+            "http://ats.apps.svc/gomod/github.com/you06/client-go/v2/com_github_you06_client_go_v2-v2.0.0-alpha.0.20240624045708-f42c9d878488.zip",
+            "https://cache.hawkingrei.com/gomod/github.com/you06/client-go/v2/com_github_you06_client_go_v2-v2.0.0-alpha.0.20240624045708-f42c9d878488.zip",
+            "https://storage.googleapis.com/pingcapmirror/gomod/github.com/you06/client-go/v2/com_github_you06_client_go_v2-v2.0.0-alpha.0.20240624045708-f42c9d878488.zip",
         ],
     )
     go_repository(
         name = "com_github_tikv_pd_client",
         build_file_proto_mode = "disable_global",
         importpath = "github.com/tikv/pd/client",
-        sha256 = "e05c2d0d2b557d7b5aa6479a48d0455c2947eb2b277519282412d63b23920c02",
-        strip_prefix = "github.com/tikv/pd/[email protected]20240603082825-a929a546a790",
+        sha256 = "107218d0f9404ed51c23fda05b4fdac1bd3dbe8c8724c63ceddb4cbc7e77ceaf",
+        strip_prefix = "github.com/tikv/pd/[email protected]20240620115049-049de1761e56",
         urls = [
-            "http://bazel-cache.pingcap.net:8080/gomod/github.com/tikv/pd/client/com_github_tikv_pd_client-v0.0.0-20240603082825-a929a546a790.zip",
-            "http://ats.apps.svc/gomod/github.com/tikv/pd/client/com_github_tikv_pd_client-v0.0.0-20240603082825-a929a546a790.zip",
-            "https://cache.hawkingrei.com/gomod/github.com/tikv/pd/client/com_github_tikv_pd_client-v0.0.0-20240603082825-a929a546a790.zip",
-            "https://storage.googleapis.com/pingcapmirror/gomod/github.com/tikv/pd/client/com_github_tikv_pd_client-v0.0.0-20240603082825-a929a546a790.zip",
+            "http://bazel-cache.pingcap.net:8080/gomod/github.com/tikv/pd/client/com_github_tikv_pd_client-v0.0.0-20240620115049-049de1761e56.zip",
+            "http://ats.apps.svc/gomod/github.com/tikv/pd/client/com_github_tikv_pd_client-v0.0.0-20240620115049-049de1761e56.zip",
+            "https://cache.hawkingrei.com/gomod/github.com/tikv/pd/client/com_github_tikv_pd_client-v0.0.0-20240620115049-049de1761e56.zip",
+            "https://storage.googleapis.com/pingcapmirror/gomod/github.com/tikv/pd/client/com_github_tikv_pd_client-v0.0.0-20240620115049-049de1761e56.zip",
         ],
     )
     go_repository(

diff --git a/br/pkg/restore/split/client.go b/br/pkg/restore/split/client.go
@@ -769,7 +769,7 @@ func (c *pdClient) ScanRegions(ctx context.Context, key, endKey []byte, limit in
 		failpoint.Return(nil, status.Error(codes.Unavailable, "not leader"))
 	})
 
-	regions, err := c.client.ScanRegions(ctx, key, endKey, limit)
+	regions, err := c.client.BatchScanRegions(ctx, []pd.KeyRange{{StartKey: key, EndKey: endKey}}, limit)
 	if err != nil {
 		return nil, errors.Trace(err)
 	}

diff --git a/br/pkg/restore/split/mock_pd_client.go b/br/pkg/restore/split/mock_pd_client.go
@@ -3,6 +3,7 @@
 package split
 
 import (
+	"bytes"
 	"context"
 	"sync"
 
@@ -117,6 +118,48 @@ func (c *MockPDClientForSplit) ScanRegions(
 	return ret, nil
 }
 
+func (c *MockPDClientForSplit) BatchScanRegions(
+	_ context.Context,
+	keyRanges []pd.KeyRange,
+	limit int,
+	_ ...pd.GetRegionOption,
+) ([]*pd.Region, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	if len(c.scanRegions.errors) > 0 {
+		err := c.scanRegions.errors[0]
+		c.scanRegions.errors = c.scanRegions.errors[1:]
+		return nil, err
+	}
+
+	if c.scanRegions.beforeHook != nil {
+		c.scanRegions.beforeHook()
+	}
+
+	regions := make([]*pd.Region, 0, len(keyRanges))
+	var lastRegion *pdtypes.Region
+	for _, keyRange := range keyRanges {
+		if lastRegion != nil {
+			if len(lastRegion.Meta.EndKey) == 0 || bytes.Compare(lastRegion.Meta.EndKey, keyRange.EndKey) >= 0 {
+				continue
+			}
+			if bytes.Compare(lastRegion.Meta.EndKey, keyRange.StartKey) > 0 {
+				keyRange.StartKey = lastRegion.Meta.EndKey
+			}
+		}
+		rs := c.Regions.ScanRange(keyRange.StartKey, keyRange.EndKey, limit)
+		for _, r := range rs {
+			lastRegion = r
+			regions = append(regions, &pd.Region{
+				Meta:   r.Meta,
+				Leader: r.Leader,
+			})
+		}
+	}
+	return regions, nil
+}
+
 func (c *MockPDClientForSplit) GetRegionByID(_ context.Context, regionID uint64, _ ...pd.GetRegionOption) (*pd.Region, error) {
 	c.mu.Lock()
 	defer c.mu.Unlock()

diff --git a/br/pkg/streamhelper/advancer_env.go b/br/pkg/streamhelper/advancer_env.go
@@ -73,7 +73,7 @@ func (c PDRegionScanner) FetchCurrentTS(ctx context.Context) (uint64, error) {
 // RegionScan gets a list of regions, starts from the region that contains key.
 // Limit limits the maximum number of regions returned.
 func (c PDRegionScanner) RegionScan(ctx context.Context, key, endKey []byte, limit int) ([]RegionWithLeader, error) {
-	rs, err := c.Client.ScanRegions(ctx, key, endKey, limit)
+	rs, err := c.Client.BatchScanRegions(ctx, []pd.KeyRange{{StartKey: key, EndKey: endKey}}, limit)
 	if err != nil {
 		return nil, err
 	}

diff --git a/br/pkg/utiltest/fake.go b/br/pkg/utiltest/fake.go
@@ -202,6 +202,34 @@ func (fpdc *FakePDClient) ScanRegions(
 	return regions, nil
 }
 
+func (fpdc *FakePDClient) BatchScanRegions(
+	ctx context.Context,
+	ranges []pd.KeyRange,
+	limit int,
+	opts ...pd.GetRegionOption,
+) ([]*pd.Region, error) {
+	regions := make([]*pd.Region, 0, len(fpdc.regions))
+	fpdc.peerStoreId = fpdc.peerStoreId + 1
+	peerStoreId := (fpdc.peerStoreId + 1) / 2
+	for _, region := range fpdc.regions {
+		inRange := false
+		for _, keyRange := range ranges {
+			if len(keyRange.EndKey) != 0 && bytes.Compare(region.Meta.StartKey, keyRange.EndKey) >= 0 {
+				continue
+			}
+			if len(region.Meta.EndKey) != 0 && bytes.Compare(region.Meta.EndKey, keyRange.StartKey) <= 0 {
+				continue
+			}
+			inRange = true
+		}
+		if inRange {
+			region.Meta.Peers = []*metapb.Peer{{StoreId: peerStoreId}}
+			regions = append(regions, region)
+		}
+	}
+	return nil, nil
+}
+
 func (fpdc *FakePDClient) GetTS(ctx context.Context) (int64, int64, error) {
 	(*fpdc.retryTimes)++
 	if *fpdc.retryTimes >= 3 { // the mock PD leader switched successfully

diff --git a/br/tests/br_key_locked/codec.go b/br/tests/br_key_locked/codec.go
@@ -63,7 +63,7 @@ func (c *codecPDClient) ScanRegions(
 		endKey = codec.EncodeBytes(nil, endKey)
 	}
 
-	regions, err := c.Client.ScanRegions(ctx, startKey, endKey, limit, opts...)
+	regions, err := c.Client.BatchScanRegions(ctx, []pd.KeyRange{{StartKey: startKey, EndKey: endKey}}, limit, opts...)
 	if err != nil {
 		return nil, errors.Trace(err)
 	}

diff --git a/go.mod b/go.mod
@@ -84,7 +84,7 @@ require (
 	github.com/pingcap/errors v0.11.5-0.20240318064555-6bd07397691f
 	github.com/pingcap/failpoint v0.0.0-20240527053858-9b3b6e34194a
 	github.com/pingcap/fn v1.0.0
-	github.com/pingcap/kvproto v0.0.0-20240513094934-d9297553c900
+	github.com/pingcap/kvproto v0.0.0-20240620063548-118a4cab53e4
 	github.com/pingcap/log v1.1.1-0.20240314023424-862ccc32f18d
 	github.com/pingcap/sysutil v1.0.1-0.20240311050922-ae81ee01f3a5
 	github.com/pingcap/tidb/pkg/parser v0.0.0-20211011031125-9b13dc409c5e
@@ -108,7 +108,7 @@ require (
 	github.com/tiancaiamao/appdash v0.0.0-20181126055449-889f96f722a2
 	github.com/tidwall/btree v1.7.0
 	github.com/tikv/client-go/v2 v2.0.8-0.20240621090319-d4f0f4cf12a9
-	github.com/tikv/pd/client v0.0.0-20240603082825-a929a546a790
+	github.com/tikv/pd/client v0.0.0-20240620115049-049de1761e56
 	github.com/timakin/bodyclose v0.0.0-20240125160201-f835fa56326a
 	github.com/twmb/murmur3 v1.1.6
 	github.com/uber/jaeger-client-go v2.22.1+incompatible
@@ -331,3 +331,5 @@ replace (
 	sourcegraph.com/sourcegraph/appdash => github.com/sourcegraph/appdash v0.0.0-20190731080439-ebfcffb1b5c0
 	sourcegraph.com/sourcegraph/appdash-data => github.com/sourcegraph/appdash-data v0.0.0-20151005221446-73f23eafcf67
 )
+
+replace github.com/tikv/client-go/v2 => github.com/you06/client-go/v2 v2.0.0-alpha.0.20240624045708-f42c9d878488
diff --git a/go.sum b/go.sum
@@ -690,8 +690,8 @@ github.com/pingcap/fn v1.0.0/go.mod h1:u9WZ1ZiOD1RpNhcI42RucFh/lBuzTu6rw88a+oF2Z
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989 h1:surzm05a8C9dN8dIUmo4Be2+pMRb6f55i+UIYrluu2E=
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989/go.mod h1:O17XtbryoCJhkKGbT62+L2OlrniwqiGLSqrmdHCMzZw=
 github.com/pingcap/kvproto v0.0.0-20191211054548-3c6b38ea5107/go.mod h1:WWLmULLO7l8IOcQG+t+ItJ3fEcrL5FxF0Wu+HrMy26w=
-github.com/pingcap/kvproto v0.0.0-20240513094934-d9297553c900 h1:snIM8DC846ufdlRclITACXfr1kvVIPU4cuQ6w3JVVY4=
-github.com/pingcap/kvproto v0.0.0-20240513094934-d9297553c900/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
+github.com/pingcap/kvproto v0.0.0-20240620063548-118a4cab53e4 h1:6aIKNB2YGAec4IUDLw6G2eDECiGiufZcgEbZSCELBx0=
+github.com/pingcap/kvproto v0.0.0-20240620063548-118a4cab53e4/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
 github.com/pingcap/log v0.0.0-20210625125904-98ed8e2eb1c7/go.mod h1:8AanEdAHATuRurdGxZXBz0At+9avep+ub7U1AGYLIMM=
 github.com/pingcap/log v1.1.0/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=
 github.com/pingcap/log v1.1.1-0.20240314023424-862ccc32f18d h1:y3EueKVfVykdpTyfUnQGqft0ud+xVFuCdp1XkVL0X1E=
@@ -846,10 +846,8 @@ github.com/tiancaiamao/gp v0.0.0-20221230034425-4025bc8a4d4a h1:J/YdBZ46WKpXsxsW
 github.com/tiancaiamao/gp v0.0.0-20221230034425-4025bc8a4d4a/go.mod h1:h4xBhSNtOeEosLJ4P7JyKXX7Cabg7AVkWCK5gV2vOrM=
 github.com/tidwall/btree v1.7.0 h1:L1fkJH/AuEh5zBnnBbmTwQ5Lt+bRJ5A8EWecslvo9iI=
 github.com/tidwall/btree v1.7.0/go.mod h1:twD9XRA5jj9VUQGELzDO4HPQTNJsoWWfYEL+EUQ2cKY=
-github.com/tikv/client-go/v2 v2.0.8-0.20240621090319-d4f0f4cf12a9 h1:hf5Tjut3Dht8ox5zB66Q1mO2GWd39P6GUtSG3xT0kAA=
-github.com/tikv/client-go/v2 v2.0.8-0.20240621090319-d4f0f4cf12a9/go.mod h1:eQKY0fe8dXTiAvdP0O7xaPF3clw/49Z2cSQqpKCDxvY=
-github.com/tikv/pd/client v0.0.0-20240603082825-a929a546a790 h1:bGmvWcMkbOlVgWpsXza2gu18Ud2dEyTz60UU2oEUSoA=
-github.com/tikv/pd/client v0.0.0-20240603082825-a929a546a790/go.mod h1:kNRekhwXqjTjNHy+kPmbZvsMmvl42zOj/UW5IIG+nP0=
+github.com/tikv/pd/client v0.0.0-20240620115049-049de1761e56 h1:7TLLfwrKoty9UeJMsSopRTXYw8ooxcF0Z1fegXhIgks=
+github.com/tikv/pd/client v0.0.0-20240620115049-049de1761e56/go.mod h1:EHHidLItrJGh0jqfdfFhIHG5vwkR8+43tFnp7v7iv1Q=
 github.com/timakin/bodyclose v0.0.0-20240125160201-f835fa56326a h1:A6uKudFIfAEpoPdaal3aSqGxBzLyU8TqyXImLwo6dIo=
 github.com/timakin/bodyclose v0.0.0-20240125160201-f835fa56326a/go.mod h1:mkjARE7Yr8qU23YcGMSALbIxTQ9r9QBVahQOBRfU460=
 github.com/tklauser/go-sysconf v0.3.9/go.mod h1:11DU/5sG7UexIrp/O6g35hrWzu0JxlwQ3LSFUzyeuhs=
@@ -882,6 +880,8 @@ github.com/xitongsys/parquet-go-source v0.0.0-20190524061010-2b72cbee77d5/go.mod
 github.com/xitongsys/parquet-go-source v0.0.0-20200817004010-026bad9b25d0 h1:a742S4V5A15F93smuVxA60LQWsrCnN8bKeWDBARU1/k=
 github.com/xitongsys/parquet-go-source v0.0.0-20200817004010-026bad9b25d0/go.mod h1:HYhIKsdns7xz80OgkbgJYrtQY7FjHWHKH6cvN7+czGE=
 github.com/xordataexchange/crypt v0.0.3-0.20170626215501-b2862e3d0a77/go.mod h1:aYKd//L2LvnjZzWKhF00oedf4jCCReLcmhLdhm1A27Q=
+github.com/you06/client-go/v2 v2.0.0-alpha.0.20240624045708-f42c9d878488 h1:+E55rN2EA1UVSAVwT4mMy76vPAOZkvyUvWuc1OThD2M=
+github.com/you06/client-go/v2 v2.0.0-alpha.0.20240624045708-f42c9d878488/go.mod h1:4HDOAx8OXAJPtqhCZ03IhChXgaFs4B3+vSrPWmiPxjg=
 github.com/yuin/goldmark v1.1.25/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.1.32/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=

diff --git a/pkg/ddl/partition.go b/pkg/ddl/partition.go
@@ -65,6 +65,7 @@ import (
 	"github.com/pingcap/tidb/pkg/util/stringutil"
 	"github.com/tikv/client-go/v2/tikv"
 	kvutil "github.com/tikv/client-go/v2/util"
+	pd "github.com/tikv/pd/client"
 	"go.uber.org/zap"
 )
 
@@ -435,9 +436,9 @@ func checkPartitionReplica(replicaCount uint64, addingDefinitions []model.Partit
 	if replicaCount > tiFlashStoreCount {
 		return false, errors.Errorf("[ddl] the tiflash replica count: %d should be less than the total tiflash server count: %d", replicaCount, tiFlashStoreCount)
 	}
-	for _, pd := range addingDefinitions {
-		startKey, endKey := tablecodec.GetTableHandleKeyRange(pd.ID)
-		regions, err := pdCli.ScanRegions(ctx, startKey, endKey, -1)
+	for _, pDef := range addingDefinitions {
+		startKey, endKey := tablecodec.GetTableHandleKeyRange(pDef.ID)
+		regions, err := pdCli.BatchScanRegions(ctx, []pd.KeyRange{{StartKey: startKey, EndKey: endKey}}, -1)
 		if err != nil {
 			return needWait, errors.Trace(err)
 		}
@@ -458,7 +459,7 @@ func checkPartitionReplica(replicaCount uint64, addingDefinitions []model.Partit
 				continue
 			}
 			needWait = true
-			logutil.DDLLogger().Info("partition replicas check failed in replica-only DDL state", zap.Int64("pID", pd.ID), zap.Uint64("wait region ID", region.Meta.Id), zap.Bool("tiflash peer at least one", tiflashPeerAtLeastOne), zap.Time("check time", time.Now()))
+			logutil.DDLLogger().Info("partition replicas check failed in replica-only DDL state", zap.Int64("pID", pDef.ID), zap.Uint64("wait region ID", region.Meta.Id), zap.Bool("tiflash peer at least one", tiflashPeerAtLeastOne), zap.Time("check time", time.Now()))
 			return needWait, nil
 		}
 	}

diff --git a/pkg/ddl/tests/tiflash/BUILD.bazel b/pkg/ddl/tests/tiflash/BUILD.bazel
@@ -38,6 +38,7 @@ go_test(
         "@com_github_tikv_client_go_v2//oracle",
         "@com_github_tikv_client_go_v2//testutils",
         "@com_github_tikv_client_go_v2//tikv",
+        "@com_github_tikv_pd_client//:client",
         "@org_uber_go_goleak//:goleak",
         "@org_uber_go_zap//:zap",
     ],

diff --git a/pkg/ddl/tests/tiflash/ddl_tiflash_test.go b/pkg/ddl/tests/tiflash/ddl_tiflash_test.go
@@ -52,6 +52,7 @@ import (
 	"github.com/tikv/client-go/v2/oracle"
 	"github.com/tikv/client-go/v2/testutils"
 	"github.com/tikv/client-go/v2/tikv"
+	pd "github.com/tikv/pd/client"
 	"go.uber.org/zap"
 )
 
@@ -1460,9 +1461,9 @@ func TestTiFlashReorgPartition(t *testing.T) {
 				_ = job.DecodeArgs(&dummy, &partInfo)
 				ctx := context.Background()
 				stores, _ := pdCli.GetAllStores(ctx)
-				for _, pd := range partInfo.Definitions {
-					startKey, endKey := tablecodec.GetTableHandleKeyRange(pd.ID)
-					regions, _ := pdCli.ScanRegions(ctx, startKey, endKey, -1)
+				for _, pDef := range partInfo.Definitions {
+					startKey, endKey := tablecodec.GetTableHandleKeyRange(pDef.ID)
+					regions, _ := pdCli.BatchScanRegions(ctx, []pd.KeyRange{{StartKey: startKey, EndKey: endKey}}, -1)
 					for i := range regions {
 						// similar as storeHasEngineTiFlashLabel
 						for _, store := range stores {

diff --git a/pkg/server/handler/tikvhandler/tikv_handler.go b/pkg/server/handler/tikvhandler/tikv_handler.go
@@ -1336,7 +1336,7 @@ func (h *TableHandler) getRegionsByID(tbl table.Table, id int64, name string) (*
 	startKey, endKey := tablecodec.GetTableHandleKeyRange(id)
 	ctx := context.Background()
 	pdCli := h.RegionCache.PDClient()
-	regions, err := pdCli.ScanRegions(ctx, startKey, endKey, -1)
+	regions, err := pdCli.BatchScanRegions(ctx, []pd.KeyRange{{StartKey: startKey, EndKey: endKey}}, -1)
 	if err != nil {
 		return nil, err
 	}
@@ -1359,7 +1359,7 @@ func (h *TableHandler) getRegionsByID(tbl table.Table, id int64, name string) (*
 		indices[i].Name = index.Meta().Name.String()
 		indices[i].ID = indexID
 		startKey, endKey := tablecodec.GetTableIndexKeyRange(id, indexID)
-		regions, err := pdCli.ScanRegions(ctx, startKey, endKey, -1)
+		regions, err := pdCli.BatchScanRegions(ctx, []pd.KeyRange{{StartKey: startKey, EndKey: endKey}}, -1)
 		if err != nil {
 			return nil, err
 		}

diff --git a/pkg/store/copr/batch_coprocessor.go b/pkg/store/copr/batch_coprocessor.go
@@ -635,7 +635,7 @@ func buildBatchCopTasksConsistentHash(
 
 	for i, ranges := range rangesForEachPhysicalTable {
 		rangesLen += ranges.Len()
-		locations, err := cache.SplitKeyRangesByLocationsWithoutBuckets(bo, ranges, UnspecifiedLimit)
+		locations, err := cache.SplitKeyRangesByLocations(bo, ranges, UnspecifiedLimit, false, false)
 		if err != nil {
 			return nil, errors.Trace(err)
 		}
@@ -932,7 +932,7 @@ func buildBatchCopTasksCore(bo *backoff.Backoffer, store *kvStore, rangesForEach
 		rangesLen = 0
 		for i, ranges := range rangesForEachPhysicalTable {
 			rangesLen += ranges.Len()
-			locations, err := cache.SplitKeyRangesByLocationsWithoutBuckets(bo, ranges, UnspecifiedLimit)
+			locations, err := cache.SplitKeyRangesByLocations(bo, ranges, UnspecifiedLimit, false, false)
 			if err != nil {
 				return nil, errors.Trace(err)
 			}
@@ -1483,7 +1483,7 @@ func buildBatchCopTasksConsistentHashForPD(bo *backoff.Backoffer,
 		splitKeyStart := time.Now()
 		for i, ranges := range rangesForEachPhysicalTable {
 			rangesLen += ranges.Len()
-			locations, err := cache.SplitKeyRangesByLocationsWithoutBuckets(bo, ranges, UnspecifiedLimit)
+			locations, err := cache.SplitKeyRangesByLocations(bo, ranges, UnspecifiedLimit, false, false)
 			if err != nil {
 				return nil, errors.Trace(err)
 			}