Use KeyRange to replace the key parameters

Signed-off-by: JmPotato <[email protected]>
tikv · Nov 27, 2023 · 520a5e8 · 520a5e8
1 parent aced6d8
commit 520a5e8
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 30 deletions.
diff --git a/client/http/client.go b/client/http/client.go
@@ -47,11 +47,11 @@ type Client interface {
 	GetRegionByID(context.Context, uint64) (*RegionInfo, error)
 	GetRegionByKey(context.Context, []byte) (*RegionInfo, error)
 	GetRegions(context.Context) (*RegionsInfo, error)
-	GetRegionsByKeyRange(context.Context, []byte, []byte, int) (*RegionsInfo, error)
+	GetRegionsByKeyRange(context.Context, *KeyRange, int) (*RegionsInfo, error)
 	GetRegionsByStoreID(context.Context, uint64) (*RegionsInfo, error)
 	GetHotReadRegions(context.Context) (*StoreHotPeersInfos, error)
 	GetHotWriteRegions(context.Context) (*StoreHotPeersInfos, error)
-	GetRegionStatusByKeyRange(context.Context, []byte, []byte) (*RegionStats, error)
+	GetRegionStatusByKeyRange(context.Context, *KeyRange) (*RegionStats, error)
 	GetStores(context.Context) (*StoresInfo, error)
 	/* Rule-related interfaces */
 	GetAllPlacementRuleBundles(context.Context) ([]*GroupBundle, error)
@@ -70,11 +70,8 @@ type Client interface {
 	SetRegionLabelRule(context.Context, *LabelRule) error
 	PatchRegionLabelRules(context.Context, *LabelRulePatch) error
 	/* Scheduling-related interfaces */
-	AccelerateSchedule(context.Context, []byte, []byte) error
-	AccelerateScheduleInBatch(context.Context, []struct {
-		StartKey []byte `json:"start_key"`
-		EndKey   []byte `json:"end_key"`
-	}) error
+	AccelerateSchedule(context.Context, *KeyRange) error
+	AccelerateScheduleInBatch(context.Context, []*KeyRange) error
 	/* Other interfaces */
 	GetMinResolvedTSByStoresIDs(context.Context, []uint64) (uint64, map[uint64]uint64, error)
 
@@ -318,10 +315,10 @@ func (c *client) GetRegions(ctx context.Context) (*RegionsInfo, error) {
 }
 
 // GetRegionsByKeyRange gets the regions info by key range. If the limit is -1, it will return all regions within the range.
-func (c *client) GetRegionsByKeyRange(ctx context.Context, startKey, endKey []byte, limit int) (*RegionsInfo, error) {
+func (c *client) GetRegionsByKeyRange(ctx context.Context, keyRange *KeyRange, limit int) (*RegionsInfo, error) {
 	var regions RegionsInfo
 	err := c.requestWithRetry(ctx,
-		"GetRegionsByKeyRange", RegionsByKey(startKey, endKey, limit),
+		"GetRegionsByKeyRange", RegionsByKey(keyRange.StartKey, keyRange.EndKey, limit),
 		http.MethodGet, http.NoBody, &regions)
 	if err != nil {
 		return nil, err
@@ -366,10 +363,10 @@ func (c *client) GetHotWriteRegions(ctx context.Context) (*StoreHotPeersInfos, e
 }
 
 // GetRegionStatusByKeyRange gets the region status by key range.
-func (c *client) GetRegionStatusByKeyRange(ctx context.Context, startKey, endKey []byte) (*RegionStats, error) {
+func (c *client) GetRegionStatusByKeyRange(ctx context.Context, keyRange *KeyRange) (*RegionStats, error) {
 	var regionStats RegionStats
 	err := c.requestWithRetry(ctx,
-		"GetRegionStatusByKeyRange", RegionStatsByKeyRange(startKey, endKey),
+		"GetRegionStatusByKeyRange", RegionStatsByKeyRange(keyRange.StartKey, keyRange.StartKey),
 		http.MethodGet, http.NoBody, &regionStats,
 	)
 	if err != nil {
@@ -560,12 +557,11 @@ func (c *client) PatchRegionLabelRules(ctx context.Context, labelRulePatch *Labe
 }
 
 // AccelerateSchedule accelerates the scheduling of the regions within the given key range.
-func (c *client) AccelerateSchedule(ctx context.Context, startKey, endKey []byte) error {
-	input := map[string]string{
-		"start_key": url.QueryEscape(hex.EncodeToString(startKey)),
-		"end_key":   url.QueryEscape(hex.EncodeToString(endKey)),
-	}
-	inputJSON, err := json.Marshal(input)
+func (c *client) AccelerateSchedule(ctx context.Context, keyRange *KeyRange) error {
+	inputJSON, err := json.Marshal(map[string]string{
+		"start_key": url.QueryEscape(hex.EncodeToString(keyRange.StartKey)),
+		"end_key":   url.QueryEscape(hex.EncodeToString(keyRange.EndKey)),
+	})
 	if err != nil {
 		return errors.Trace(err)
 	}
@@ -575,15 +571,12 @@ func (c *client) AccelerateSchedule(ctx context.Context, startKey, endKey []byte
 }
 
 // AccelerateScheduleInBatch accelerates the scheduling of the regions within the given key ranges in batch.
-func (c *client) AccelerateScheduleInBatch(ctx context.Context, ranges []struct {
-	StartKey []byte `json:"start_key"`
-	EndKey   []byte `json:"end_key"`
-}) error {
-	input := make([]map[string]string, 0, len(ranges))
-	for _, r := range ranges {
+func (c *client) AccelerateScheduleInBatch(ctx context.Context, keyRanges []*KeyRange) error {
+	input := make([]map[string]string, 0, len(keyRanges))
+	for _, keyRange := range keyRanges {
 		input = append(input, map[string]string{
-			"start_key": url.QueryEscape(hex.EncodeToString(r.StartKey)),
-			"end_key":   url.QueryEscape(hex.EncodeToString(r.EndKey)),
+			"start_key": url.QueryEscape(hex.EncodeToString(keyRange.StartKey)),
+			"end_key":   url.QueryEscape(hex.EncodeToString(keyRange.EndKey)),
 		})
 	}
 	inputJSON, err := json.Marshal(input)

diff --git a/tests/integrations/client/http_client_test.go b/tests/integrations/client/http_client_test.go
@@ -280,17 +280,16 @@ func (suite *httpClientTestSuite) TestAccelerateSchedule() {
 	}
 	suspectRegions := raftCluster.GetSuspectRegions()
 	re.Len(suspectRegions, 0)
-	err := suite.client.AccelerateSchedule(suite.ctx, []byte("a1"), []byte("a2"))
+	err := suite.client.AccelerateSchedule(suite.ctx, &pd.KeyRange{
+		StartKey: []byte("a1"),
+		EndKey:   []byte("a2")})
 	re.NoError(err)
 	suspectRegions = raftCluster.GetSuspectRegions()
 	re.Len(suspectRegions, 1)
 	raftCluster.ClearSuspectRegions()
 	suspectRegions = raftCluster.GetSuspectRegions()
 	re.Len(suspectRegions, 0)
-	err = suite.client.AccelerateScheduleInBatch(suite.ctx, []struct {
-		StartKey []byte `json:"start_key"`
-		EndKey   []byte `json:"end_key"`
-	}{
+	err = suite.client.AccelerateScheduleInBatch(suite.ctx, []*pd.KeyRange{
 		{
 			StartKey: []byte("a1"),
 			EndKey:   []byte("a2"),