Add keyspace meta in codec (#1323)

Signed-off-by: y_static_y@sina.com <y_static_y@sina.com>

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/genproto/googleapis/rpc v0.0.0-20240401170217-c3f982113cda // indirect
-	google.golang.org/grpc v1.62.1 // 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
 )

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/genproto/googleapis/rpc v0.0.0-20240401170217-c3f982113cda // indirect
-	google.golang.org/grpc v1.62.1 // 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
 )

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/genproto/googleapis/rpc v0.0.0-20240401170217-c3f982113cda // indirect
-	google.golang.org/grpc v1.62.1 // 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
 )

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/genproto/googleapis/rpc v0.0.0-20240401170217-c3f982113cda // indirect
-	google.golang.org/grpc v1.62.1 // 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
 )

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/genproto/googleapis/rpc v0.0.0-20240401170217-c3f982113cda // indirect
-	google.golang.org/grpc v1.62.1 // 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
 )

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/genproto/googleapis/rpc v0.0.0-20240401170217-c3f982113cda // indirect
-	google.golang.org/grpc v1.62.1 // 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
 )

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/genproto/googleapis/rpc v0.0.0-20240401170217-c3f982113cda // indirect
-	google.golang.org/grpc v1.62.1 // 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
 )

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/genproto/googleapis/rpc v0.0.0-20240401170217-c3f982113cda // indirect
-	google.golang.org/grpc v1.62.1 // 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
 )

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)

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

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
-	prefix     []byte
+	endKey       []byte
-	endKey     []byte
+	memCodec     memCodec
-	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 {

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,
+			keyspaceID: math.MaxUint32,
-			shouldErr: true,
+			shouldErr:  true,
 		},
 		{
 			// Bad mode should result in error.
-			mode:      Mode(99),
+			mode:       Mode(99),
-			spaceID:   DefaultKeyspaceID,
+			keyspaceID: DefaultKeyspaceID,
-			shouldErr: true,
+			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)
 	}

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