refactor(machine): shim.Start support ctx for krunkit part

cmd/bauklotze/machine/start.go

@@ -129,73 +129,62 @@ func start(cmd *cobra.Command, args []string) error {
 	//    this will block the current goroutine, until the network provider and Hypervisor exit
 	// 3. if the network provider or Hypervisor got exit, clean up files and kill all child process
 	g.Go(func() error {
-		logrus.Infof("Starting machine %q", mc.VMName)
-		// Start the network and hypervisor, shim.Start is non-block func
-		_, err = shim.Start(ctx, mc)
-		defer cleanUp(mc) // Clean tmp files at the end
-		if err != nil {
-			return fmt.Errorf("failed to start machine: %w", err)
-		}
+		return startMachine(ctx, mc)
+	})
 
-		// Test the machine ssh connection by consume a string with '\n'
-		if shim.ConductVMReadinessCheck(ctx, mc) {
-			logrus.Infof("Machine %s SSH is ready, using sshkey %s with %s, listen in %d",
-				mc.VMName, mc.SSH.IdentityPath, mc.SSH.RemoteUsername, mc.SSH.Port)
-		} else {
-			return fmt.Errorf("machine ssh is not ready")
-		}
+	defer cleanUp(mc) // Clean tmp files at the end
+	return g.Wait()   //nolint:wrapcheck
+}
 
-		// Test the podman api which forwarded from host to guest
-		err = machine.WaitPodmanReady(ctx, mc.GvProxy.ForwardInfo["forward-sock"][0])
-		if err != nil {
-			return fmt.Errorf("failed to ping podman api: %w", err)
-		}
-		events.NotifyRun(events.Ready)
-		logrus.Infof("Machine %s Podman API listened in %q", mc.VMName, mc.GvProxy.ForwardInfo["forward-sock"][0])
+func startMachine(parentCtx context.Context, mc *vmconfig.MachineConfig) error {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	context.AfterFunc(parentCtx, func() {
+		SyncDisk(mc)
+		cancel()
+	})
 
-		// Start Sleep Notifier and dispatch tasks
-		logrus.Infof("Start Sleep Notifier and dispatch tasks")
-		go shim.SleepNotifier(ctx, mc)
+	logrus.Infof("Starting machine %q", mc.VMName)
+	// Start the network and hypervisor, shim.Start is non-block func
+	_, err := shim.Start(ctx, mc)
 
-		if err = mc.UpdateLastBoot(); err != nil {
-			logrus.Errorf("failed to update last boot time: %v", err)
-		}
+	if err != nil {
+		return fmt.Errorf("failed to start machine: %w", err)
+	}
 
-		logrus.Infof("Machine start successful, wait the network provider and hypervisor to exit")
-		gvpErrChan := make(chan error, 1)
-		go func() {
-			if mc.GvpCmd != nil {
-				logrus.Infof("Waiting for gvisor network provider to exit")
-				gvpErrChan <- mc.GvpCmd.Wait()
-			}
-		}()
+	// Test the machine ssh connection by consume a string with '\n'
+	if !shim.ConductVMReadinessCheck(ctx, mc) {
+		return fmt.Errorf("machine ssh is not ready")
+	}
+	logrus.Infof("Machine %s SSH is ready, using sshkey %s with %s, listen in %d",
+		mc.VMName, mc.SSH.IdentityPath, mc.SSH.RemoteUsername, mc.SSH.Port)
 
-		vmmErrChan := make(chan error, 1)
-		go func() {
-			if mc.VmmCmd != nil {
-				logrus.Infof("Waiting for hypervisor to exit")
-				vmmErrChan <- mc.VmmCmd.Wait()
-			}
-		}()
+	// Test the podman api which forwarded from host to guest
+	err = machine.WaitPodmanReady(ctx, mc.GvProxy.ForwardInfo["forward-sock"][0])
+	if err != nil {
+		return fmt.Errorf("failed to ping podman api: %w", err)
+	}
+	events.NotifyRun(events.Ready)
+	logrus.Infof("Machine %s Podman API listened in %q", mc.VMName, mc.GvProxy.ForwardInfo["forward-sock"][0])
 
-		select {
-		case <-ctx.Done():
-			return context.Cause(ctx) //nolint:wrapcheck
-		case err = <-gvpErrChan:
-			return fmt.Errorf("gvproxy exited with error: %w", err)
-		case err = <-vmmErrChan:
-			return fmt.Errorf("hypervisor exited with error: %w", err)
-		}
-	})
+	// Start Sleep Notifier and dispatch tasks
+	logrus.Infof("Start Sleep Notifier and dispatch tasks")
+	go shim.SleepNotifier(ctx, mc)
 
-	return g.Wait() //nolint:wrapcheck
+	if err = mc.UpdateLastBoot(); err != nil {
+		logrus.Warnf("failed to update last boot time: %v", err)
+	}
+
+	logrus.Infof("Machine start successful")
+
+	return shim.Wait(ctx, mc) //nolint:wrapcheck
 }
 
 // cleanUp deletes the temporary socks file and terminates the child process using
 // cmd.Process.Kill()
 func cleanUp(mc *vmconfig.MachineConfig) {
 	events.NotifyRun(events.SyncMachineDisk)
-	SyncDisk(mc) // Sync the disk to make sure all data is written to the disk
+
 	logrus.Infof("Start clean up files")
 	gvpBackendSocket, _ := mc.GVProxyNetworkBackendSocks()
 	_ = gvpBackendSocket.Delete(true)
@@ -208,9 +197,6 @@ func cleanUp(mc *vmconfig.MachineConfig) {
 
 	gvpPidFile := &io.VMFile{Path: mc.GvProxy.PidFile}
 	_ = gvpPidFile.Delete(true)
-
-	logrus.Infof("Start clean up process")
-	system.KillCmdWithWarn(mc.VmmCmd, mc.GvpCmd)
 }
 
 func SyncDisk(mc *vmconfig.MachineConfig) {

pkg/machine/krunkit/machine.go

@@ -6,18 +6,12 @@
 package krunkit
 
 import (
-	"bauklotze/pkg/machine/events"
-	"bauklotze/pkg/machine/helper"
-	"bauklotze/pkg/machine/vmconfig"
 	"fmt"
-	"io"
-	"os"
 
-	"bauklotze/pkg/machine/ignition"
-	mypty "bauklotze/pkg/pty"
+	"bauklotze/pkg/machine/helper"
+	"bauklotze/pkg/machine/vmconfig"
 
 	vfConfig "github.com/crc-org/vfkit/pkg/config"
-	"github.com/sirupsen/logrus"
 )
 
 const applehvMACAddress = "5a:94:ef:e4:0c:ee"
@@ -63,55 +57,3 @@ func setupDevices(mc *vmconfig.MachineConfig) ([]vfConfig.VirtioDevice, error) {
 
 	return devices, nil
 }
-
-func startKrunkit(mc *vmconfig.MachineConfig) error {
-	bootloader := mc.AppleKrunkitHypervisor.Krunkit.VirtualMachine.Bootloader
-	if bootloader == nil {
-		return fmt.Errorf("unable to determine boot loader for this machine")
-	}
-
-	vmc := vfConfig.NewVirtualMachine(uint(mc.Resources.CPUs), uint64(mc.Resources.Memory), bootloader)
-
-	defaultDevices, err := setupDevices(mc)
-	if err != nil {
-		return fmt.Errorf("failed to get default devices: %w", err)
-	}
-	vmc.Devices = append(vmc.Devices, defaultDevices...)
-
-	krunkitBin := mc.Dirs.Hypervisor.Bin.GetPath()
-	logrus.Infof("krunkit binary path is: %s", krunkitBin)
-
-	krunkitCmd, err := vmc.Cmd(krunkitBin)
-	if err != nil {
-		return fmt.Errorf("failed to create krunkit command: %w", err)
-	}
-	libsDir := mc.Dirs.Hypervisor.LibsDir.GetPath()
-	krunkitCmd.Env = append(krunkitCmd.Env, fmt.Sprintf("DYLD_LIBRARY_PATH=%s", libsDir))
-
-	// Add the "krun-log-level" allflag for setting up the desired log level for libkrun's debug facilities.
-	// Log level for libkrun (0=off, 1=error, 2=warn, 3=info, 4=debug, 5 or higher=trace)
-	krunkitCmd.Args = append(krunkitCmd.Args, "--krun-log-level", "3")
-
-	err = ignition.GenerateIgnScripts(mc)
-	if err != nil {
-		return fmt.Errorf("failed to generate ignition scripts: %w", err)
-	}
-
-	logrus.Infof("FULL KRUNKIT CMDLINE:%q", krunkitCmd.Args)
-	events.NotifyRun(events.StartVMProvider, "krunkit staring...")
-
-	// Run krunkit in pty, the pty should never close because the krunkit is a background process
-	ptmx, err := mypty.RunInPty(krunkitCmd)
-	if err != nil {
-		return fmt.Errorf("failed to run krunkit in pty: %w", err)
-	}
-	mc.VmmCmd = krunkitCmd
-
-	go func() {
-		_, _ = io.Copy(os.Stdout, ptmx)
-	}()
-
-	events.NotifyRun(events.StartVMProvider, "krunkit started")
-
-	return nil
-}

pkg/machine/krunkit/stubber.go

@@ -6,14 +6,22 @@
 package krunkit
 
 import (
+	"context"
+	"fmt"
+	"io"
+	"os"
+	"os/exec"
+	"strconv"
+
 	"bauklotze/pkg/decompress"
 	"bauklotze/pkg/machine/defconfig"
 	"bauklotze/pkg/machine/define"
+	"bauklotze/pkg/machine/events"
+	"bauklotze/pkg/machine/ignition"
 	"bauklotze/pkg/machine/vmconfig"
 	"bauklotze/pkg/machine/volumes"
 	"bauklotze/pkg/port"
-	"fmt"
-	"strconv"
+	mypty "bauklotze/pkg/pty"
 
 	gvproxy "github.com/containers/gvisor-tap-vsock/pkg/types"
 	vfConfig "github.com/crc-org/vfkit/pkg/config"
@@ -72,6 +80,55 @@ func (l LibKrunStubber) VMType() defconfig.VMType {
 	return defconfig.LibKrun
 }
 
-func (l LibKrunStubber) StartVM(mc *vmconfig.MachineConfig) error {
-	return startKrunkit(mc)
+func (l LibKrunStubber) StartVM(ctx context.Context, mc *vmconfig.MachineConfig) error {
+	bootloader := mc.AppleKrunkitHypervisor.Krunkit.VirtualMachine.Bootloader
+	if bootloader == nil {
+		return fmt.Errorf("unable to determine boot loader for this machine")
+	}
+
+	vmc := vfConfig.NewVirtualMachine(uint(mc.Resources.CPUs), uint64(mc.Resources.Memory), bootloader)
+
+	defaultDevices, err := setupDevices(mc)
+	if err != nil {
+		return fmt.Errorf("failed to get default devices: %w", err)
+	}
+	vmc.Devices = append(vmc.Devices, defaultDevices...)
+
+	krunkitBin := mc.Dirs.Hypervisor.Bin.GetPath()
+	logrus.Infof("krunkit binary path is: %s", krunkitBin)
+
+	cmd, err := vmc.Cmd(krunkitBin)
+	if err != nil {
+		return fmt.Errorf("failed to create krunkit command: %w", err)
+	}
+	libsDir := mc.Dirs.Hypervisor.LibsDir.GetPath()
+
+	// Add the "krun-log-level" allflag for setting up the desired log level for libkrun's debug facilities.
+	// Log level for libkrun (0=off, 1=error, 2=warn, 3=info, 4=debug, 5 or higher=trace)
+	cmd.Args = append(cmd.Args, "--krun-log-level", "3")
+
+	myKrunKitCmd := exec.CommandContext(ctx, krunkitBin, cmd.Args[1:]...)
+	myKrunKitCmd.Env = append(myKrunKitCmd.Env, fmt.Sprintf("DYLD_LIBRARY_PATH=%s", libsDir))
+
+	if err = ignition.GenerateIgnScripts(mc); err != nil {
+		return fmt.Errorf("failed to generate ignition scripts: %w", err)
+	}
+
+	logrus.Infof("FULL KRUNKIT CMDLINE: %q", myKrunKitCmd.Args)
+	events.NotifyRun(events.StartVMProvider, "krunkit staring...")
+
+	// Run krunkit in pty, the pty should never close because the krunkit is a background process
+	ptmx, err := mypty.RunInPty(myKrunKitCmd)
+	if err != nil {
+		return fmt.Errorf("failed to run krunkit in pty: %w", err)
+	}
+	mc.VmmCmd = myKrunKitCmd
+
+	go func() {
+		_, _ = io.Copy(os.Stdout, ptmx)
+	}()
+
+	events.NotifyRun(events.StartVMProvider, "krunkit started")
+
+	return nil
 }

pkg/machine/shim/host.go

@@ -147,15 +147,46 @@ func getMCsOverProviders(stubs []vmconfig.VMProvider) (map[string]*vmconfig.Mach
 	return mcs, nil
 }
 
-func startNetworkProvider(mc *vmconfig.MachineConfig) error {
+func Wait(ctx context.Context, mc *vmconfig.MachineConfig) error {
+	errChanGvp := make(chan error, 1) //nolint:gomnd
+	go func() {
+		if mc.GvpCmd != nil {
+			logrus.Infoln("Waiting for gvisor network provider to exit")
+			errChanGvp <- mc.GvpCmd.Wait()
+		} else {
+			logrus.Warnf("GvpCmd is nil, gvp network provider is not started")
+		}
+	}()
+
+	errChanVmm := make(chan error, 1) //nolint:gomnd
+	go func() {
+		if mc.VmmCmd != nil {
+			logrus.Infoln("Waiting for hypervisor to exit")
+			errChanVmm <- mc.VmmCmd.Wait()
+		} else {
+			logrus.Warnf("VmmCmd is nil, hypervisor is not started")
+		}
+	}()
+
+	select {
+	case <-ctx.Done():
+		return fmt.Errorf("context done: %w", context.Cause(ctx))
+	case err := <-errChanGvp:
+		return fmt.Errorf("network provider exit: %w", err)
+	case err := <-errChanVmm:
+		return fmt.Errorf("hypervisor exit: %w", err)
+	}
+}
+
+func startNetworkProvider(ctx context.Context, mc *vmconfig.MachineConfig) error {
 	// p is the port maybe changed different from the default port which 6123
 	p, err := port.GetFree(mc.SSH.Port)
 	if err != nil {
 		return fmt.Errorf("failed to get free port: %w", err)
 	}
 	// set the port to the machine configure
 	mc.SSH.Port = p
-	podmanAPISocksInHost, podmanAPISocksInGuest, err := startNetworking(mc)
+	podmanAPISocksInHost, podmanAPISocksInGuest, err := startNetworking(ctx, mc)
 	if err != nil {
 		return fmt.Errorf("failed to start network provider: %w", err)
 	}
@@ -173,11 +204,11 @@ func tryKillHyperVisorBeforeRun(mc *vmconfig.MachineConfig) {
 	}
 }
 
-func startVMProvider(mc *vmconfig.MachineConfig) error {
+func startVMProvider(ctx context.Context, mc *vmconfig.MachineConfig) error {
 	tryKillHyperVisorBeforeRun(mc)
 	provider := mc.VMProvider
 	logrus.Infof("Start VM provider: %s", provider.VMType())
-	return provider.StartVM(mc) //nolint:wrapcheck
+	return provider.StartVM(ctx, mc) //nolint:wrapcheck
 }
 
 // Start the machine It will start network provider and hypervisor:
@@ -189,13 +220,13 @@ func startVMProvider(mc *vmconfig.MachineConfig) error {
 // Note: this function is a non-block function, it will return immediately after start the network provider and hypervisor
 func Start(ctx context.Context, mc *vmconfig.MachineConfig) (context.Context, error) {
 	// First start network provider which provided by gvproxy
-	err := startNetworkProvider(mc)
+	err := startNetworkProvider(ctx, mc)
 	if err != nil {
 		return nil, fmt.Errorf("failed to start network provider: %w", err)
 	}
 
 	// Start HyperVisor which provided by krunkit(arm64)/vfkit(x86_64)
-	err = startVMProvider(mc)
+	err = startVMProvider(ctx, mc)
 	if err != nil {
 		return nil, fmt.Errorf("failed to start vm provider: %w", err)
 	}