tso: validate callee id for tso requests (tikv#439)

* tso: validate callee id for tso requests

Signed-off-by: iosmanthus <myosmanthustree@gmail.com>

* attach CalleeId to tsopb.RequestHeader

Signed-off-by: iosmanthus <myosmanthustree@gmail.com>

* update kvproto to tidbcloud/kvproto/release-8.5-keyspace

Signed-off-by: iosmanthus <myosmanthustree@gmail.com>

* compare callee id with url.Parse

Signed-off-by: iosmanthus <myosmanthustree@gmail.com>

* resolve comments from @rleungx

Signed-off-by: iosmanthus <myosmanthustree@gmail.com>

* make linter happy

Signed-off-by: iosmanthus <myosmanthustree@gmail.com>

---------

Signed-off-by: iosmanthus <myosmanthustree@gmail.com>

Author: iosmanthus

client/errs/errno.go

@@ -34,6 +34,8 @@ const (
 	RetryTimeoutErr = "retry timeout"
 	// NotPrimaryErr indicates the non-primary member received the requests which should be received by primary.
 	NotPrimaryErr = "not primary"
+	// MismatchCalleeIDErr indicates the callee ID is mismatched, usually caused by the stale DNS cache.
+	MismatchCalleeIDErr = "mismatch callee id"
 )
 
 // client errors

client/errs/errs.go

@@ -38,6 +38,15 @@ func IsLeaderChange(err error) bool {
 		strings.Contains(errMsg, NotPrimaryErr)
 }
 
+// IsCalleeMismatch checks whether the error is caused by callee ID mismatch.
+func IsCalleeMismatch(err error) bool {
+	if err == nil {
+		return false
+	}
+	errMsg := err.Error()
+	return strings.Contains(errMsg, MismatchCalleeIDErr)
+}
+
 // ZapError is used to make the log output easier.
 func ZapError(err error, causeError ...error) zap.Field {
 	if err == nil {

client/go.mod

@@ -44,3 +44,5 @@ require (
 	gopkg.in/natefinch/lumberjack.v2 v2.2.1 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
+
+replace github.com/pingcap/kvproto => github.com/tidbcloud/kvproto v0.0.0-20251121150221-d07b34e68f9f

client/go.sum

@@ -48,8 +48,6 @@ github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c h1:xpW9bvK+HuuTm
 github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c/go.mod h1:X2r9ueLEUZgtx2cIogM0v4Zj5uvvzhuuiu7Pn8HzMPg=
 github.com/pingcap/failpoint v0.0.0-20210918120811-547c13e3eb00 h1:C3N3itkduZXDZFh4N3vQ5HEtld3S+Y+StULhWVvumU0=
 github.com/pingcap/failpoint v0.0.0-20210918120811-547c13e3eb00/go.mod h1:4qGtCB0QK0wBzKtFEGDhxXnSnbQApw1gc9siScUl8ew=
-github.com/pingcap/kvproto v0.0.0-20250224053625-b6a98c6bf02d h1:52qhTQG8G8V/pHo/w7F4d2Tw98KMk2C+gAe3U8SWRAg=
-github.com/pingcap/kvproto v0.0.0-20250224053625-b6a98c6bf02d/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 h1:HR/ylkkLmGdSSDaD8IDP+SZrdhV1Kibl9KrHxJ9eciw=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=
 github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
@@ -79,6 +77,8 @@ github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
 github.com/stretchr/testify v1.8.2 h1:+h33VjcLVPDHtOdpUCuF+7gSuG3yGIftsP1YvFihtJ8=
 github.com/stretchr/testify v1.8.2/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
+github.com/tidbcloud/kvproto v0.0.0-20251121150221-d07b34e68f9f h1:d2HVdcW8LM+6s0vZGz/vusUqHlAK/q82x/X35whRcGk=
+github.com/tidbcloud/kvproto v0.0.0-20251121150221-d07b34e68f9f/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.3.5/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k=

client/mock_pd_service_discovery.go

@@ -88,6 +88,9 @@ func (*mockPDServiceDiscovery) GetOrCreateGRPCConn(string) (*grpc.ClientConn, er
 	return nil, nil
 }
 
+// RemoveClientConn implements the ServiceDiscovery interface.
+func (*mockPDServiceDiscovery) RemoveClientConn(string) {}
+
 // ScheduleCheckMemberChanged implements the ServiceDiscovery interface.
 func (*mockPDServiceDiscovery) ScheduleCheckMemberChanged() {}
 

client/pd_service_discovery.go

@@ -105,6 +105,8 @@ type ServiceDiscovery interface {
 	GetAllServiceClients() []ServiceClient
 	// GetOrCreateGRPCConn returns the corresponding grpc client connection of the given url.
 	GetOrCreateGRPCConn(url string) (*grpc.ClientConn, error)
+	// RemoveClientConn removes and closes the gRPC connection of the given url.
+	RemoveClientConn(url string)
 	// ScheduleCheckMemberChanged is used to trigger a check to see if there is any membership change
 	// among the leader/followers in a quorum-based cluster or among the primary/secondaries in a
 	// primary/secondary configured cluster.
@@ -1129,6 +1131,17 @@ func (c *pdServiceDiscovery) GetOrCreateGRPCConn(url string) (*grpc.ClientConn,
 	return grpcutil.GetOrCreateGRPCConn(c.ctx, &c.clientConns, url, c.tlsCfg, c.option.gRPCDialOptions...)
 }
 
+// RemoveClientConn removes and closes the grpc client connection of the given URL.
+func (c *pdServiceDiscovery) RemoveClientConn(url string) {
+	cc, ok := c.clientConns.Load(url)
+	if !ok {
+		return
+	}
+	if err := cc.(*grpc.ClientConn).Close(); err != nil {
+		log.Error("[pd] failed to close grpc clientConn", errs.ZapError(errs.ErrCloseGRPCConn, err))
+	}
+}
+
 func addrsToURLs(addrs []string, tlsCfg *tls.Config) []string {
 	// Add default schema "http://" to addrs.
 	urls := make([]string, 0, len(addrs))

client/tso_client.go

@@ -231,15 +231,21 @@ func (c *tsoClient) GetTSOAllocatorServingURLByDCLocation(dcLocation string) (st
 
 // GetTSOAllocatorClientConnByDCLocation returns the TSO allocator gRPC client connection of the given dcLocation.
 func (c *tsoClient) GetTSOAllocatorClientConnByDCLocation(dcLocation string) (*grpc.ClientConn, string) {
-	url, ok := c.tsoAllocators.Load(dcLocation)
+	u, ok := c.tsoAllocators.Load(dcLocation)
 	if !ok {
 		log.Fatal("[tso] the allocator leader should exist", zap.String("dc-location", dcLocation))
 	}
-	cc, ok := c.svcDiscovery.GetClientConns().Load(url)
-	if !ok {
-		return nil, url.(string)
+	url := u.(string)
+	cc, err := c.svcDiscovery.GetOrCreateGRPCConn(url)
+	if err != nil {
+		log.Error("[tso] failed to get tso allocator client connection",
+			zap.String("dc-location", dcLocation),
+			zap.String("serving-url", url),
+			errs.ZapError(err),
+		)
+		return nil, url
 	}
-	return cc.(*grpc.ClientConn), url.(string)
+	return cc, url
 }
 
 // AddTSOAllocatorServingURLSwitchedCallback adds callbacks which will be called

client/tso_dispatcher.go

@@ -453,8 +453,12 @@ func (td *tsoDispatcher) handleProcessRequestError(ctx context.Context, bo *retr
 	// Set `stream` to nil and remove this stream from the `connectionCtxs` due to error.
 	td.connectionCtxs.Delete(streamURL)
 	streamCancelFunc()
-	// Because ScheduleCheckMemberChanged is asynchronous, if the leader changes, we better call `updateMember` ASAP.
-	if errs.IsLeaderChange(err) {
+	switch {
+	case errs.IsCalleeMismatch(err):
+		// If callee ID mismatches, the gRPC connection is no longer valid, we need to enforce reconnecting.
+		svcDiscovery.RemoveClientConn(streamURL)
+	case errs.IsLeaderChange(err):
+		// Because ScheduleCheckMemberChanged is asynchronous, if the leader changes, we better call `updateMember` ASAP.
 		if err := bo.Exec(ctx, svcDiscovery.CheckMemberChanged); err != nil {
 			select {
 			case <-ctx.Done():

client/tso_service_discovery.go

@@ -317,6 +317,17 @@ func (c *tsoServiceDiscovery) GetOrCreateGRPCConn(url string) (*grpc.ClientConn,
 	return grpcutil.GetOrCreateGRPCConn(c.ctx, &c.clientConns, url, c.tlsCfg, c.option.gRPCDialOptions...)
 }
 
+// RemoveClientConn removes and closes the gRPC connection of the given url.
+func (c *tsoServiceDiscovery) RemoveClientConn(url string) {
+	cc, ok := c.clientConns.Load(url)
+	if !ok {
+		return
+	}
+	if err := cc.(*grpc.ClientConn).Close(); err != nil {
+		log.Error("[tso] failed to close gRPC clientConn", errs.ZapError(errs.ErrCloseGRPCConn, err))
+	}
+}
+
 // ScheduleCheckMemberChanged is used to trigger a check to see if there is any change in service endpoints.
 func (c *tsoServiceDiscovery) ScheduleCheckMemberChanged() {
 	select {

client/tso_stream.go

@@ -89,7 +89,10 @@ func (b *tsoTSOStreamBuilder) build(
 	stream, err := b.client.Tso(ctx)
 	done <- struct{}{}
 	if err == nil {
-		return newTSOStream(ctx, b.serverURL, tsoTSOStreamAdapter{stream}), nil
+		return newTSOStream(ctx, b.serverURL, tsoTSOStreamAdapter{
+			stream:   stream,
+			calleeID: b.serverURL,
+		}), nil
 	}
 	return nil, err
 }
@@ -152,7 +155,8 @@ func (s pdTSOStreamAdapter) Recv() (tsoRequestResult, error) {
 }
 
 type tsoTSOStreamAdapter struct {
-	stream tsopb.TSO_TsoClient
+	stream   tsopb.TSO_TsoClient
+	calleeID string
 }
 
 // Send implements the grpcTSOStreamAdapter interface.
@@ -162,6 +166,7 @@ func (s tsoTSOStreamAdapter) Send(clusterID uint64, keyspaceID, keyspaceGroupID
 			ClusterId:       clusterID,
 			KeyspaceId:      keyspaceID,
 			KeyspaceGroupId: keyspaceGroupID,
+			CalleeId:        s.calleeID,
 		},
 		Count:      uint32(count),
 		DcLocation: dcLocation,