mcs: fix panic of forwarding request

server/grpc_service.go

@@ -69,6 +69,7 @@
 	ErrNotStarted                       = status.Errorf(codes.Unavailable, "server not started")
 	ErrSendHeartbeatTimeout             = status.Errorf(codes.DeadlineExceeded, "send heartbeat timeout")
 	ErrNotFoundTSOAddr                  = status.Errorf(codes.NotFound, "not found tso address")
+	ErrNotFoundSchedulingAddr           = status.Errorf(codes.NotFound, "not found scheduling address")
 	ErrForwardTSOTimeout                = status.Errorf(codes.DeadlineExceeded, "forward tso request timeout")
 	ErrMaxCountTSOProxyRoutinesExceeded = status.Errorf(codes.ResourceExhausted, "max count of concurrent tso proxy routines exceeded")
 	ErrTSOProxyRecvFromClientTimeout    = status.Errorf(codes.DeadlineExceeded, "tso proxy timeout when receiving from client; stream closed by server")
@@ -1002,16 +1003,16 @@
 		s.handleDamagedStore(request.GetStats())
 		storeHeartbeatHandleDuration.WithLabelValues(storeAddress, storeLabel).Observe(time.Since(start).Seconds())
 		if s.IsAPIServiceMode() {
-			s.updateSchedulingClient(ctx)
-			if s.schedulingClient.Load() != nil {
+			forwardCli, _ := s.updateSchedulingClient(ctx)
+			if forwardCli != nil {
 				req := &schedulingpb.StoreHeartbeatRequest{
 					Header: &schedulingpb.RequestHeader{
 						ClusterId: request.GetHeader().GetClusterId(),
 						SenderId:  request.GetHeader().GetSenderId(),
 					},
 					Stats: request.GetStats(),
 				}
-				if _, err := s.schedulingClient.Load().(*schedulingClient).getClient().StoreHeartbeat(ctx, req); err != nil {
+				if _, err := forwardCli.StoreHeartbeat(ctx, req); err != nil {
 					// reset to let it be updated in the next request
 					s.schedulingClient.Store(&schedulingClient{})
 				}
@@ -1030,19 +1031,22 @@
 	return resp, nil
 }
 
-func (s *GrpcServer) updateSchedulingClient(ctx context.Context) {
+func (s *GrpcServer) updateSchedulingClient(ctx context.Context) (schedulingpb.SchedulingClient, error) {
 	forwardedHost, _ := s.GetServicePrimaryAddr(ctx, utils.SchedulingServiceName)
 	pre := s.schedulingClient.Load()
 	if forwardedHost != "" && ((pre == nil) || (pre != nil && forwardedHost != pre.(*schedulingClient).getPrimaryAddr())) {
 		client, err := s.getDelegateClient(ctx, forwardedHost)
 		if err != nil {
 			log.Error("get delegate client failed", zap.Error(err))
 		}
-		s.schedulingClient.Store(&schedulingClient{
+		forwardCli := &schedulingClient{
 			client:      schedulingpb.NewSchedulingClient(client),
 			lastPrimary: forwardedHost,
-		})
+		}
+		s.schedulingClient.Store(forwardCli)
+		return forwardCli.getClient(), nil
 	}
+	return nil, ErrNotFoundSchedulingAddr
 }
 
 // bucketHeartbeatServer wraps PD_ReportBucketsServer to ensure when any error
@@ -1771,31 +1775,36 @@
 // ScatterRegion implements gRPC PDServer.
 func (s *GrpcServer) ScatterRegion(ctx context.Context, request *pdpb.ScatterRegionRequest) (*pdpb.ScatterRegionResponse, error) {
 	if s.IsAPIServiceMode() {
-		s.updateSchedulingClient(ctx)
-		if s.schedulingClient.Load() != nil {
+		forwardCli, err := s.updateSchedulingClient(ctx)
+		if err != nil {
+			return &pdpb.ScatterRegionResponse{
+				Header: s.wrapErrorToHeader(pdpb.ErrorType_UNKNOWN, err.Error()),
+			}, nil
+		}
+		if forwardCli != nil {
 			regionsID := request.GetRegionsId()
 			if len(regionsID) == 0 {
 				return &pdpb.ScatterRegionResponse{
 					Header: s.invalidValue("regions id is required"),
 				}, nil
 			}
 			req := &schedulingpb.ScatterRegionsRequest{
 				Header: &schedulingpb.RequestHeader{
 					ClusterId: request.GetHeader().GetClusterId(),
 					SenderId:  request.GetHeader().GetSenderId(),
 				},
 				RegionsId:      regionsID,
 				Group:          request.GetGroup(),
 				RetryLimit:     request.GetRetryLimit(),
 				SkipStoreLimit: request.GetSkipStoreLimit(),
 			}
-			resp, err := s.schedulingClient.Load().(*schedulingClient).getClient().ScatterRegions(ctx, req)
+			resp, err := forwardCli.ScatterRegions(ctx, req)
 			if err != nil {
 				// reset to let it be updated in the next request
 				s.schedulingClient.Store(&schedulingClient{})
 				return s.convertScatterResponse(resp), err
 			}
 			return s.convertScatterResponse(resp), nil
 		}
 	}
 
@@ -1990,22 +1999,27 @@
 // GetOperator gets information about the operator belonging to the specify region.
 func (s *GrpcServer) GetOperator(ctx context.Context, request *pdpb.GetOperatorRequest) (*pdpb.GetOperatorResponse, error) {
 	if s.IsAPIServiceMode() {
-		s.updateSchedulingClient(ctx)
-		if s.schedulingClient.Load() != nil {
+		forwardCli, err := s.updateSchedulingClient(ctx)
+		if err != nil {
+			return &pdpb.GetOperatorResponse{
+				Header: s.wrapErrorToHeader(pdpb.ErrorType_UNKNOWN, err.Error()),
+			}, nil
+		}
+		if forwardCli != nil {
 			req := &schedulingpb.GetOperatorRequest{
 				Header: &schedulingpb.RequestHeader{
 					ClusterId: request.GetHeader().GetClusterId(),
 					SenderId:  request.GetHeader().GetSenderId(),
 				},
 				RegionId: request.GetRegionId(),
 			}
-			resp, err := s.schedulingClient.Load().(*schedulingClient).getClient().GetOperator(ctx, req)
+			resp, err := forwardCli.GetOperator(ctx, req)
 			if err != nil {
 				// reset to let it be updated in the next request
 				s.schedulingClient.Store(&schedulingClient{})
 				return s.convertOperatorResponse(resp), err
 			}
 			return s.convertOperatorResponse(resp), nil
 		}
 	}
 	fn := func(ctx context.Context, client *grpc.ClientConn) (interface{}, error) {
@@ -2098,41 +2112,41 @@
 }
 
 func (s *GrpcServer) convertHeader(header *schedulingpb.ResponseHeader) *pdpb.ResponseHeader {
 	switch header.GetError().GetType() {
 	case schedulingpb.ErrorType_UNKNOWN:
 		return &pdpb.ResponseHeader{
 			ClusterId: header.GetClusterId(),
 			Error: &pdpb.Error{
 				Type:    pdpb.ErrorType_UNKNOWN,
 				Message: header.GetError().GetMessage(),
 			},
 		}
 	default:
 		return &pdpb.ResponseHeader{ClusterId: header.GetClusterId()}
 	}
 }
 
 func (s *GrpcServer) convertSplitResponse(resp *schedulingpb.SplitRegionsResponse) *pdpb.SplitRegionsResponse {
 	return &pdpb.SplitRegionsResponse{
 		Header:             s.convertHeader(resp.GetHeader()),
 		FinishedPercentage: resp.GetFinishedPercentage(),
 	}
 }
 
 func (s *GrpcServer) convertScatterResponse(resp *schedulingpb.ScatterRegionsResponse) *pdpb.ScatterRegionResponse {
 	return &pdpb.ScatterRegionResponse{
 		Header:             s.convertHeader(resp.GetHeader()),
 		FinishedPercentage: resp.GetFinishedPercentage(),
 	}
 }
 
 func (s *GrpcServer) convertOperatorResponse(resp *schedulingpb.GetOperatorResponse) *pdpb.GetOperatorResponse {
 	return &pdpb.GetOperatorResponse{
 		Header:   s.convertHeader(resp.GetHeader()),
 		RegionId: resp.GetRegionId(),
 		Desc:     resp.GetDesc(),
 		Kind:     resp.GetKind(),
 		Status:   resp.GetStatus(),
 	}
 }
 
@@ -2240,24 +2254,29 @@
 
 // SplitRegions split regions by the given split keys
 func (s *GrpcServer) SplitRegions(ctx context.Context, request *pdpb.SplitRegionsRequest) (*pdpb.SplitRegionsResponse, error) {
 	if s.IsAPIServiceMode() {
-		s.updateSchedulingClient(ctx)
-		if s.schedulingClient.Load() != nil {
+		forwardCli, err := s.updateSchedulingClient(ctx)
+		if err != nil {
+			return &pdpb.SplitRegionsResponse{
+				Header: s.wrapErrorToHeader(pdpb.ErrorType_UNKNOWN, err.Error()),
+			}, nil
+		}
+		if forwardCli != nil {
 			req := &schedulingpb.SplitRegionsRequest{
 				Header: &schedulingpb.RequestHeader{
 					ClusterId: request.GetHeader().GetClusterId(),
 					SenderId:  request.GetHeader().GetSenderId(),
 				},
 				SplitKeys:  request.GetSplitKeys(),
 				RetryLimit: request.GetRetryLimit(),
 			}
-			resp, err := s.schedulingClient.Load().(*schedulingClient).getClient().SplitRegions(ctx, req)
+			resp, err := forwardCli.SplitRegions(ctx, req)
 			if err != nil {
 				// reset to let it be updated in the next request
 				s.schedulingClient.Store(&schedulingClient{})
 				return s.convertSplitResponse(resp), err
 			}
 			return s.convertSplitResponse(resp), nil
 		}
 	}