vmdriver(vz): wait for Start() to complete RPC and fix race condition for AdditionalSetupForSSH()

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

Author: unsuman

pkg/driver/external/client/methods.go

@@ -7,6 +7,7 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
+	"io"
 	"net"
 	"time"
 
@@ -72,7 +73,9 @@ func (d *DriverClient) Start(ctx context.Context) (chan error, error) {
 		for {
 			errorStream, err := stream.Recv()
 			if err != nil {
-				d.logger.Errorf("Error receiving response from driver: %v", err)
+				if !errors.Is(err, io.EOF) {
+					d.logger.Errorf("Error receiving response from driver: %v", err)
+				}
 				return
 			}
 			d.logger.Debugf("Received response: %v", errorStream)

pkg/driver/external/server/methods.go

@@ -23,35 +23,39 @@ import (
 
 func (s *DriverServer) Start(_ *emptypb.Empty, stream pb.Driver_StartServer) error {
 	s.logger.Debug("Received Start request")
-	errChan, err := s.driver.Start(stream.Context())
+	errChan, err := s.driver.Start(context.Background())
 	if err != nil {
 		s.logger.Errorf("Start failed: %v", err)
 		return status.Errorf(codes.Internal, "failed to start driver: %v", err)
 	}
 
-	for {
-		select {
-		case err, ok := <-errChan:
-			if !ok {
-				s.logger.Debug("Start error channel closed")
-				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)
+	go func() {
+		for {
+			select {
+			case err, ok := <-errChan:
+				if !ok {
+					s.logger.Debug("Start error channel closed")
+					if err := stream.Send(&pb.StartResponse{Success: true}); err != nil {
+						s.logger.Errorf("Failed to send success response: %v", err)
+						return
+					}
+					return
 				}
-				return nil
-			}
-			if err != nil {
-				s.logger.Errorf("Error during Start: %v", err)
-				if err := stream.Send(&pb.StartResponse{Error: err.Error(), Success: false}); err != nil {
-					s.logger.Errorf("Failed to send error response: %v", err)
-					return status.Errorf(codes.Internal, "failed to send error response: %v", err)
+				if err != nil {
+					s.logger.Errorf("Error during Start: %v", err)
+					if err := stream.Send(&pb.StartResponse{Error: err.Error(), Success: false}); err != nil {
+						s.logger.Errorf("Failed to send error response: %v", err)
+						return
+					}
 				}
+			case <-stream.Context().Done():
+				s.logger.Debug("Stream context done, stopping Start")
+				return
 			}
-		case <-stream.Context().Done():
-			s.logger.Debug("Stream context done, stopping Start")
-			return nil
 		}
-	}
+	}()
+
+	return nil
 }
 
 func (s *DriverServer) Configure(_ context.Context, req *pb.SetConfigRequest) (*emptypb.Empty, error) {

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:
+		}
+	}
 }