From 0626a9313233daf1f7510f36dba757fcdcde98a3 Mon Sep 17 00:00:00 2001
From: Ryan Leung <rleungx@gmail.com>
Date: Wed, 31 Jan 2024 15:02:05 +0800
Subject: [PATCH] address comments

Signed-off-by: Ryan Leung <rleungx@gmail.com>
---
 pkg/utils/grpcutil/grpcutil.go | 5 +----
 server/forward.go              | 2 +-
 2 files changed, 2 insertions(+), 5 deletions(-)

diff --git a/pkg/utils/grpcutil/grpcutil.go b/pkg/utils/grpcutil/grpcutil.go
index 9823adf2cfb..4d2f250d53e 100644
--- a/pkg/utils/grpcutil/grpcutil.go
+++ b/pkg/utils/grpcutil/grpcutil.go
@@ -252,10 +252,7 @@ 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 {
-	if err == nil {
-		return false
-	}
-	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.
diff --git a/server/forward.go b/server/forward.go
index 5be798f095d..82a145e47b7 100644
--- a/server/forward.go
+++ b/server/forward.go
@@ -159,7 +159,7 @@ func (s *GrpcServer) forwardTSO(stream pdpb.PD_TsoServer) error {
 		tsopbResp, err := s.forwardTSORequestWithDeadLine(forwardCtx, cancelForward, forwardStream, request, tsDeadlineCh)
 		if err != nil {
 			tsoStreamErr = errors.WithStack(err)
-			return errors.WithStack(err)
+			return tsoStreamErr
 		}
 
 		// The error types defined for tsopb and pdpb are different, so we need to convert them.