Add keyspace meta in codec (#1323)

Signed-off-by: y_static_y@sina.com <y_static_y@sina.com>
This commit is contained in:
ystaticy 2024-04-28 12:20:55 +08:00 committed by GitHub
parent 3ed6550f3e
commit 9c05078660
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
13 changed files with 85 additions and 55 deletions

View File

@ -39,13 +39,13 @@ require (
go.uber.org/multierr v1.11.0 // indirect
go.uber.org/zap v1.26.0 // indirect
golang.org/x/exp v0.0.0-20230711005742-c3f37128e5a4 // indirect
golang.org/x/net v0.23.0 // indirect
golang.org/x/net v0.24.0 // indirect
golang.org/x/sync v0.6.0 // indirect
golang.org/x/sys v0.18.0 // indirect
golang.org/x/sys v0.19.0 // indirect
golang.org/x/text v0.14.0 // indirect
google.golang.org/genproto/googleapis/api v0.0.0-20240304212257-790db918fca8 // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20240308144416-29370a3891b7 // indirect
google.golang.org/grpc v1.62.1 // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20240401170217-c3f982113cda // indirect
google.golang.org/grpc v1.63.2 // indirect
google.golang.org/protobuf v1.33.0 // indirect
gopkg.in/natefinch/lumberjack.v2 v2.2.1 // indirect
)

View File

@ -39,13 +39,13 @@ require (
go.uber.org/multierr v1.11.0 // indirect
go.uber.org/zap v1.26.0 // indirect
golang.org/x/exp v0.0.0-20230711005742-c3f37128e5a4 // indirect
golang.org/x/net v0.23.0 // indirect
golang.org/x/net v0.24.0 // indirect
golang.org/x/sync v0.6.0 // indirect
golang.org/x/sys v0.18.0 // indirect
golang.org/x/sys v0.19.0 // indirect
golang.org/x/text v0.14.0 // indirect
google.golang.org/genproto/googleapis/api v0.0.0-20240304212257-790db918fca8 // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20240308144416-29370a3891b7 // indirect
google.golang.org/grpc v1.62.1 // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20240401170217-c3f982113cda // indirect
google.golang.org/grpc v1.63.2 // indirect
google.golang.org/protobuf v1.33.0 // indirect
gopkg.in/natefinch/lumberjack.v2 v2.2.1 // indirect
)

View File

@ -39,13 +39,13 @@ require (
go.uber.org/multierr v1.11.0 // indirect
go.uber.org/zap v1.26.0 // indirect
golang.org/x/exp v0.0.0-20230711005742-c3f37128e5a4 // indirect
golang.org/x/net v0.23.0 // indirect
golang.org/x/net v0.24.0 // indirect
golang.org/x/sync v0.6.0 // indirect
golang.org/x/sys v0.18.0 // indirect
golang.org/x/sys v0.19.0 // indirect
golang.org/x/text v0.14.0 // indirect
google.golang.org/genproto/googleapis/api v0.0.0-20240304212257-790db918fca8 // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20240308144416-29370a3891b7 // indirect
google.golang.org/grpc v1.62.1 // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20240401170217-c3f982113cda // indirect
google.golang.org/grpc v1.63.2 // indirect
google.golang.org/protobuf v1.33.0 // indirect
gopkg.in/natefinch/lumberjack.v2 v2.2.1 // indirect
)

View File

@ -39,13 +39,13 @@ require (
go.uber.org/multierr v1.11.0 // indirect
go.uber.org/zap v1.26.0 // indirect
golang.org/x/exp v0.0.0-20230711005742-c3f37128e5a4 // indirect
golang.org/x/net v0.23.0 // indirect
golang.org/x/net v0.24.0 // indirect
golang.org/x/sync v0.6.0 // indirect
golang.org/x/sys v0.18.0 // indirect
golang.org/x/sys v0.19.0 // indirect
golang.org/x/text v0.14.0 // indirect
google.golang.org/genproto/googleapis/api v0.0.0-20240304212257-790db918fca8 // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20240308144416-29370a3891b7 // indirect
google.golang.org/grpc v1.62.1 // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20240401170217-c3f982113cda // indirect
google.golang.org/grpc v1.63.2 // indirect
google.golang.org/protobuf v1.33.0 // indirect
gopkg.in/natefinch/lumberjack.v2 v2.2.1 // indirect
)

View File

@ -39,13 +39,13 @@ require (
go.uber.org/multierr v1.11.0 // indirect
go.uber.org/zap v1.26.0 // indirect
golang.org/x/exp v0.0.0-20230711005742-c3f37128e5a4 // indirect
golang.org/x/net v0.23.0 // indirect
golang.org/x/net v0.24.0 // indirect
golang.org/x/sync v0.6.0 // indirect
golang.org/x/sys v0.18.0 // indirect
golang.org/x/sys v0.19.0 // indirect
golang.org/x/text v0.14.0 // indirect
google.golang.org/genproto/googleapis/api v0.0.0-20240304212257-790db918fca8 // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20240308144416-29370a3891b7 // indirect
google.golang.org/grpc v1.62.1 // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20240401170217-c3f982113cda // indirect
google.golang.org/grpc v1.63.2 // indirect
google.golang.org/protobuf v1.33.0 // indirect
gopkg.in/natefinch/lumberjack.v2 v2.2.1 // indirect
)

View File

@ -39,13 +39,13 @@ require (
go.uber.org/multierr v1.11.0 // indirect
go.uber.org/zap v1.26.0 // indirect
golang.org/x/exp v0.0.0-20230711005742-c3f37128e5a4 // indirect
golang.org/x/net v0.23.0 // indirect
golang.org/x/net v0.24.0 // indirect
golang.org/x/sync v0.6.0 // indirect
golang.org/x/sys v0.18.0 // indirect
golang.org/x/sys v0.19.0 // indirect
golang.org/x/text v0.14.0 // indirect
google.golang.org/genproto/googleapis/api v0.0.0-20240304212257-790db918fca8 // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20240308144416-29370a3891b7 // indirect
google.golang.org/grpc v1.62.1 // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20240401170217-c3f982113cda // indirect
google.golang.org/grpc v1.63.2 // indirect
google.golang.org/protobuf v1.33.0 // indirect
gopkg.in/natefinch/lumberjack.v2 v2.2.1 // indirect
)

View File

@ -39,13 +39,13 @@ require (
go.uber.org/multierr v1.11.0 // indirect
go.uber.org/zap v1.26.0 // indirect
golang.org/x/exp v0.0.0-20230711005742-c3f37128e5a4 // indirect
golang.org/x/net v0.23.0 // indirect
golang.org/x/net v0.24.0 // indirect
golang.org/x/sync v0.6.0 // indirect
golang.org/x/sys v0.18.0 // indirect
golang.org/x/sys v0.19.0 // indirect
golang.org/x/text v0.14.0 // indirect
google.golang.org/genproto/googleapis/api v0.0.0-20240304212257-790db918fca8 // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20240308144416-29370a3891b7 // indirect
google.golang.org/grpc v1.62.1 // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20240401170217-c3f982113cda // indirect
google.golang.org/grpc v1.63.2 // indirect
google.golang.org/protobuf v1.33.0 // indirect
gopkg.in/natefinch/lumberjack.v2 v2.2.1 // indirect
)

View File

@ -39,13 +39,13 @@ require (
go.uber.org/multierr v1.11.0 // indirect
go.uber.org/zap v1.26.0 // indirect
golang.org/x/exp v0.0.0-20230711005742-c3f37128e5a4 // indirect
golang.org/x/net v0.23.0 // indirect
golang.org/x/net v0.24.0 // indirect
golang.org/x/sync v0.6.0 // indirect
golang.org/x/sys v0.18.0 // indirect
golang.org/x/sys v0.19.0 // indirect
golang.org/x/text v0.14.0 // indirect
google.golang.org/genproto/googleapis/api v0.0.0-20240304212257-790db918fca8 // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20240308144416-29370a3891b7 // indirect
google.golang.org/grpc v1.62.1 // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20240401170217-c3f982113cda // indirect
google.golang.org/grpc v1.63.2 // indirect
google.golang.org/protobuf v1.33.0 // indirect
gopkg.in/natefinch/lumberjack.v2 v2.2.1 // indirect
)

View File

@ -4,6 +4,7 @@ import (
"encoding/binary"
"github.com/pingcap/errors"
"github.com/pingcap/kvproto/pkg/keyspacepb"
"github.com/pingcap/kvproto/pkg/kvrpcpb"
"github.com/tikv/client-go/v2/tikvrpc"
)
@ -48,6 +49,8 @@ type Codec interface {
GetKeyspace() []byte
// GetKeyspaceID return the keyspace id of the codec.
GetKeyspaceID() KeyspaceID
// GetKeyspaceMeta return the keyspace meta of the codec.
GetKeyspaceMeta() *keyspacepb.KeyspaceMeta
// EncodeRequest encodes with the given Codec.
// NOTE: req is reused on retry. MUST encode on cloned request, other than overwrite the original.
EncodeRequest(req *tikvrpc.Request) (*tikvrpc.Request, error)

View File

@ -3,6 +3,7 @@ package apicodec
import (
"github.com/pingcap/errors"
"github.com/pingcap/kvproto/pkg/errorpb"
"github.com/pingcap/kvproto/pkg/keyspacepb"
"github.com/pingcap/kvproto/pkg/kvrpcpb"
"github.com/tikv/client-go/v2/tikvrpc"
)
@ -215,3 +216,7 @@ func (c *codecV1) DecodeBucketKeys(keys [][]byte) ([][]byte, error) {
}
return ks, nil
}
func (c *codecV1) GetKeyspaceMeta() *keyspacepb.KeyspaceMeta {
return nil
}

View File

@ -7,6 +7,7 @@ import (
"github.com/pingcap/kvproto/pkg/coprocessor"
"github.com/pingcap/kvproto/pkg/errorpb"
"github.com/pingcap/kvproto/pkg/keyspacepb"
"github.com/pingcap/kvproto/pkg/kvrpcpb"
"github.com/pingcap/kvproto/pkg/metapb"
"github.com/pkg/errors"
@ -49,15 +50,16 @@ func BuildKeyspaceName(name string) string {
// codecV2 is used to encode/decode keys and request into APIv2 format.
type codecV2 struct {
keyspaceID KeyspaceID
prefix []byte
endKey []byte
memCodec memCodec
prefix []byte
endKey []byte
memCodec memCodec
keyspaceMeta *keyspacepb.KeyspaceMeta
}
// NewCodecV2 returns a codec that can be used to encode/decode
// keys and requests to and from APIv2 format.
func NewCodecV2(mode Mode, keyspaceID uint32) (Codec, error) {
func NewCodecV2(mode Mode, keyspaceMeta *keyspacepb.KeyspaceMeta) (Codec, error) {
keyspaceID := keyspaceMeta.Id
if keyspaceID > maxKeyspaceID {
return nil, errors.Errorf("keyspaceID %d is out of range, maximum is %d", keyspaceID, maxKeyspaceID)
}
@ -66,9 +68,9 @@ func NewCodecV2(mode Mode, keyspaceID uint32) (Codec, error) {
return nil, err
}
codec := &codecV2{
keyspaceID: KeyspaceID(keyspaceID),
// Region keys in CodecV2 are always encoded in memory comparable form.
memCodec: &memComparableCodec{},
memCodec: &memComparableCodec{},
keyspaceMeta: keyspaceMeta,
}
codec.prefix = make([]byte, 4)
codec.endKey = make([]byte, 4)
@ -106,7 +108,11 @@ func (c *codecV2) GetKeyspace() []byte {
}
func (c *codecV2) GetKeyspaceID() KeyspaceID {
return c.keyspaceID
return KeyspaceID(c.keyspaceMeta.Id)
}
func (c *codecV2) GetKeyspaceMeta() *keyspacepb.KeyspaceMeta {
return c.keyspaceMeta
}
func (c *codecV2) GetAPIVersion() kvrpcpb.APIVersion {

View File

@ -6,6 +6,7 @@ import (
"github.com/pingcap/kvproto/pkg/coprocessor"
"github.com/pingcap/kvproto/pkg/errorpb"
"github.com/pingcap/kvproto/pkg/keyspacepb"
"github.com/pingcap/kvproto/pkg/kvrpcpb"
"github.com/pingcap/kvproto/pkg/metapb"
"github.com/pingcap/kvproto/pkg/mpp"
@ -38,7 +39,10 @@ func TestCodecV2(t *testing.T) {
}
func (suite *testCodecV2Suite) SetupSuite() {
codec, err := NewCodecV2(ModeRaw, testKeyspaceID)
testKeyspaceMeta := keyspacepb.KeyspaceMeta{
Id: testKeyspaceID,
}
codec, err := NewCodecV2(ModeRaw, &testKeyspaceMeta)
suite.NoError(err)
suite.Equal(keyspacePrefix, codec.GetKeyspace())
suite.codec = codec.(*codecV2)
@ -110,7 +114,7 @@ func (suite *testCodecV2Suite) TestNewCodecV2() {
re := suite.Require()
testCases := []struct {
mode Mode
spaceID uint32
keyspaceID uint32
shouldErr bool
expectedPrefix []byte
expectedEnd []byte
@ -118,54 +122,57 @@ func (suite *testCodecV2Suite) TestNewCodecV2() {
{
mode: ModeRaw,
// A too large keyspaceID should result in error.
spaceID: math.MaxUint32,
shouldErr: true,
keyspaceID: math.MaxUint32,
shouldErr: true,
},
{
// Bad mode should result in error.
mode: Mode(99),
spaceID: DefaultKeyspaceID,
shouldErr: true,
mode: Mode(99),
keyspaceID: DefaultKeyspaceID,
shouldErr: true,
},
{
mode: ModeRaw,
spaceID: 1<<24 - 2,
keyspaceID: 1<<24 - 2,
expectedPrefix: []byte{'r', 255, 255, 254},
expectedEnd: []byte{'r', 255, 255, 255},
},
{
// EndKey should be able to carry over increment from lower byte.
mode: ModeTxn,
spaceID: 1<<8 - 1,
keyspaceID: 1<<8 - 1,
expectedPrefix: []byte{'x', 0, 0, 255},
expectedEnd: []byte{'x', 0, 1, 0},
},
{
// EndKey should be able to carry over increment from lower byte.
mode: ModeTxn,
spaceID: 1<<16 - 1,
keyspaceID: 1<<16 - 1,
expectedPrefix: []byte{'x', 0, 255, 255},
expectedEnd: []byte{'x', 1, 0, 0},
},
{
// If prefix is the last keyspace, then end should change the mode byte.
mode: ModeRaw,
spaceID: 1<<24 - 1,
keyspaceID: 1<<24 - 1,
expectedPrefix: []byte{'r', 255, 255, 255},
expectedEnd: []byte{'s', 0, 0, 0},
},
}
for _, testCase := range testCases {
keyspaceMeta := &keyspacepb.KeyspaceMeta{Id: testCase.keyspaceID}
if testCase.shouldErr {
_, err := NewCodecV2(testCase.mode, testCase.spaceID)
_, err := NewCodecV2(testCase.mode, keyspaceMeta)
re.Error(err)
continue
}
codec, err := NewCodecV2(testCase.mode, testCase.spaceID)
codec, err := NewCodecV2(testCase.mode, keyspaceMeta)
re.NoError(err)
v2Codec, ok := codec.(*codecV2)
re.True(ok)
re.Equal(keyspaceMeta, v2Codec.keyspaceMeta)
re.Equal(testCase.expectedPrefix, v2Codec.prefix)
re.Equal(testCase.expectedEnd, v2Codec.endKey)
}

View File

@ -60,11 +60,11 @@ func NewCodecPDClient(mode apicodec.Mode, client pd.Client) *CodecPDClient {
// NewCodecPDClientWithKeyspace creates a CodecPDClient in API v2 with keyspace name.
func NewCodecPDClientWithKeyspace(mode apicodec.Mode, client pd.Client, keyspace string) (*CodecPDClient, error) {
id, err := GetKeyspaceID(client, keyspace)
keyspaceMeta, err := GetKeyspaceMeta(client, keyspace)
if err != nil {
return nil, err
}
codec, err := apicodec.NewCodecV2(mode, id)
codec, err := apicodec.NewCodecV2(mode, keyspaceMeta)
if err != nil {
return nil, err
}
@ -85,6 +85,15 @@ func GetKeyspaceID(client pd.Client, name string) (uint32, error) {
return meta.Id, nil
}
// GetKeyspaceMeta attempts to retrieve keyspace meta corresponding to the given keyspace name from PD.
func GetKeyspaceMeta(client pd.Client, name string) (*keyspacepb.KeyspaceMeta, error) {
meta, err := client.LoadKeyspace(context.Background(), apicodec.BuildKeyspaceName(name))
if err != nil {
return nil, err
}
return meta, nil
}
// GetCodec returns CodecPDClient's codec.
func (c *CodecPDClient) GetCodec() apicodec.Codec {
return c.codec