Use KeyRange to differentiate the encoding format of the key

Signed-off-by: JmPotato <[email protected]>

client/http/api.go

@@ -84,12 +84,9 @@ func RegionByKey(key []byte) string {
 }
 
 // RegionsByKey returns the path of PD HTTP API to scan regions with given start key, end key and limit parameters.
-func RegionsByKey(startKey, endKey []byte, limit int) string {
+func RegionsByKey(startKey, endKey string, limit int) string {
 	return fmt.Sprintf("%s?start_key=%s&end_key=%s&limit=%d",
-		regionsByKey,
-		url.QueryEscape(string(startKey)),
-		url.QueryEscape(string(endKey)),
-		limit)
+		regionsByKey, startKey, endKey, limit)
 }
 
 // RegionsByStoreID returns the path of PD HTTP API to get regions by store ID.
@@ -98,11 +95,9 @@ func RegionsByStoreID(storeID uint64) string {
 }
 
 // RegionStatsByKeyRange returns the path of PD HTTP API to get region stats by start key and end key.
-func RegionStatsByKeyRange(startKey, endKey []byte) string {
+func RegionStatsByKeyRange(startKey, endKey string) string {
 	return fmt.Sprintf("%s?start_key=%s&end_key=%s",
-		StatsRegion,
-		url.QueryEscape(string(startKey)),
-		url.QueryEscape(string(endKey)))
+		StatsRegion, startKey, endKey)
 }
 
 // StoreByID returns the store API with store ID parameter.

client/http/client.go

@@ -18,7 +18,6 @@ import (
 	"bytes"
 	"context"
 	"crypto/tls"
-	"encoding/hex"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -315,10 +314,14 @@ 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.
+// The keys in the key range should be encoded in the UTF-8 bytes format.
 func (c *client) GetRegionsByKeyRange(ctx context.Context, keyRange *KeyRange, limit int) (*RegionsInfo, error) {
-	var regions RegionsInfo
+	var (
+		regions          RegionsInfo
+		startKey, endKey = keyRange.EscapeAsUTF8Str()
+	)
 	err := c.requestWithRetry(ctx,
-		"GetRegionsByKeyRange", RegionsByKey(keyRange.StartKey, keyRange.EndKey, limit),
+		"GetRegionsByKeyRange", RegionsByKey(startKey, endKey, limit),
 		http.MethodGet, http.NoBody, &regions)
 	if err != nil {
 		return nil, err
@@ -363,10 +366,14 @@ func (c *client) GetHotWriteRegions(ctx context.Context) (*StoreHotPeersInfos, e
 }
 
 // GetRegionStatusByKeyRange gets the region status by key range.
+// The keys in the key range should be encoded in the UTF-8 bytes format.
 func (c *client) GetRegionStatusByKeyRange(ctx context.Context, keyRange *KeyRange) (*RegionStats, error) {
-	var regionStats RegionStats
+	var (
+		regionStats      RegionStats
+		startKey, endKey = keyRange.EscapeAsUTF8Str()
+	)
 	err := c.requestWithRetry(ctx,
-		"GetRegionStatusByKeyRange", RegionStatsByKeyRange(keyRange.StartKey, keyRange.StartKey),
+		"GetRegionStatusByKeyRange", RegionStatsByKeyRange(startKey, endKey),
 		http.MethodGet, http.NoBody, &regionStats,
 	)
 	if err != nil {
@@ -557,10 +564,12 @@ func (c *client) PatchRegionLabelRules(ctx context.Context, labelRulePatch *Labe
 }
 
 // AccelerateSchedule accelerates the scheduling of the regions within the given key range.
+// The keys in the key range should be encoded in the hex bytes format (without encoding to the UTF-8 bytes).
 func (c *client) AccelerateSchedule(ctx context.Context, keyRange *KeyRange) error {
+	startKey, endKey := keyRange.EscapeAsHexStr()
 	inputJSON, err := json.Marshal(map[string]string{
-		"start_key": url.QueryEscape(hex.EncodeToString(keyRange.StartKey)),
-		"end_key":   url.QueryEscape(hex.EncodeToString(keyRange.EndKey)),
+		"start_key": url.QueryEscape(startKey),
+		"end_key":   url.QueryEscape(endKey),
 	})
 	if err != nil {
 		return errors.Trace(err)
@@ -571,12 +580,14 @@ func (c *client) AccelerateSchedule(ctx context.Context, keyRange *KeyRange) err
 }
 
 // AccelerateScheduleInBatch accelerates the scheduling of the regions within the given key ranges in batch.
+// The keys in the key ranges should be encoded in the hex bytes format (without encoding to the UTF-8 bytes).
 func (c *client) AccelerateScheduleInBatch(ctx context.Context, keyRanges []*KeyRange) error {
 	input := make([]map[string]string, 0, len(keyRanges))
 	for _, keyRange := range keyRanges {
+		startKey, endKey := keyRange.EscapeAsHexStr()
 		input = append(input, map[string]string{
-			"start_key": url.QueryEscape(hex.EncodeToString(keyRange.StartKey)),
-			"end_key":   url.QueryEscape(hex.EncodeToString(keyRange.EndKey)),
+			"start_key": url.QueryEscape(startKey),
+			"end_key":   url.QueryEscape(endKey),
 		})
 	}
 	inputJSON, err := json.Marshal(input)

client/http/types.go

@@ -15,14 +15,42 @@
 package http
 
 import (
+	"encoding/hex"
 	"encoding/json"
+	"net/url"
 	"time"
 )
 
-// KeyRange defines a range of keys.
+// KeyRange defines a range of keys in bytes.
 type KeyRange struct {
-	StartKey []byte `json:"start_key"`
-	EndKey   []byte `json:"end_key"`
+	startKey []byte
+	endKey   []byte
+}
+
+// NewKeyRange creates a new key range structure with the given start key and end key bytes.
+// Notice: the actual encoding of the key range is not specified here. It should be either UTF-8 or hex.
+//   - UTF-8 means the key has already been encoded into a string with UTF-8 encoding, like:
+//     []byte{52 56 54 53 54 99 54 99 54 102 50 48 53 55 54 102 55 50 54 99 54 52}, which will later be converted to "48656c6c6f20576f726c64"
+//     by using `string()` method.
+//   - Hex means the key is just a raw hex bytes without encoding to a UTF-8 string, like:
+//     []byte{72, 101, 108, 108, 111, 32, 87, 111, 114, 108, 100}, which will later be converted to "48656c6c6f20576f726c64"
+//     by using `hex.EncodeToString()` method.
+func NewKeyRange(startKey, endKey []byte) *KeyRange {
+	return &KeyRange{startKey, endKey}
+}
+
+// EscapeAsUTF8Str returns the URL escaped key strings as they are UTF-8 encoded.
+func (r *KeyRange) EscapeAsUTF8Str() (startKeyStr, endKeyStr string) {
+	startKeyStr = url.QueryEscape(string(r.startKey))
+	endKeyStr = url.QueryEscape(string(r.endKey))
+	return
+}
+
+// EscapeAsHexStr returns the URL escaped key strings as they are hex encoded.
+func (r *KeyRange) EscapeAsHexStr() (startKeyStr, endKeyStr string) {
+	startKeyStr = url.QueryEscape(hex.EncodeToString(r.startKey))
+	endKeyStr = url.QueryEscape(hex.EncodeToString(r.endKey))
+	return
 }
 
 // NOTICE: the structures below are copied from the PD API definitions.

tests/integrations/client/http_client_test.go

@@ -280,24 +280,16 @@ func (suite *httpClientTestSuite) TestAccelerateSchedule() {
 	}
 	suspectRegions := raftCluster.GetSuspectRegions()
 	re.Len(suspectRegions, 0)
-	err := suite.client.AccelerateSchedule(suite.ctx, &pd.KeyRange{
-		StartKey: []byte("a1"),
-		EndKey:   []byte("a2")})
+	err := suite.client.AccelerateSchedule(suite.ctx, pd.NewKeyRange([]byte("a1"), []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, []*pd.KeyRange{
-		{
-			StartKey: []byte("a1"),
-			EndKey:   []byte("a2"),
-		},
-		{
-			StartKey: []byte("a2"),
-			EndKey:   []byte("a3"),
-		},
+		pd.NewKeyRange([]byte("a1"), []byte("a2")),
+		pd.NewKeyRange([]byte("a2"), []byte("a3")),
 	})
 	re.NoError(err)
 	suspectRegions = raftCluster.GetSuspectRegions()