diff --git a/client/http/client.go b/client/http/client.go index ae0f7fee347c..1d8c2d5c427b 100644 --- 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, ®ions) 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, ®ionStats, ) 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 index 25625ac10334..dc901b1d2908 100644 --- 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"),