tikv
/
client-go
mirror of https://github.com/tikv/client-go.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

refactor decoding buckets with new interface method: DecodeBucketKeys (#760) 

Signed-off-by: iosmanthus <myosmanthustree@gmail.com>
This commit is contained in:
iosmanthus 2023-03-31 16:06:44 +08:00 committed by GitHub
parent 6a92aeec92
commit 4ac424aea1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 103 additions and 53 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
internal/apicodec/codec.go
View File

@ -57,8 +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)
// DecodeBucketKeys decode region bucket's key
DecodeBucketKeys(keys [][]byte) ([][]byte, error)
// EncodeRegionRange encode region's start and end.
EncodeRegionRange(start, end []byte) ([]byte, []byte)
// DecodeRegionRange decode region's start and end.

13
internal/apicodec/codec_v1.go
View File

@ -1,6 +1,7 @@
package apicodec
import (
"github.com/pingcap/errors"
"github.com/pingcap/kvproto/pkg/errorpb"
"github.com/pingcap/kvproto/pkg/kvrpcpb"
"github.com/tikv/client-go/v2/tikvrpc"
@ -203,6 +204,14 @@ func (c *codecV1) DecodeKey(key []byte) ([]byte, error) {
return key, nil
}
func (c *codecV1) DecodeBucketKey(key []byte) ([]byte, error) {
return c.DecodeRegionKey(key)
func (c *codecV1) DecodeBucketKeys(keys [][]byte) ([][]byte, error) {
ks := make([][]byte, 0, len(keys))
for _, key := range keys {
k, err := c.DecodeRegionKey(key)
if err != nil {
return nil, errors.Trace(err)
}
ks = append(ks, k)
}
return ks, nil
}

15
internal/apicodec/codec_v1_test.go
View File

@ -2,22 +2,7 @@ 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)
}

28
internal/apicodec/codec_v2.go
View File

@ -954,15 +954,31 @@ 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)
func (c *codecV2) DecodeBucketKeys(keys [][]byte) ([][]byte, error) {
ks := make([][]byte, 0, len(keys))
for i, key := range keys {
var (
k []byte
err error
)
if len(key) > 0 {
k, err = c.memCodec.decodeKey(key)
}
if err != nil {
return nil, err
}
if bytes.Compare(key, c.endKey) >= 0 || bytes.Compare(key, c.prefix) < 0 {
return []byte{}, nil
if i == 0 && bytes.Compare(k, c.prefix) < 0 {
ks = append(ks, []byte{})
} else if i == len(keys)-1 && (len(k) == 0 || bytes.Compare(k, c.endKey) >= 0) {
ks = append(ks, []byte{})
} else if bytes.HasPrefix(k, c.prefix) {
raw := k[len(c.prefix):]
if len(raw) == 0 && len(ks) > 0 && len(ks[0]) == 0 {
continue
}
return key[len(c.prefix):], nil
ks = append(ks, raw)
}
}
return ks, nil
}

85
internal/apicodec/codec_v2_test.go
View File

@ -296,22 +296,71 @@ func (suite *testCodecV2Suite) TestEncodeMPPRequest() {
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)
func (suite *testCodecV2Suite) TestDecodeBucketKeys() {
encodeWithPrefix := func(prefix, key []byte) []byte {
return codec.EncodeBytes(nil, append(prefix, key...))
}
bucketKeys := [][]byte{
encodeWithPrefix(prevKeyspacePrefix, []byte("a")),
encodeWithPrefix(prevKeyspacePrefix, []byte("b")),
encodeWithPrefix(prevKeyspacePrefix, []byte("c")),
suite.codec.EncodeRegionKey([]byte("a")),
suite.codec.EncodeRegionKey([]byte("b")),
suite.codec.EncodeRegionKey([]byte("c")),
encodeWithPrefix(keyspaceEndKey, []byte("")),
encodeWithPrefix(keyspaceEndKey, []byte("a")),
encodeWithPrefix(keyspaceEndKey, []byte("b")),
encodeWithPrefix(keyspaceEndKey, []byte("c")),
}
keys, err := suite.codec.DecodeBucketKeys(bucketKeys)
suite.Nil(err)
suite.Equal([][]byte{
{},
[]byte("a"),
[]byte("b"),
[]byte("c"),
{},
}, keys)
bucketKeys = [][]byte{
encodeWithPrefix(prevKeyspacePrefix, []byte("a")),
encodeWithPrefix(prevKeyspacePrefix, []byte("b")),
encodeWithPrefix(prevKeyspacePrefix, []byte("c")),
suite.codec.EncodeRegionKey([]byte("")),
suite.codec.EncodeRegionKey([]byte("a")),
suite.codec.EncodeRegionKey([]byte("b")),
suite.codec.EncodeRegionKey([]byte("c")),
{},
}
keys, err = suite.codec.DecodeBucketKeys(bucketKeys)
suite.Nil(err)
suite.Equal([][]byte{
{},
[]byte("a"),
[]byte("b"),
[]byte("c"),
{},
}, keys)
bucketKeys = [][]byte{
{},
encodeWithPrefix(prevKeyspacePrefix, []byte("a")),
encodeWithPrefix(prevKeyspacePrefix, []byte("b")),
encodeWithPrefix(prevKeyspacePrefix, []byte("c")),
suite.codec.EncodeRegionKey([]byte("")),
suite.codec.EncodeRegionKey([]byte("a")),
suite.codec.EncodeRegionKey([]byte("b")),
suite.codec.EncodeRegionKey([]byte("c")),
{},
}
keys, err = suite.codec.DecodeBucketKeys(bucketKeys)
suite.Nil(err)
suite.Equal([][]byte{
{},
[]byte("a"),
[]byte("b"),
[]byte("c"),
{},
}, keys)
}

13
internal/locate/pd_codec.go
View File

@ -163,16 +163,7 @@ func (c *CodecPDClient) decodeRegionKeyInPlace(r *pd.Region) error {
r.Meta.StartKey = decodedStart
r.Meta.EndKey = decodedEnd
if r.Buckets != nil {
for i, k := range r.Buckets.Keys {
if len(k) == 0 {
continue
r.Buckets.Keys, err = c.codec.DecodeBucketKeys(r.Buckets.Keys)
}
decoded, err := c.codec.DecodeBucketKey(k)
if err != nil {
return errors.WithStack(err)
}
r.Buckets.Keys[i] = decoded
}
}
return nil
return err
}