Skip to content

Commit

Permalink
fix api v2 decode of region bucket's keys (#757)
Browse files Browse the repository at this point in the history
Signed-off-by: iosmanthus <myosmanthustree@gmail.com>
  • Loading branch information
iosmanthus authored Mar 31, 2023
1 parent fd1b055 commit 6a92aee
Show file tree
Hide file tree
Showing 6 changed files with 71 additions and 7 deletions.
2 changes: 2 additions & 0 deletions internal/apicodec/codec.go
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,8 @@ type Codec interface {
EncodeRegionKey(key []byte) []byte
// DecodeRegionKey decode region's key
DecodeRegionKey(encodedKey []byte) ([]byte, error)
// DecodeBucketKey decode region bucket's key
DecodeBucketKey(encodedKey []byte) ([]byte, error)
// EncodeRegionRange encode region's start and end.
EncodeRegionRange(start, end []byte) ([]byte, []byte)
// DecodeRegionRange decode region's start and end.
Expand Down
4 changes: 4 additions & 0 deletions internal/apicodec/codec_v1.go
Original file line number Diff line number Diff line change
Expand Up @@ -202,3 +202,7 @@ func (c *codecV1) DecodeRange(start, end []byte) ([]byte, []byte, error) {
func (c *codecV1) DecodeKey(key []byte) ([]byte, error) {
return key, nil
}

func (c *codecV1) DecodeBucketKey(key []byte) ([]byte, error) {
return c.DecodeRegionKey(key)
}
23 changes: 23 additions & 0 deletions internal/apicodec/codec_v1_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
package apicodec

import (
"testing"

"github.com/stretchr/testify/assert"
"github.com/tikv/client-go/v2/util/codec"
)

func TestV1DecodeBucketKey(t *testing.T) {
c := NewCodecV1(ModeTxn)
raw := []byte("test")
encoded := codec.EncodeBytes(nil, raw)
key, err := c.DecodeBucketKey(encoded)
assert.Nil(t, err)
assert.Equal(t, raw, key)

raw = []byte{}
encoded = codec.EncodeBytes(nil, raw)
key, err = c.DecodeBucketKey(encoded)
assert.Nil(t, err)
assert.Equal(t, raw, key)
}
13 changes: 13 additions & 0 deletions internal/apicodec/codec_v2.go
Original file line number Diff line number Diff line change
Expand Up @@ -953,3 +953,16 @@ func (c *codecV2) decodeLockInfos(locks []*kvrpcpb.LockInfo) ([]*kvrpcpb.LockInf
}
return locks, nil
}

func (c *codecV2) DecodeBucketKey(encodedKey []byte) ([]byte, error) {
key, err := c.memCodec.decodeKey(encodedKey)
if err != nil {
return nil, err
}

if bytes.Compare(key, c.endKey) >= 0 || bytes.Compare(key, c.prefix) < 0 {
return []byte{}, nil
}

return key[len(c.prefix):], nil
}
34 changes: 28 additions & 6 deletions internal/apicodec/codec_v2_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,19 +11,21 @@ import (
"github.com/pingcap/kvproto/pkg/mpp"
"github.com/stretchr/testify/suite"
"github.com/tikv/client-go/v2/tikvrpc"
"github.com/tikv/client-go/v2/util/codec"
)

var (
testKeyspaceID = uint32(4242)
// Keys below are ordered as following:
// beforePrefix, keyspacePrefix, insideLeft, insideRight, keyspaceEndKey, afterEndKey
// where valid keyspace range is [keyspacePrefix, keyspaceEndKey)
keyspacePrefix = []byte{'r', 0, 16, 146}
keyspaceEndKey = []byte{'r', 0, 16, 147}
beforePrefix = []byte{'r', 0, 0, 1}
afterEndKey = []byte{'r', 1, 0, 0}
insideLeft = []byte{'r', 0, 16, 146, 100}
insideRight = []byte{'r', 0, 16, 146, 200}
prevKeyspacePrefix = []byte{'r', 0, 16, 145}
keyspacePrefix = []byte{'r', 0, 16, 146}
keyspaceEndKey = []byte{'r', 0, 16, 147}
beforePrefix = []byte{'r', 0, 0, 1}
afterEndKey = []byte{'r', 1, 0, 0}
insideLeft = []byte{'r', 0, 16, 146, 100}
insideRight = []byte{'r', 0, 16, 146, 200}
)

type testCodecV2Suite struct {
Expand Down Expand Up @@ -293,3 +295,23 @@ func (suite *testCodecV2Suite) TestEncodeMPPRequest() {
suite.Equal(task.Regions[0].Ranges[0].Start, suite.codec.EncodeKey([]byte("a")))
suite.Equal(task.Regions[0].Ranges[0].End, suite.codec.EncodeKey([]byte("b")))
}

func (suite *testCodecV2Suite) TestDecodeBucketKey() {
raw := []byte("a")
key := suite.codec.EncodeRegionKey(raw)
bucketKey, err := suite.codec.DecodeBucketKey(key)
suite.Nil(err)
suite.Equal(raw, bucketKey)

raw = keyspaceEndKey
key = codec.EncodeBytes([]byte{}, raw)
bucketKey, err = suite.codec.DecodeBucketKey(key)
suite.Nil(err)
suite.Empty(bucketKey)

raw = append(prevKeyspacePrefix, []byte("a")...)
key = codec.EncodeBytes([]byte{}, raw)
bucketKey, err = suite.codec.DecodeBucketKey(key)
suite.Nil(err)
suite.Empty(bucketKey)
}
2 changes: 1 addition & 1 deletion internal/locate/pd_codec.go
Original file line number Diff line number Diff line change
Expand Up @@ -167,7 +167,7 @@ func (c *CodecPDClient) decodeRegionKeyInPlace(r *pd.Region) error {
if len(k) == 0 {
continue
}
decoded, err := c.codec.DecodeRegionKey(k)
decoded, err := c.codec.DecodeBucketKey(k)
if err != nil {
return errors.WithStack(err)
}
Expand Down

0 comments on commit 6a92aee

Please sign in to comment.