Hygiene round: context discipline on cleanup paths, comment budget, wg.Go

cmd/vm/clone.go

@@ -74,17 +74,17 @@ func (h *Handler) clone(cmd *cobra.Command, srcRec *record, name string) (retErr
 			return err
 		}
 	}
+	ctx := cliutil.CommandContext(cmd)
 	dir, overlay, ovmfVars, digest, err := scaffoldVM(cmd, name, srcRec.Image, srcRec.OVMFVars, filepath.Base(srcRec.OVMFVars))
 	if err != nil {
 		return err
 	}
 	var r *record
 	defer func() {
 		if retErr != nil {
-			retErr = errors.Join(retErr, cleanupFailedVM(cmd, dir, r))
+			retErr = errors.Join(retErr, cleanupFailedVM(ctx, cmd, dir, r))
 		}
 	}()
-	ctx := cliutil.CommandContext(cmd)
 	storage, err = resizeSystemDisk(ctx, overlay, storage)
 	if err != nil {
 		return err

cmd/vm/lifecycle.go

@@ -125,12 +125,8 @@ func (h *Handler) RM(cmd *cobra.Command, args []string) error {
 				if err := teardownNet(ctx, cmd, r); err != nil {
 					return err
 				}
-			} else {
-				cleanupCtx, cancel := context.WithTimeout(context.Background(), vmCleanupTimeout)
-				defer cancel()
-				if cleanupErr := reapStrayHelpers(cleanupCtx, dir); cleanupErr != nil {
-					return cleanupErr
-				}
+			} else if cleanupErr := reapStrayHelpers(ctx, dir); cleanupErr != nil {
+				return cleanupErr
 			}
 			if err := os.RemoveAll(dir); err != nil {
 				return fmt.Errorf("remove vm dir: %w", err)
@@ -150,7 +146,8 @@ func (h *Handler) createVM(cmd *cobra.Command, image string, launch bool) error
 	if err != nil {
 		return err
 	}
-	return withVMLock(cliutil.CommandContext(cmd), dir, func() error {
+	ctx := cliutil.CommandContext(cmd)
+	return withVMLock(ctx, dir, func() error {
 		r, err := h.create(cmd, image, name)
 		if err != nil {
 			return err
@@ -160,7 +157,7 @@ func (h *Handler) createVM(cmd *cobra.Command, image string, launch bool) error
 			return nil
 		}
 		if err := h.launch(cmd, dir, r); err != nil {
-			return errors.Join(err, cleanupFailedVM(cmd, dir, r))
+			return errors.Join(err, cleanupFailedVM(ctx, cmd, dir, r))
 		}
 		fmt.Printf("%s (pid %d)\n", r.Name, r.PID)
 		return nil
@@ -191,16 +188,16 @@ func (h *Handler) create(cmd *cobra.Command, image, name string) (r *record, ret
 	if err != nil {
 		return nil, err
 	}
+	ctx := cliutil.CommandContext(cmd)
 	dir, overlay, ovmfVars, digest, err := scaffoldVM(cmd, name, image, varsTmpl, "OVMF_VARS.fd")
 	if err != nil {
 		return nil, err
 	}
 	defer func() {
 		if retErr != nil {
-			retErr = errors.Join(retErr, cleanupFailedVM(cmd, dir, r))
+			retErr = errors.Join(retErr, cleanupFailedVM(ctx, cmd, dir, r))
 		}
 	}()
-	ctx := cliutil.CommandContext(cmd)
 	storage, err = resizeSystemDisk(ctx, overlay, storage)
 	if err != nil {
 		return nil, err
@@ -258,7 +255,7 @@ func (h *Handler) launch(cmd *cobra.Command, dir string, r *record) error {
 	pidfile := filepath.Join(dir, "qemu.pid")
 	args := append(spec.Args(), "-daemonize", "-pidfile", pidfile)
 	stopVNCProxy(ctx, dir)
-	ensureNetnsLoopback(ctx, r) // CNI: a fresh netns has lo DOWN, so qemu's -vnc 127.0.0.1 would fail to bind
+	ensureNetnsLoopback(ctx, r)
 	if r.Netns != "" {
 		logger.Debugf(ctx, "running qemu in netns %s via `ip netns exec`", filepath.Base(r.Netns))
 	}
@@ -267,7 +264,7 @@ func (h *Handler) launch(cmd *cobra.Command, dir string, r *record) error {
 	if err := saveRec(dir, r); err != nil {
 		return err
 	}
-	c := launchCmd(r, args) // CNI: wraps in `ip netns exec` so -netdev tap finds the in-netns TAP
+	c := launchCmd(r, args)
 	c.Stdout, c.Stderr = os.Stdout, os.Stderr
 	if err := c.Run(); err != nil {
 		return fmt.Errorf("launch qemu: %w", err)
@@ -312,8 +309,8 @@ func validateMacOSCPUs(cpus int) error {
 }
 
 // cleanupFailedVM uses an uncanceled bounded context so cancellation still reaps helpers, networking and QEMU.
-func cleanupFailedVM(cmd *cobra.Command, dir string, r *record) error {
-	ctx, cancel := context.WithTimeout(context.Background(), vmCleanupTimeout)
+func cleanupFailedVM(ctx context.Context, cmd *cobra.Command, dir string, r *record) error {
+	ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), vmCleanupTimeout)
 	defer cancel()
 	var errs []error
 	if r != nil {

cmd/vm/utils.go

@@ -55,7 +55,7 @@ func withVMLock(ctx context.Context, dir string, fn func() error) error {
 	if err := l.Lock(ctx); err != nil {
 		return fmt.Errorf("lock vm: %w", err)
 	}
-	defer func() { _ = l.Unlock(context.Background()) }()
+	defer func() { _ = l.Unlock(context.WithoutCancel(ctx)) }()
 	return fn()
 }
 
@@ -141,12 +141,12 @@ func scaffoldVM(cmd *cobra.Command, name, image, varsSrc, varsName string) (dir,
 	} else if !os.IsNotExist(statErr) {
 		return "", "", "", "", fmt.Errorf("stat vm record: %w", statErr)
 	}
-	cleanupCtx, cancel := context.WithTimeout(context.Background(), vmCleanupTimeout)
+	ctx := cliutil.CommandContext(cmd)
+	cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), vmCleanupTimeout)
 	defer cancel()
 	if err = resetIncompleteVMDir(cleanupCtx, dir); err != nil {
 		return "", "", "", "", err
 	}
-	ctx := cliutil.CommandContext(cmd)
 	base, digest, err := resolveBase(ctx, cmd, image, name)
 	if err != nil {
 		return "", "", "", "", err

cmd/vm/utils_test.go

@@ -1,6 +1,7 @@
 package vm
 
 import (
+	"context"
 	"net"
 	"os"
 	"path/filepath"
@@ -53,32 +54,28 @@ func TestWithVMLockSurvivesVMDirRemoval(t *testing.T) {
 	acquired := make(chan struct{})
 	release := make(chan struct{})
 	var wg sync.WaitGroup
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
+	wg.Go(func() {
 		if err := withVMLock(t.Context(), dir, func() error {
 			close(acquired)
 			<-release
 			return nil
 		}); err != nil {
 			t.Errorf("first lock: %v", err)
 		}
-	}()
+	})
 	<-acquired
 	if err := os.RemoveAll(dir); err != nil {
 		t.Fatal(err)
 	}
 	second := make(chan struct{})
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
+	wg.Go(func() {
 		if err := withVMLock(t.Context(), dir, func() error {
 			close(second)
 			return nil
 		}); err != nil {
 			t.Errorf("second lock: %v", err)
 		}
-	}()
+	})
 	select {
 	case <-second:
 		t.Fatal("second operation acquired while first still held the VM lock")
@@ -178,6 +175,29 @@ func TestHMPRepliedFlagsAnyMessage(t *testing.T) {
 	}
 }
 
+func TestWithVMLockUnlocksAfterCallerCancel(t *testing.T) {
+	dir := filepath.Join(t.TempDir(), "vms", "demo")
+	ctx, cancel := context.WithCancel(t.Context())
+	if err := withVMLock(ctx, dir, func() error {
+		cancel()
+		return nil
+	}); err != nil {
+		t.Fatalf("first lock: %v", err)
+	}
+	relocked := make(chan error, 1)
+	go func() {
+		relocked <- withVMLock(t.Context(), dir, func() error { return nil })
+	}()
+	select {
+	case err := <-relocked:
+		if err != nil {
+			t.Fatalf("second lock: %v", err)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("second lock blocked: the cancelled caller context suppressed Unlock")
+	}
+}
+
 func hmpTranscript(chunks ...string) (string, bool) {
 	client, monitor := net.Pipe()
 	go func() {

qemu/inject.go

@@ -34,24 +34,25 @@ type nbdConnection struct {
 	ocPath  string
 }
 
-// fresh context, not the caller's: disconnect must still run when the caller is unwinding after a timeout.
-func (c *nbdConnection) disconnect() error {
+// detached from the caller's cancellation: disconnect must still run when the caller is unwinding after a timeout.
+func (c *nbdConnection) disconnect(ctx context.Context) error {
 	logger := log.WithFunc("qemu.nbdConnection.disconnect")
+	ctx = context.WithoutCancel(ctx)
 	var commandErr error
-	disconnectCtx, disconnectCancel := context.WithTimeout(context.Background(), nbdCommandTimeout)
+	disconnectCtx, disconnectCancel := context.WithTimeout(ctx, nbdCommandTimeout)
 	if out, err := exec.CommandContext(disconnectCtx, "qemu-nbd", "--disconnect", c.device).CombinedOutput(); err != nil {
 		logger.Warnf(disconnectCtx, "disconnect %s (output: %s): %v", c.device, out, err)
 		commandErr = fmt.Errorf("disconnect %s (output: %s): %w", c.device, out, err)
 	}
 	disconnectCancel()
-	waitCtx, waitCancel := context.WithTimeout(context.Background(), nbdCleanupTimeout)
+	waitCtx, waitCancel := context.WithTimeout(ctx, nbdCleanupTimeout)
 	defer waitCancel()
 	if err := utils.WaitFor(waitCtx, 10*time.Second, 100*time.Millisecond, func() (bool, error) {
 		held, scanErr := isFileHeld(c.ocPath)
 		return !held, scanErr
 	}); err != nil {
 		logger.Warnf(waitCtx, "qcow2 %s still held after nbd disconnect: %v", c.ocPath, err)
-		killCtx, killCancel := context.WithTimeout(context.Background(), nbdCleanupTimeout)
+		killCtx, killCancel := context.WithTimeout(ctx, nbdCleanupTimeout)
 		defer killCancel()
 		var cleanupErrs []error
 		if cleanupErr := utils.TerminateProcess(killCtx, c.pid, "qemu-nbd", c.ocPath, time.Second); cleanupErr != nil {
@@ -83,7 +84,7 @@ func InjectConfig(ctx context.Context, ocPath string, sm *SMBIOS) error {
 		if connectErr != nil {
 			return connectErr
 		}
-		defer func() { retErr = errors.Join(retErr, conn.disconnect()) }()
+		defer func() { retErr = errors.Join(retErr, conn.disconnect(ctx)) }()
 		if waitErr := waitForPart(ctx, conn.device); waitErr != nil {
 			return waitErr
 		}
@@ -92,7 +93,7 @@ func InjectConfig(ctx context.Context, ocPath string, sm *SMBIOS) error {
 			return fmt.Errorf("create mount dir: %w", err)
 		}
 		mounted := false
-		defer func() { retErr = errors.Join(retErr, cleanupNBDMount(mnt, mounted)) }()
+		defer func() { retErr = errors.Join(retErr, cleanupNBDMount(ctx, mnt, mounted)) }()
 		var mountErr error
 		for _, p := range []string{conn.device + "p1", conn.device + "p2", conn.device} {
 			if mountErr = exec.CommandContext(ctx, "mount", p, mnt).Run(); mountErr == nil {
@@ -113,7 +114,7 @@ func withNBDLease(ctx context.Context, path string, fn func() error) error {
 	if err := l.Lock(ctx); err != nil {
 		return fmt.Errorf("lock qemu-nbd lease: %w", err)
 	}
-	defer func() { _ = l.Unlock(context.Background()) }()
+	defer func() { _ = l.Unlock(context.WithoutCancel(ctx)) }()
 	return fn()
 }
 
@@ -130,14 +131,14 @@ func waitForPart(ctx context.Context, nbd string) error {
 	return nil
 }
 
-func cleanupNBDMount(mnt string, mounted bool) error {
+func cleanupNBDMount(ctx context.Context, mnt string, mounted bool) error {
 	if !mounted {
 		if err := os.Remove(mnt); err != nil && !os.IsNotExist(err) {
 			return fmt.Errorf("remove mount dir %s: %w", mnt, err)
 		}
 		return nil
 	}
-	ctx, cancel := context.WithTimeout(context.Background(), nbdCleanupTimeout)
+	ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), nbdCleanupTimeout)
 	defer cancel()
 	if out, err := exec.CommandContext(ctx, "umount", mnt).CombinedOutput(); err != nil {
 		return fmt.Errorf("unmount %s (output: %s): %w", mnt, out, err)
@@ -185,14 +186,14 @@ func connectFreeNBD(ctx context.Context, ocPath string) (*nbdConnection, error)
 			if err == nil {
 				return &nbdConnection{device: nbd, pid: pid, pidFile: pidFile, ocPath: ocPath}, nil
 			}
-			cleanupCtx, cancel := context.WithTimeout(context.Background(), nbdCleanupTimeout)
+			cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), nbdCleanupTimeout)
 			_ = exec.CommandContext(cleanupCtx, "qemu-nbd", "--disconnect", nbd).Run()
 			_ = cleanupNBDForPath(cleanupCtx, ocPath)
 			cancel()
 			return nil, fmt.Errorf("read qemu-nbd pid file %s: %w", pidFile, err)
 		}
 		lastErr = fmt.Errorf("connect qemu-nbd %s (output: %s): %w", nbd, out, cerr)
-		cleanupCtx, cancel := context.WithTimeout(context.Background(), nbdCleanupTimeout)
+		cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), nbdCleanupTimeout)
 		_ = cleanupNBDForPath(cleanupCtx, ocPath)
 		cancel()
 		if ctx.Err() != nil {

qemu/inject_test.go

@@ -67,9 +67,7 @@ func TestWithNBDLeaseSerializesConcurrentInjectors(t *testing.T) {
 	errCh := make(chan error, workers)
 	var wg sync.WaitGroup
 	for range workers {
-		wg.Add(1)
-		go func() {
-			defer wg.Done()
+		wg.Go(func() {
 			errCh <- withNBDLease(t.Context(), lockPath, func() error {
 				n := active.Add(1)
 				defer active.Add(-1)
@@ -78,7 +76,7 @@ func TestWithNBDLeaseSerializesConcurrentInjectors(t *testing.T) {
 				time.Sleep(time.Millisecond)
 				return nil
 			})
-		}()
+		})
 	}
 	wg.Wait()
 	close(errCh)