From 9c05078660d9e5c4cb4d2dc748e1fe4d99cf3c5a Mon Sep 17 00:00:00 2001 From: ystaticy Date: Sun, 28 Apr 2024 12:20:55 +0800 Subject: [PATCH] Add keyspace meta in codec (#1323) Signed-off-by: y_static_y@sina.com --- examples/gcworker/go.mod | 8 +++--- examples/rawkv/go.mod | 8 +++--- examples/txnkv/1pc_txn/go.mod | 8 +++--- examples/txnkv/async_commit/go.mod | 8 +++--- examples/txnkv/delete_range/go.mod | 8 +++--- examples/txnkv/go.mod | 8 +++--- examples/txnkv/pessimistic_txn/go.mod | 8 +++--- examples/txnkv/unsafedestoryrange/go.mod | 8 +++--- internal/apicodec/codec.go | 3 +++ internal/apicodec/codec_v1.go | 5 ++++ internal/apicodec/codec_v2.go | 22 ++++++++++------ internal/apicodec/codec_v2_test.go | 33 ++++++++++++++---------- internal/locate/pd_codec.go | 13 ++++++++-- 13 files changed, 85 insertions(+), 55 deletions(-) diff --git a/examples/gcworker/go.mod b/examples/gcworker/go.mod index 59e984ce..c735b5df 100644 --- a/examples/gcworker/go.mod +++ b/examples/gcworker/go.mod @@ -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 ) diff --git a/examples/rawkv/go.mod b/examples/rawkv/go.mod index b973b799..ce612172 100644 --- a/examples/rawkv/go.mod +++ b/examples/rawkv/go.mod @@ -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 ) diff --git a/examples/txnkv/1pc_txn/go.mod b/examples/txnkv/1pc_txn/go.mod index ba52b506..dae17ded 100644 --- a/examples/txnkv/1pc_txn/go.mod +++ b/examples/txnkv/1pc_txn/go.mod @@ -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 ) diff --git a/examples/txnkv/async_commit/go.mod b/examples/txnkv/async_commit/go.mod index db10d224..12714073 100644 --- a/examples/txnkv/async_commit/go.mod +++ b/examples/txnkv/async_commit/go.mod @@ -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 ) diff --git a/examples/txnkv/delete_range/go.mod b/examples/txnkv/delete_range/go.mod index 6142997c..44cd1df2 100644 --- a/examples/txnkv/delete_range/go.mod +++ b/examples/txnkv/delete_range/go.mod @@ -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 ) diff --git a/examples/txnkv/go.mod b/examples/txnkv/go.mod index bbfb25ea..9cd62c7f 100644 --- a/examples/txnkv/go.mod +++ b/examples/txnkv/go.mod @@ -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 ) diff --git a/examples/txnkv/pessimistic_txn/go.mod b/examples/txnkv/pessimistic_txn/go.mod index 4b9089c0..f38c9160 100644 --- a/examples/txnkv/pessimistic_txn/go.mod +++ b/examples/txnkv/pessimistic_txn/go.mod @@ -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 ) diff --git a/examples/txnkv/unsafedestoryrange/go.mod b/examples/txnkv/unsafedestoryrange/go.mod index 8eeb1bc1..fe52ad51 100644 --- a/examples/txnkv/unsafedestoryrange/go.mod +++ b/examples/txnkv/unsafedestoryrange/go.mod @@ -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 ) diff --git a/internal/apicodec/codec.go b/internal/apicodec/codec.go index 54b5840c..7426b04a 100644 --- a/internal/apicodec/codec.go +++ b/internal/apicodec/codec.go @@ -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) diff --git a/internal/apicodec/codec_v1.go b/internal/apicodec/codec_v1.go index 74e32df9..03ebd1f2 100644 --- a/internal/apicodec/codec_v1.go +++ b/internal/apicodec/codec_v1.go @@ -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 +} diff --git a/internal/apicodec/codec_v2.go b/internal/apicodec/codec_v2.go index 25313d61..b23b1565 100644 --- a/internal/apicodec/codec_v2.go +++ b/internal/apicodec/codec_v2.go @@ -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 { diff --git a/internal/apicodec/codec_v2_test.go b/internal/apicodec/codec_v2_test.go index 04180b80..67d9009d 100644 --- a/internal/apicodec/codec_v2_test.go +++ b/internal/apicodec/codec_v2_test.go @@ -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) } diff --git a/internal/locate/pd_codec.go b/internal/locate/pd_codec.go index c2d76873..88e187e0 100644 --- a/internal/locate/pd_codec.go +++ b/internal/locate/pd_codec.go @@ -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