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

Signed-off-by: Ansuman Sahoo <anshumansahoo500@gmail.com>

Author: unsuman

pkg/driver/external/client/methods.go

@@ -67,8 +67,22 @@ 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.
+	for {
+		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 {
+			break
+		}
+	}
+
 	errCh := make(chan error, 1)
 	go func() {
+		defer stream.CloseSend()
 		for {
 			errorStream, err := stream.Recv()
 			if err != nil {

pkg/driver/external/server/methods.go

@@ -29,6 +29,13 @@ func (s *DriverServer) Start(_ *emptypb.Empty, stream pb.Driver_StartServer) err
 		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,24 @@ func (l *LimaVzDriver) ForwardGuestAgent() bool {
 	return l.vSockPort == 0 && l.virtioPort == ""
 }
 
-func (l *LimaVzDriver) AdditionalSetupForSSH(_ context.Context) error {
-	<-l.waitSSHLocalPortAccessible
-	return nil
+func (l *LimaVzDriver) AdditionalSetupForSSH(ctx context.Context) error {
+	ticker := time.NewTicker(100 * time.Millisecond)
+	defer ticker.Stop()
+
+	timeout := time.After(60 * time.Second)
+
+	for {
+		if l.waitSSHLocalPortAccessible != nil {
+			<-l.waitSSHLocalPortAccessible
+			return nil
+		}
+
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-timeout:
+			return errors.New("timeout waiting for Start() to initialize")
+		case <-ticker.C:
+		}
+	}
 }