Add API V3 keyspace identity to PD client

client/client.go

@@ -28,6 +28,7 @@ import (
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/failpoint"
+	"github.com/pingcap/kvproto/pkg/apipb"
 	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/pingcap/kvproto/pkg/pdpb"
 	"github.com/pingcap/kvproto/pkg/routerpb"
@@ -231,6 +232,20 @@ func NewClientWithKeyspace(
 		svrAddrs, security, opts...)
 }
 
+// NewClientWithKeyspaceIdentity creates a client with context and the specified API V3 keyspace identity.
+func NewClientWithKeyspaceIdentity(
+	ctx context.Context,
+	callerComponent caller.Component,
+	identity *apipb.KeyspaceIdentity, svrAddrs []string,
+	security SecurityOption, opts ...opt.ClientOption,
+) (Client, error) {
+	if err := validateKeyspaceIdentity(identity); err != nil {
+		return nil, err
+	}
+	return createClientWithKeyspaceIdentity(ctx, callerComponent, identity,
+		svrAddrs, security, opts...)
+}
+
 // createClientWithKeyspace creates a client with context and the specified keyspace id.
 func createClientWithKeyspace(
 	ctx context.Context,
@@ -272,6 +287,73 @@ func createClientWithKeyspace(
 	return c, c.inner.init(nil)
 }
 
+// createClientWithKeyspaceIdentity creates a client with context and the specified API V3 keyspace identity.
+func createClientWithKeyspaceIdentity(
+	ctx context.Context,
+	callerComponent caller.Component,
+	identity *apipb.KeyspaceIdentity, svrAddrs []string,
+	security SecurityOption, opts ...opt.ClientOption,
+) (Client, error) {
+	tlsCfg, err := tlsutil.TLSConfig{
+		CAPath:   security.CAPath,
+		CertPath: security.CertPath,
+		KeyPath:  security.KeyPath,
+
+		SSLCABytes:   security.SSLCABytes,
+		SSLCertBytes: security.SSLCertBytes,
+		SSLKEYBytes:  security.SSLKEYBytes,
+	}.ToTLSConfig()
+	if err != nil {
+		return nil, err
+	}
+	clientCtx, clientCancel := context.WithCancel(ctx)
+	c := &client{
+		callerComponent: adjustCallerComponent(callerComponent),
+		inner: &innerClient{
+			keyspaceID:              identity.GetKeyspaceId(),
+			keyspaceIdentity:        cloneKeyspaceIdentity(identity),
+			svrUrls:                 svrAddrs,
+			updateTokenConnectionCh: make(chan struct{}, 1),
+			ctx:                     clientCtx,
+			cancel:                  clientCancel,
+			tlsCfg:                  tlsCfg,
+			option:                  opt.NewOption(),
+		},
+	}
+
+	// Inject the client options.
+	for _, opt := range opts {
+		opt(c.inner.option)
+	}
+
+	return c, c.inner.init(nil)
+}
+
+func validateKeyspaceIdentity(identity *apipb.KeyspaceIdentity) error {
+	if identity == nil {
+		return errors.New("missing keyspace identity")
+	}
+	if identity.GetNamespaceId() == 0 {
+		return errors.New("keyspace identity namespace id must be non-zero")
+	}
+	keyspaceID := identity.GetKeyspaceId()
+	if keyspaceID == 0 || keyspaceID > constants.MaxKeyspaceID {
+		return errors.Errorf("invalid keyspace id %d. It must be in the range of [1, %d]",
+			keyspaceID, constants.MaxKeyspaceID)
+	}
+	return nil
+}
+
+func cloneKeyspaceIdentity(identity *apipb.KeyspaceIdentity) *apipb.KeyspaceIdentity {
+	if identity == nil {
+		return nil
+	}
+	return &apipb.KeyspaceIdentity{
+		NamespaceId: identity.GetNamespaceId(),
+		KeyspaceId:  identity.GetKeyspaceId(),
+	}
+}
+
 // APIVersion is the API version the server and the client is using.
 // See more details in https://github.com/tikv/rfcs/blob/master/text/0069-api-v2.md#kvproto
 type APIVersion int
@@ -282,6 +364,7 @@ const (
 	V1 APIVersion = iota
 	_
 	V2
+	V3
 )
 
 // APIContext is the context for API version.
@@ -290,6 +373,10 @@ type APIContext interface {
 	GetKeyspaceName() (keyspaceName string)
 }
 
+type keyspaceIdentityAPIContext interface {
+	GetKeyspaceIdentity() *apipb.KeyspaceIdentity
+}
+
 type apiContextV1 struct{}
 
 // NewAPIContextV1 creates a API context for V1.
@@ -329,6 +416,34 @@ func (apiCtx *apiContextV2) GetKeyspaceName() (keyspaceName string) {
 	return apiCtx.keyspaceName
 }
 
+type apiContextV3 struct {
+	keyspaceName     string
+	keyspaceIdentity *apipb.KeyspaceIdentity
+}
+
+// NewAPIContextV3 creates an API context with the specified keyspace identity for V3.
+func NewAPIContextV3(keyspaceName string, identity *apipb.KeyspaceIdentity) APIContext {
+	return &apiContextV3{
+		keyspaceName:     keyspaceName,
+		keyspaceIdentity: cloneKeyspaceIdentity(identity),
+	}
+}
+
+// GetAPIVersion returns the API version.
+func (*apiContextV3) GetAPIVersion() (version APIVersion) {
+	return V3
+}
+
+// GetKeyspaceName returns the keyspace name.
+func (apiCtx *apiContextV3) GetKeyspaceName() (keyspaceName string) {
+	return apiCtx.keyspaceName
+}
+
+// GetKeyspaceIdentity returns the keyspace identity.
+func (apiCtx *apiContextV3) GetKeyspaceIdentity() *apipb.KeyspaceIdentity {
+	return cloneKeyspaceIdentity(apiCtx.keyspaceIdentity)
+}
+
 // NewClientWithAPIContext creates a client according to the API context.
 func NewClientWithAPIContext(
 	ctx context.Context, apiCtx APIContext,
@@ -344,6 +459,13 @@ func NewClientWithAPIContext(
 	case V2:
 		return newClientWithKeyspaceName(ctx, callerComponent,
 			keyspaceName, svrAddrs, security, opts...)
+	case V3:
+		identityCtx, ok := apiCtx.(keyspaceIdentityAPIContext)
+		if !ok {
+			return nil, errors.New("[pd] API V3 context missing keyspace identity")
+		}
+		return NewClientWithKeyspaceIdentity(ctx, callerComponent,
+			identityCtx.GetKeyspaceIdentity(), svrAddrs, security, opts...)
 	default:
 		return nil, errors.Errorf("[pd] invalid API version %d", apiVersion)
 	}

client/clients/tso/dispatcher.go

@@ -431,6 +431,7 @@ func (td *tsoDispatcher) processRequests(
 		svcDiscovery       = td.provider.getServiceDiscovery()
 		clusterID          = svcDiscovery.GetClusterID()
 		keyspaceID         = svcDiscovery.GetKeyspaceID()
+		keyspaceIdentity   = svcDiscovery.GetKeyspaceIdentity()
 		reqKeyspaceGroupID = svcDiscovery.GetKeyspaceGroupID()
 	)
 
@@ -465,7 +466,7 @@ func (td *tsoDispatcher) processRequests(
 	}
 
 	err := stream.processRequests(
-		clusterID, keyspaceID, reqKeyspaceGroupID,
+		clusterID, keyspaceID, reqKeyspaceGroupID, keyspaceIdentity,
 		count, tbc.GetExtraBatchingStartTime(), cb)
 	if err != nil {
 		close(done)

client/clients/tso/dispatcher_test.go

@@ -29,6 +29,7 @@ import (
 	"google.golang.org/grpc"
 
 	"github.com/pingcap/failpoint"
+	"github.com/pingcap/kvproto/pkg/apipb"
 	"github.com/pingcap/log"
 
 	"github.com/tikv/pd/client/errs"
@@ -199,6 +200,8 @@ func (*countingServiceDiscovery) Close()
 func (*countingServiceDiscovery) GetClusterID() uint64                               { return 0 }
 func (*countingServiceDiscovery) GetKeyspaceID() uint32                              { return 0 }
 func (*countingServiceDiscovery) SetKeyspaceID(uint32)                               {}
+func (*countingServiceDiscovery) GetKeyspaceIdentity() *apipb.KeyspaceIdentity       { return nil }
+func (*countingServiceDiscovery) SetKeyspaceIdentity(*apipb.KeyspaceIdentity)        {}
 func (*countingServiceDiscovery) GetKeyspaceGroupID() uint32                         { return 0 }
 func (*countingServiceDiscovery) GetServiceURLs() []string                           { return nil }
 func (*countingServiceDiscovery) GetServingEndpointClientConn() *grpc.ClientConn     { return nil }

client/clients/tso/stream.go

@@ -29,6 +29,7 @@ import (
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/failpoint"
+	"github.com/pingcap/kvproto/pkg/apipb"
 	"github.com/pingcap/kvproto/pkg/pdpb"
 	"github.com/pingcap/kvproto/pkg/tsopb"
 	"github.com/pingcap/log"
@@ -124,7 +125,7 @@ type tsoRequestResult struct {
 }
 
 type grpcTSOStreamAdapter interface {
-	Send(clusterID uint64, keyspaceID, keyspaceGroupID uint32, count int64) error
+	Send(clusterID uint64, keyspaceID, keyspaceGroupID uint32, keyspaceIdentity *apipb.KeyspaceIdentity, count int64) error
 	Recv() (tsoRequestResult, error)
 }
 
@@ -133,12 +134,13 @@ type pdTSOStreamAdapter struct {
 }
 
 // Send implements the grpcTSOStreamAdapter interface.
-func (s pdTSOStreamAdapter) Send(clusterID uint64, _, _ uint32, count int64) error {
+func (s pdTSOStreamAdapter) Send(clusterID uint64, _, _ uint32, keyspaceIdentity *apipb.KeyspaceIdentity, count int64) error {
 	req := &pdpb.TsoRequest{
 		Header: &pdpb.RequestHeader{
 			ClusterId: clusterID,
 		},
-		Count: uint32(count),
+		Count:            uint32(count),
+		KeyspaceIdentity: cloneKeyspaceIdentity(keyspaceIdentity),
 	}
 	return s.stream.Send(req)
 }
@@ -163,15 +165,24 @@ type tsoTSOStreamAdapter struct {
 }
 
 // Send implements the grpcTSOStreamAdapter interface.
-func (s tsoTSOStreamAdapter) Send(clusterID uint64, keyspaceID, keyspaceGroupID uint32, count int64) error {
+func (s tsoTSOStreamAdapter) Send(clusterID uint64, keyspaceID, keyspaceGroupID uint32, keyspaceIdentity *apipb.KeyspaceIdentity, count int64) error {
+	header := &tsopb.RequestHeader{
+		ClusterId:       clusterID,
+		KeyspaceGroupId: keyspaceGroupID,
+		CalleeId:        s.calleeID,
+	}
+	if keyspaceIdentity != nil {
+		header.Keyspace = &tsopb.RequestHeader_KeyspaceIdentity{
+			KeyspaceIdentity: cloneKeyspaceIdentity(keyspaceIdentity),
+		}
+	} else {
+		header.Keyspace = &tsopb.RequestHeader_KeyspaceId{
+			KeyspaceId: keyspaceID,
+		}
+	}
 	req := &tsopb.TsoRequest{
-		Header: &tsopb.RequestHeader{
-			ClusterId:       clusterID,
-			KeyspaceId:      keyspaceID,
-			KeyspaceGroupId: keyspaceGroupID,
-			CalleeId:        s.calleeID,
-		},
-		Count: uint32(count),
+		Header: header,
+		Count:  uint32(count),
 	}
 	return s.stream.Send(req)
 }
@@ -274,7 +285,8 @@ func (s *tsoStream) getServerURL() string {
 // It's guaranteed that the `callback` will be called, but when the request is failed to be scheduled, the callback
 // will be ignored.
 func (s *tsoStream) processRequests(
-	clusterID uint64, keyspaceID, keyspaceGroupID uint32, count int64, batchStartTime time.Time, callback onFinishedCallback,
+	clusterID uint64, keyspaceID, keyspaceGroupID uint32, keyspaceIdentity *apipb.KeyspaceIdentity,
+	count int64, batchStartTime time.Time, callback onFinishedCallback,
 ) error {
 	start := time.Now()
 
@@ -312,7 +324,7 @@ func (s *tsoStream) processRequests(
 	s.state.Store(prevState)
 
 	failpoint.InjectCall("pauseAfterTSORequestAttachedToStream")
-	if err := s.stream.Send(clusterID, keyspaceID, keyspaceGroupID, count); err != nil {
+	if err := s.stream.Send(clusterID, keyspaceID, keyspaceGroupID, keyspaceIdentity, count); err != nil {
 		// As the request is already put into `pendingRequests`, the request should finally be canceled by the recvLoop.
 		// So skip returning error here to avoid
 		// if err == io.EOF {
@@ -327,6 +339,16 @@ func (s *tsoStream) processRequests(
 	return nil
 }
 
+func cloneKeyspaceIdentity(identity *apipb.KeyspaceIdentity) *apipb.KeyspaceIdentity {
+	if identity == nil {
+		return nil
+	}
+	return &apipb.KeyspaceIdentity{
+		NamespaceId: identity.GetNamespaceId(),
+		KeyspaceId:  identity.GetKeyspaceId(),
+	}
+}
+
 func (s *tsoStream) recvLoop(ctx context.Context) {
 	var finishWithErr error
 	var currentReq batchedRequests

client/clients/tso/stream_test.go

@@ -26,6 +26,7 @@ import (
 	"go.uber.org/zap/zapcore"
 
 	"github.com/pingcap/errors"
+	"github.com/pingcap/kvproto/pkg/apipb"
 	"github.com/pingcap/log"
 
 	"github.com/tikv/pd/client/errs"
@@ -79,7 +80,7 @@ func newMockTSOStreamImpl(ctx context.Context, resultMode resultMode) *mockTSOSt
 	}
 }
 
-func (s *mockTSOStreamImpl) Send(clusterID uint64, _keyspaceID, keyspaceGroupID uint32, count int64) error {
+func (s *mockTSOStreamImpl) Send(clusterID uint64, _keyspaceID, keyspaceGroupID uint32, _keyspaceIdentity *apipb.KeyspaceIdentity, count int64) error {
 	select {
 	case <-s.ctx.Done():
 		return s.ctx.Err()
@@ -306,7 +307,7 @@ func (s *testTSOStreamSuite) getResult(ch <-chan callbackInvocation) callbackInv
 
 func (s *testTSOStreamSuite) processRequestWithResultCh(count int64) (<-chan callbackInvocation, error) {
 	ch := make(chan callbackInvocation, 1)
-	err := s.stream.processRequests(1, 2, 3, count, time.Now(), func(result tsoRequestResult, reqKeyspaceGroupID uint32, err error) {
+	err := s.stream.processRequests(1, 2, 3, nil, count, time.Now(), func(result tsoRequestResult, reqKeyspaceGroupID uint32, err error) {
 		if err == nil {
 			s.re.Equal(uint32(3), reqKeyspaceGroupID)
 		}
@@ -358,7 +359,7 @@ func (s *testTSOStreamSuite) TestTSOStreamBasic() {
 	// After an error from the (simulated) RPC stream, the tsoStream should be in a broken status and can't accept
 	// new request anymore.
 	testutil.Eventually(s.re, func() bool {
-		return s.stream.processRequests(1, 2, 3, 1, time.Now(), func(_result tsoRequestResult, _reqKeyspaceGroupID uint32, err error) {
+		return s.stream.processRequests(1, 2, 3, nil, 1, time.Now(), func(_result tsoRequestResult, _reqKeyspaceGroupID uint32, err error) {
 			s.re.Error(err)
 		}) != nil
 	})
@@ -622,7 +623,7 @@ func BenchmarkTSOStreamSendRecv(b *testing.B) {
 
 	b.ResetTimer()
 	for range b.N {
-		err := stream.processRequests(1, 1, 1, 1, now, func(result tsoRequestResult, _ uint32, err error) {
+		err := stream.processRequests(1, 1, 1, nil, 1, now, func(result tsoRequestResult, _ uint32, err error) {
 			if err != nil {
 				panic(err)
 			}