From fdeda3f4c7ba011c1a38001e8441ab3ee21129f8 Mon Sep 17 00:00:00 2001 From: JmPotato Date: Mon, 27 Nov 2023 14:53:50 +0800 Subject: [PATCH] Use KeyRange to differentiate the encoding format of the key Signed-off-by: JmPotato --- client/http/api.go | 17 ++++------ client/http/client.go | 20 ++++++----- client/http/types.go | 34 +++++++++++++++++-- tests/integrations/client/http_client_test.go | 14 ++------ 4 files changed, 53 insertions(+), 32 deletions(-) diff --git a/client/http/api.go b/client/http/api.go index c6d4f2dfb74..6b317330b61 100644 --- a/client/http/api.go +++ b/client/http/api.go @@ -83,13 +83,11 @@ func RegionByKey(key []byte) string { return fmt.Sprintf("%s/%s", regionByKey, url.QueryEscape(string(key))) } -// 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 { +// RegionsByKeyRange returns the path of PD HTTP API to scan regions with given start key, end key and limit parameters. +func RegionsByKeyRange(keyRange *KeyRange, limit int) string { + startKeyStr, endKeyStr := keyRange.EscapeAsUTF8Str() return fmt.Sprintf("%s?start_key=%s&end_key=%s&limit=%d", - regionsByKey, - url.QueryEscape(string(startKey)), - url.QueryEscape(string(endKey)), - limit) + regionsByKey, startKeyStr, endKeyStr, limit) } // RegionsByStoreID returns the path of PD HTTP API to get regions by store ID. @@ -98,11 +96,10 @@ 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(keyRange *KeyRange) string { + startKeyStr, endKeyStr := keyRange.EscapeAsUTF8Str() return fmt.Sprintf("%s?start_key=%s&end_key=%s", - StatsRegion, - url.QueryEscape(string(startKey)), - url.QueryEscape(string(endKey))) + StatsRegion, startKeyStr, endKeyStr) } // StoreByID returns the store API with store ID parameter. diff --git a/client/http/client.go b/client/http/client.go index 1d8c2d5c427..cd7be205c12 100644 --- a/client/http/client.go +++ b/client/http/client.go @@ -18,12 +18,10 @@ import ( "bytes" "context" "crypto/tls" - "encoding/hex" "encoding/json" "fmt" "io" "net/http" - "net/url" "strings" "time" @@ -315,10 +313,11 @@ 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 err := c.requestWithRetry(ctx, - "GetRegionsByKeyRange", RegionsByKey(keyRange.StartKey, keyRange.EndKey, limit), + "GetRegionsByKeyRange", RegionsByKeyRange(keyRange, limit), http.MethodGet, http.NoBody, ®ions) if err != nil { return nil, err @@ -363,10 +362,11 @@ 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 err := c.requestWithRetry(ctx, - "GetRegionStatusByKeyRange", RegionStatsByKeyRange(keyRange.StartKey, keyRange.StartKey), + "GetRegionStatusByKeyRange", RegionStatsByKeyRange(keyRange), http.MethodGet, http.NoBody, ®ionStats, ) if err != nil { @@ -557,10 +557,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": startKey, + "end_key": endKey, }) if err != nil { return errors.Trace(err) @@ -571,12 +573,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": startKey, + "end_key": endKey, }) } inputJSON, err := json.Marshal(input) diff --git a/client/http/types.go b/client/http/types.go index 56d59bafa58..df448e7e20d 100644 --- a/client/http/types.go +++ b/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. diff --git a/tests/integrations/client/http_client_test.go b/tests/integrations/client/http_client_test.go index dc901b1d290..3a52e91e1f8 100644 --- a/tests/integrations/client/http_client_test.go +++ b/tests/integrations/client/http_client_test.go @@ -280,9 +280,7 @@ 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) @@ -290,14 +288,8 @@ func (suite *httpClientTestSuite) TestAccelerateSchedule() { 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()