vmdriver(vz): Wait for Start() to complete on server side

Signed-off-by: Ansuman Sahoo <[email protected]>

Author: unsuman

pkg/driver/external/client/methods.go

@@ -67,19 +67,37 @@ func (d *DriverClient) Start(ctx context.Context) (chan error, error) {
 		return nil, err
 	}
 
+	// Blocking to receive an initial response to ensure Start() is initiated
+	// at the server-side.
+	initialResp, err := stream.Recv()
+	if err != nil {
+		d.logger.Errorf("Error receiving initial response from driver start: %v", err)
+		return nil, err
+	}
+	if !initialResp.Success {
+		return nil, errors.New(initialResp.Error)
+	}
+
+	go func() {
+		<-ctx.Done()
+		if closeErr := stream.CloseSend(); closeErr != nil {
+			d.logger.Warnf("Failed to close stream: %v", closeErr)
+		}
+	}()
+
 	errCh := make(chan error, 1)
 	go func() {
 		for {
-			errorStream, err := stream.Recv()
+			respStream, err := stream.Recv()
 			if err != nil {
-				d.logger.Errorf("Error receiving response from driver: %v", err)
+				d.logger.Infof("Error receiving response from driver: %v", err)
 				return
 			}
-			d.logger.Debugf("Received response: %v", errorStream)
-			if !errorStream.Success {
-				errCh <- errors.New(errorStream.Error)
+			d.logger.Debugf("Received response: %v", respStream)
+			if !respStream.Success {
+				errCh <- errors.New(respStream.Error)
 			} else {
-				errCh <- nil
+				close(errCh)
 				return
 			}
 		}

pkg/driver/external/server/methods.go

@@ -26,9 +26,20 @@ func (s *DriverServer) Start(_ *emptypb.Empty, stream pb.Driver_StartServer) err
 	errChan, err := s.driver.Start(stream.Context())
 	if err != nil {
 		s.logger.Errorf("Start failed: %v", err)
+		if sendErr := stream.Send(&pb.StartResponse{Success: false, Error: err.Error()}); sendErr != nil {
+			s.logger.Errorf("Failed to send error response: %v", sendErr)
+			return status.Errorf(codes.Internal, "failed to send error response: %v", sendErr)
+		}
 		return status.Errorf(codes.Internal, "failed to start driver: %v", err)
 	}
 
+	// First send a success response upon receiving the errChan to unblock the client
+	// and start receiving errors (if any).
+	if err := stream.Send(&pb.StartResponse{Success: true}); err != nil {
+		s.logger.Errorf("Failed to send success response: %v", err)
+		return status.Errorf(codes.Internal, "failed to send success response: %v", err)
+	}
+
 	for {
 		select {
 		case err, ok := <-errChan:

pkg/driver/vz/vz_driver_darwin.go

@@ -440,7 +440,7 @@ func (l *LimaVzDriver) ForwardGuestAgent() bool {
 	return l.vSockPort == 0 && l.virtioPort == ""
 }
 
-func (l *LimaVzDriver) AdditionalSetupForSSH(_ context.Context) error {
+func (l *LimaVzDriver) AdditionalSetupForSSH(ctx context.Context) error {
 	<-l.waitSSHLocalPortAccessible
 	return nil
 }