Update the PD client and adopt the latest changes

Signed-off-by: JmPotato <[email protected]>
tikv · Nov 30, 2023 · aedb7a4 · aedb7a4
1 parent 812a791
commit aedb7a4
Show file tree

Hide file tree

Showing 6 changed files with 11 additions and 9 deletions.
diff --git a/go.mod b/go.mod
@@ -21,7 +21,7 @@ require (
 	github.com/prometheus/client_model v0.3.0
 	github.com/stretchr/testify v1.8.2
 	github.com/tiancaiamao/gp v0.0.0-20221230034425-4025bc8a4d4a
-	github.com/tikv/pd/client v0.0.0-20231116034417-4e600c227e83
+	github.com/tikv/pd/client v0.0.0-20231130081618-862eee18738e
 	github.com/twmb/murmur3 v1.1.3
 	go.etcd.io/etcd/api/v3 v3.5.10
 	go.etcd.io/etcd/client/v3 v3.5.10

diff --git a/go.sum b/go.sum
@@ -114,6 +114,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/tikv/pd/client v0.0.0-20231116034417-4e600c227e83 h1:wPZW67pTWz2yEboCw2dK+z4p/tGFVCTqRl5WkNCK0lc=
 github.com/tikv/pd/client v0.0.0-20231116034417-4e600c227e83/go.mod h1:cd6zBqRM9aogxf26K8NnFRPVtq9BnRE59tKEpX8IaWQ=
+github.com/tikv/pd/client v0.0.0-20231130081618-862eee18738e h1:11cWLLmEreKof/VJi6LLQ+Jkav5ZqPJgeI+KX4pc/DE=
+github.com/tikv/pd/client v0.0.0-20231130081618-862eee18738e/go.mod h1:cd6zBqRM9aogxf26K8NnFRPVtq9BnRE59tKEpX8IaWQ=
 github.com/twmb/murmur3 v1.1.3 h1:D83U0XYKcHRYwYIpBKf3Pks91Z0Byda/9SJ8B6EMRcA=
 github.com/twmb/murmur3 v1.1.3/go.mod h1:Qq/R7NUyOfr65zD+6Q5IHKsJLwP7exErjN6lyyq3OSQ=
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=

diff --git a/internal/locate/pd_codec.go b/internal/locate/pd_codec.go
@@ -115,9 +115,9 @@ func (c *CodecPDClient) GetRegionByID(ctx context.Context, regionID uint64, opts
 
 // ScanRegions encodes the key before send requests to pd-server and decodes the
 // returned StartKey && EndKey from pd-server.
-func (c *CodecPDClient) ScanRegions(ctx context.Context, startKey []byte, endKey []byte, limit int) ([]*pd.Region, error) {
+func (c *CodecPDClient) ScanRegions(ctx context.Context, startKey []byte, endKey []byte, limit int, opts ...pd.GetRegionOption) ([]*pd.Region, error) {
 	startKey, endKey = c.codec.EncodeRegionRange(startKey, endKey)
-	regions, err := c.Client.ScanRegions(ctx, startKey, endKey, limit)
+	regions, err := c.Client.ScanRegions(ctx, startKey, endKey, limit, opts...)
 	if err != nil {
 		return nil, errors.WithStack(err)
 	}

diff --git a/internal/mockstore/mocktikv/cluster.go b/internal/mockstore/mocktikv/cluster.go
@@ -349,7 +349,7 @@ func (c *Cluster) GetRegionByID(regionID uint64) (*metapb.Region, *metapb.Peer,
 }
 
 // ScanRegions returns at most `limit` regions from given `key` and their leaders.
-func (c *Cluster) ScanRegions(startKey, endKey []byte, limit int) []*pd.Region {
+func (c *Cluster) ScanRegions(startKey, endKey []byte, limit int, _ ...pd.GetRegionOption) []*pd.Region {
 	c.RLock()
 	defer c.RUnlock()
 

diff --git a/internal/mockstore/mocktikv/pd.go b/internal/mockstore/mocktikv/pd.go
@@ -209,7 +209,7 @@ func (c *pdClient) GetRegion(ctx context.Context, key []byte, opts ...pd.GetRegi
 	return &pd.Region{Meta: region, Leader: peer, Buckets: buckets, DownPeers: downPeers}, nil
 }
 
-func (c *pdClient) GetRegionFromMember(ctx context.Context, key []byte, memberURLs []string) (*pd.Region, error) {
+func (c *pdClient) GetRegionFromMember(ctx context.Context, key []byte, memberURLs []string, opts ...pd.GetRegionOption) (*pd.Region, error) {
 	return &pd.Region{}, nil
 }
 
@@ -226,8 +226,8 @@ func (c *pdClient) GetRegionByID(ctx context.Context, regionID uint64, opts ...p
 	return &pd.Region{Meta: region, Leader: peer, Buckets: buckets, DownPeers: downPeers}, nil
 }
 
-func (c *pdClient) ScanRegions(ctx context.Context, startKey []byte, endKey []byte, limit int) ([]*pd.Region, error) {
-	regions := c.cluster.ScanRegions(startKey, endKey, limit)
+func (c *pdClient) ScanRegions(ctx context.Context, startKey []byte, endKey []byte, limit int, opts ...pd.GetRegionOption) ([]*pd.Region, error) {
+	regions := c.cluster.ScanRegions(startKey, endKey, limit, opts...)
 	return regions, nil
 }
 

diff --git a/util/pd_interceptor.go b/util/pd_interceptor.go
@@ -119,9 +119,9 @@ func (m InterceptedPDClient) GetRegionByID(ctx context.Context, regionID uint64,
 }
 
 // ScanRegions implements pd.Client#ScanRegions.
-func (m InterceptedPDClient) ScanRegions(ctx context.Context, key, endKey []byte, limit int) ([]*pd.Region, error) {
+func (m InterceptedPDClient) ScanRegions(ctx context.Context, key, endKey []byte, limit int, opts ...pd.GetRegionOption) ([]*pd.Region, error) {
 	start := time.Now()
-	r, err := m.Client.ScanRegions(ctx, key, endKey, limit)
+	r, err := m.Client.ScanRegions(ctx, key, endKey, limit, opts...)
 	recordPDWaitTime(ctx, start)
 	return r, err
 }