grpc: fix the wrong error handler. (#5374) (#5377)

close #5373, ref #5374 Signed-off-by: bufferflies <[email protected]> Signed-off-by: nolouch <[email protected]> Co-authored-by: bufferflies <[email protected]> Co-authored-by: nolouch <[email protected]>
tikv · Jul 29, 2022 · 989134e · 989134e
1 parent f44b998
commit 989134e
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/server/grpc_service.go b/server/grpc_service.go
@@ -745,10 +745,20 @@ func (s *GrpcServer) ReportBuckets(stream pdpb.PD_ReportBucketsServer) error {
 	}()
 	for {
 		request, err := server.Recv()
+		failpoint.Inject("grpcClientClosed", func() {
+			err = status.Error(codes.Canceled, "grpc client closed")
+			request = nil
+		})
 		if err == io.EOF {
 			return nil
 		}
+		if err != nil {
+			return errors.WithStack(err)
+		}
 		forwardedHost := getForwardedHost(stream.Context())
+		failpoint.Inject("grpcClientClosed", func() {
+			forwardedHost = s.GetMember().Member().GetClientUrls()[0]
+		})
 		if !s.isLocalRequest(forwardedHost) {
 			if forwardStream == nil || lastForwardedHost != forwardedHost {
 				if cancel != nil {
@@ -1774,6 +1784,9 @@ func getForwardedHost(ctx context.Context) string {
 }
 
 func (s *GrpcServer) isLocalRequest(forwardedHost string) bool {
+	failpoint.Inject("useForwardRequest", func() {
+		failpoint.Return(false)
+	})
 	if forwardedHost == "" {
 		return true
 	}

diff --git a/tests/client/client_test.go b/tests/client/client_test.go
@@ -894,6 +894,13 @@ func (suite *clientTestSuite) TestGetRegion() {
 		return r.Buckets == nil
 	})
 	config.EnableRegionBucket = true
+
+	suite.NoError(failpoint.Enable("github.com/tikv/pd/server/grpcClientClosed", `return(true)`))
+	suite.NoError(failpoint.Enable("github.com/tikv/pd/server/useForwardRequest", `return(true)`))
+	suite.NoError(suite.reportBucket.Send(breq))
+	suite.Error(suite.reportBucket.RecvMsg(breq))
+	suite.NoError(failpoint.Disable("github.com/tikv/pd/server/grpcClientClosed"))
+	suite.NoError(failpoint.Disable("github.com/tikv/pd/server/useForwardRequest"))
 }
 
 func (suite *clientTestSuite) TestGetPrevRegion() {