refactor decoding buckets with new interface method: DecodeBucketKeys

internal/apicodec/codec.go

@@ -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.

internal/apicodec/codec_v1.go

@@ -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
 }

internal/apicodec/codec_v1_test.go

@@ -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)
 }

internal/apicodec/codec_v2.go

@@ -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)
-	if err != nil {
-		return nil, err
-	}
+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
+			}
+			ks = append(ks, raw)
+		}
 	}
-
-	return key[len(c.prefix):], nil
+	return ks, nil
 }

internal/apicodec/codec_v2_test.go

@@ -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)
+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(raw, bucketKey)
+	suite.Equal([][]byte{
+		{},
+		[]byte("a"),
+		[]byte("b"),
+		[]byte("c"),
+		{},
+	}, keys)
 
-	raw = keyspaceEndKey
-	key = codec.EncodeBytes([]byte{}, raw)
-	bucketKey, err = suite.codec.DecodeBucketKey(key)
+	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.Empty(bucketKey)
+	suite.Equal([][]byte{
+		{},
+		[]byte("a"),
+		[]byte("b"),
+		[]byte("c"),
+		{},
+	}, keys)
 
-	raw = append(prevKeyspacePrefix, []byte("a")...)
-	key = codec.EncodeBytes([]byte{}, raw)
-	bucketKey, err = suite.codec.DecodeBucketKey(key)
+	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.Empty(bucketKey)
+	suite.Equal([][]byte{
+		{},
+		[]byte("a"),
+		[]byte("b"),
+		[]byte("c"),
+		{},
+	}, keys)
 }

internal/locate/pd_codec.go

@@ -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
-			}
-			decoded, err := c.codec.DecodeBucketKey(k)
-			if err != nil {
-				return errors.WithStack(err)
-			}
-			r.Buckets.Keys[i] = decoded
-		}
+		r.Buckets.Keys, err = c.codec.DecodeBucketKeys(r.Buckets.Keys)
 	}
-	return nil
+	return err
 }