mcs: close the connection to avoid retrying to connect

pkg/utils/grpcutil/grpcutil.go

@@ -252,11 +252,11 @@ func CheckStream(ctx context.Context, cancel context.CancelFunc, done chan struc
 
 // NeedRebuildConnection checks if the error is a connection error.
 func NeedRebuildConnection(err error) bool {
-	return err == io.EOF ||
+	return (err != nil) && (err == io.EOF ||
 		strings.Contains(err.Error(), codes.Unavailable.String()) || // Unavailable indicates the service is currently unavailable. This is a most likely a transient condition.
 		strings.Contains(err.Error(), codes.DeadlineExceeded.String()) || // DeadlineExceeded means operation expired before completion.
 		strings.Contains(err.Error(), codes.Internal.String()) || // Internal errors.
 		strings.Contains(err.Error(), codes.Unknown.String()) || // Unknown error.
-		strings.Contains(err.Error(), codes.ResourceExhausted.String()) // ResourceExhausted is returned when either the client or the server has exhausted their resources.
+		strings.Contains(err.Error(), codes.ResourceExhausted.String())) // ResourceExhausted is returned when either the client or the server has exhausted their resources.
 	// Besides, we don't need to rebuild the connection if the code is Canceled, which means the client cancelled the request.
 }

server/forward.go

@@ -89,13 +89,17 @@ func (s *GrpcServer) forwardTSO(stream pdpb.PD_TsoServer) error {
 		forwardStream     tsopb.TSO_TsoClient
 		forwardCtx        context.Context
 		cancelForward     context.CancelFunc
+		tsoStreamErr      error
 		lastForwardedHost string
 	)
 	defer func() {
 		s.concurrentTSOProxyStreamings.Add(-1)
 		if cancelForward != nil {
 			cancelForward()
 		}
+		if grpcutil.NeedRebuildConnection(tsoStreamErr) {
+			s.closeDelegateClient(lastForwardedHost)
+		}
 	}()
 
 	maxConcurrentTSOProxyStreamings := int32(s.GetMaxConcurrentTSOProxyStreamings())
@@ -131,7 +135,8 @@ func (s *GrpcServer) forwardTSO(stream pdpb.PD_TsoServer) error {
 
 		forwardedHost, ok := s.GetServicePrimaryAddr(stream.Context(), utils.TSOServiceName)
 		if !ok || len(forwardedHost) == 0 {
-			return errors.WithStack(ErrNotFoundTSOAddr)
+			tsoStreamErr = errors.WithStack(ErrNotFoundTSOAddr)
+			return tsoStreamErr
 		}
 		if forwardStream == nil || lastForwardedHost != forwardedHost {
 			if cancelForward != nil {
@@ -140,18 +145,21 @@ func (s *GrpcServer) forwardTSO(stream pdpb.PD_TsoServer) error {
 
 			clientConn, err := s.getDelegateClient(s.ctx, forwardedHost)
 			if err != nil {
-				return errors.WithStack(err)
+				tsoStreamErr = errors.WithStack(err)
+				return tsoStreamErr
 			}
 			forwardStream, forwardCtx, cancelForward, err = s.createTSOForwardStream(stream.Context(), clientConn)
 			if err != nil {
-				return errors.WithStack(err)
+				tsoStreamErr = errors.WithStack(err)
+				return tsoStreamErr
 			}
 			lastForwardedHost = forwardedHost
 		}
 
 		tsopbResp, err := s.forwardTSORequestWithDeadLine(forwardCtx, cancelForward, forwardStream, request, tsDeadlineCh)
 		if err != nil {
-			return errors.WithStack(err)
+			tsoStreamErr = errors.WithStack(err)
+			return tsoStreamErr
 		}
 
 		// The error types defined for tsopb and pdpb are different, so we need to convert them.
@@ -363,25 +371,13 @@ func (s *GrpcServer) getDelegateClient(ctx context.Context, forwardedHost string
 	return conn.(*grpc.ClientConn), nil
 }
 
-func (s *GrpcServer) getForwardedHost(ctx, streamCtx context.Context, serviceName ...string) (forwardedHost string, err error) {
-	if s.IsAPIServiceMode() {
-		var ok bool
-		if len(serviceName) == 0 {
-			return "", ErrNotFoundService
-		}
-		forwardedHost, ok = s.GetServicePrimaryAddr(ctx, serviceName[0])
-		if !ok || len(forwardedHost) == 0 {
-			switch serviceName[0] {
-			case utils.TSOServiceName:
-				return "", ErrNotFoundTSOAddr
-			case utils.SchedulingServiceName:
-				return "", ErrNotFoundSchedulingAddr
-			}
-		}
-	} else if fh := grpcutil.GetForwardedHost(streamCtx); !s.isLocalRequest(fh) {
-		forwardedHost = fh
+func (s *GrpcServer) closeDelegateClient(forwardedHost string) {
+	client, ok := s.clientConns.LoadAndDelete(forwardedHost)
+	if !ok {
+		return
 	}
-	return forwardedHost, nil
+	client.(*grpc.ClientConn).Close()
+	log.Debug("close delegate client connection", zap.String("forwarded-host", forwardedHost))
 }
 
 func (s *GrpcServer) isLocalRequest(host string) bool {

server/grpc_service.go

@@ -547,9 +547,8 @@ func (s *GrpcServer) Tso(stream pdpb.PD_TsoServer) error {
 			return errors.WithStack(err)
 		}
 
-		if forwardedHost, err := s.getForwardedHost(ctx, stream.Context(), utils.TSOServiceName); err != nil {
-			return err
-		} else if len(forwardedHost) > 0 {
+		forwardedHost := grpcutil.GetForwardedHost(stream.Context())
+		if !s.isLocalRequest(forwardedHost) {
 			clientConn, err := s.getDelegateClient(s.ctx, forwardedHost)
 			if err != nil {
 				return errors.WithStack(err)
@@ -1332,6 +1331,7 @@ func (s *GrpcServer) RegionHeartbeat(stream pdpb.PD_RegionHeartbeatServer) error
 				if cancel != nil {
 					cancel()
 				}
+
 				client, err := s.getDelegateClient(s.ctx, forwardedSchedulingHost)
 				if err != nil {
 					errRegionHeartbeatClient.Inc()
@@ -1370,6 +1370,9 @@ func (s *GrpcServer) RegionHeartbeat(stream pdpb.PD_RegionHeartbeatServer) error
 			}
 			if err := forwardSchedulingStream.Send(schedulingpbReq); err != nil {
 				forwardSchedulingStream = nil
+				if grpcutil.NeedRebuildConnection(err) {
+					s.closeDelegateClient(lastForwardedSchedulingHost)
+				}
 				errRegionHeartbeatSend.Inc()
 				log.Error("failed to send request to scheduling service", zap.Error(err))
 			}