refactor(talosctl): propagate command context throughout, handle interrupts

Currently, WithClientMaintenance[1] and WithClientNoNodes[2] will spawn a new context.Background on each invocation. This effectively severs the chain of context, with respect to the the root/base Cobra command context [3]. In some cases, this was known to break handling of SIGINT (Ctrl^C), as multiple such signals were required to terminate the running CLI (one SIGINT for every separate context).

What I've done to solve this:

* cmd/talosctl/cmd/root.go provides a root-level Ctrl^C cancelable context to all of the talosctl commands.
* Every talosctl command now reuses the root context, which is propagated through the Cobra cmd.Context().
* No talosctl command creates another SIGINT handler, like they did before on an individual basis.
* Termination of processes launched via *-launch is also improved. They'll attempt a graceful shutdown as soon as cmd.Context().Done() is closed. This includes dhcpd, dnsd, loadbalancer, kms, and tftpd.

Signed-off-by: Maja Bojarska <maja.bojarska@siderolabs.com>

Author: majabojarska

cmd/installer/cmd/imager/root.go

@@ -57,174 +57,174 @@ var rootCmd = &cobra.Command{
 	Args:         cobra.ExactArgs(1),
 	SilenceUsage: true,
 	RunE: func(cmd *cobra.Command, args []string) error {
-		return cli.WithContext(context.Background(), func(ctx context.Context) error {
-			report := reporter.New()
-			report.Report(reporter.Update{
-				Message: "assembling the finalized profile...",
-				Status:  reporter.StatusRunning,
-			})
-
-			baseProfile := args[0]
-
-			var prof profile.Profile
+		ctx := cmd.Context()
 
-			if baseProfile == "-" {
-				if err := yaml.NewDecoder(os.Stdin).Decode(&prof); err != nil {
-					return err
-				}
-			} else {
-				prof = profile.Profile{
-					BaseProfileName: baseProfile,
-					Arch:            cmdFlags.Arch,
-					Platform:        cmdFlags.Platform,
-					Customization: profile.CustomizationProfile{
-						ExtraKernelArgs: cmdFlags.ExtraKernelArgs,
-						MetaContents:    cmdFlags.MetaValues.GetMetaValues(),
-					},
-				}
-
-				extraOverlayOptions := overlay.ExtraOptions{}
-
-				for _, option := range cmdFlags.OverlayOptions {
-					if strings.HasPrefix(option, "@") {
-						data, err := os.ReadFile(option[1:])
-						if err != nil {
-							return err
-						}
-
-						decoder := yaml.NewDecoder(bytes.NewReader(data))
-						decoder.KnownFields(true)
+		report := reporter.New()
+		report.Report(reporter.Update{
+			Message: "assembling the finalized profile...",
+			Status:  reporter.StatusRunning,
+		})
 
-						if err := decoder.Decode(&extraOverlayOptions); err != nil {
-							return err
-						}
+		baseProfile := args[0]
 
-						continue
-					}
+		var prof profile.Profile
 
-					k, v, _ := strings.Cut(option, "=")
+		if baseProfile == "-" {
+			if err := yaml.NewDecoder(os.Stdin).Decode(&prof); err != nil {
+				return err
+			}
+		} else {
+			prof = profile.Profile{
+				BaseProfileName: baseProfile,
+				Arch:            cmdFlags.Arch,
+				Platform:        cmdFlags.Platform,
+				Customization: profile.CustomizationProfile{
+					ExtraKernelArgs: cmdFlags.ExtraKernelArgs,
+					MetaContents:    cmdFlags.MetaValues.GetMetaValues(),
+				},
+			}
 
-					if strings.HasPrefix(v, "@") {
-						data, err := os.ReadFile(v[1:])
-						if err != nil {
-							return err
-						}
+			extraOverlayOptions := overlay.ExtraOptions{}
 
-						v = string(data)
+			for _, option := range cmdFlags.OverlayOptions {
+				if strings.HasPrefix(option, "@") {
+					data, err := os.ReadFile(option[1:])
+					if err != nil {
+						return err
 					}
 
-					extraOverlayOptions[k] = v
-				}
+					decoder := yaml.NewDecoder(bytes.NewReader(data))
+					decoder.KnownFields(true)
 
-				if cmdFlags.OverlayName != "" || cmdFlags.OverlayImage != "" {
-					prof.Overlay = &profile.OverlayOptions{
-						Name: cmdFlags.OverlayName,
-						Image: profile.ContainerAsset{
-							ImageRef: cmdFlags.OverlayImage,
-						},
-						ExtraOptions: extraOverlayOptions,
+					if err := decoder.Decode(&extraOverlayOptions); err != nil {
+						return err
 					}
 
-					prof.Input.OverlayInstaller.ImageRef = cmdFlags.OverlayImage
+					continue
 				}
 
-				prof.Input.SystemExtensions = xslices.Map(
-					cmdFlags.SystemExtensionImages,
-					func(imageRef string) profile.ContainerAsset {
-						return profile.ContainerAsset{
-							ImageRef:      imageRef,
-							ForceInsecure: cmdFlags.Insecure,
-						}
-					},
-				)
+				k, v, _ := strings.Cut(option, "=")
 
-				if cmdFlags.OutputKind != "" {
-					outKind, err := profile.OutputKindString(cmdFlags.OutputKind)
+				if strings.HasPrefix(v, "@") {
+					data, err := os.ReadFile(v[1:])
 					if err != nil {
 						return err
 					}
 
-					prof.Output.Kind = outKind
+					v = string(data)
 				}
 
-				if cmdFlags.BaseInstallerImage != "" {
-					prof.Input.BaseInstaller = profile.ContainerAsset{
-						ImageRef: cmdFlags.BaseInstallerImage,
-					}
+				extraOverlayOptions[k] = v
+			}
+
+			if cmdFlags.OverlayName != "" || cmdFlags.OverlayImage != "" {
+				prof.Overlay = &profile.OverlayOptions{
+					Name: cmdFlags.OverlayName,
+					Image: profile.ContainerAsset{
+						ImageRef: cmdFlags.OverlayImage,
+					},
+					ExtraOptions: extraOverlayOptions,
 				}
 
-				if cmdFlags.ImageCache != "" {
-					parseOpts := []name.Option{name.StrictValidation}
+				prof.Input.OverlayInstaller.ImageRef = cmdFlags.OverlayImage
+			}
 
-					if cmdFlags.Insecure {
-						parseOpts = append(parseOpts, name.Insecure)
+			prof.Input.SystemExtensions = xslices.Map(
+				cmdFlags.SystemExtensionImages,
+				func(imageRef string) profile.ContainerAsset {
+					return profile.ContainerAsset{
+						ImageRef:      imageRef,
+						ForceInsecure: cmdFlags.Insecure,
 					}
+				},
+			)
 
-					if _, err := name.ParseReference(cmdFlags.ImageCache, parseOpts...); err == nil {
-						prof.Input.ImageCache = profile.ContainerAsset{
-							ImageRef: cmdFlags.ImageCache,
-						}
-					} else {
-						prof.Input.ImageCache = profile.ContainerAsset{
-							OCIPath: cmdFlags.ImageCache,
-						}
-					}
+			if cmdFlags.OutputKind != "" {
+				outKind, err := profile.OutputKindString(cmdFlags.OutputKind)
+				if err != nil {
+					return err
 				}
 
-				if cmdFlags.Insecure {
-					prof.Input.BaseInstaller.ForceInsecure = cmdFlags.Insecure
-					prof.Input.ImageCache.ForceInsecure = cmdFlags.Insecure
+				prof.Output.Kind = outKind
+			}
+
+			if cmdFlags.BaseInstallerImage != "" {
+				prof.Input.BaseInstaller = profile.ContainerAsset{
+					ImageRef: cmdFlags.BaseInstallerImage,
 				}
+			}
 
-				if cmdFlags.SecurebootIncludeWellKnownCerts {
-					if prof.Input.SecureBoot == nil {
-						prof.Input.SecureBoot = &profile.SecureBootAssets{}
-					}
+			if cmdFlags.ImageCache != "" {
+				parseOpts := []name.Option{name.StrictValidation}
 
-					prof.Input.SecureBoot.IncludeWellKnownCerts = true
+				if cmdFlags.Insecure {
+					parseOpts = append(parseOpts, name.Insecure)
 				}
 
-				if cmdFlags.EmbeddedConfigPath != "" {
-					data, err := os.ReadFile(cmdFlags.EmbeddedConfigPath)
-					if err != nil {
-						return fmt.Errorf("error reading embedded config file: %w", err)
+				if _, err := name.ParseReference(cmdFlags.ImageCache, parseOpts...); err == nil {
+					prof.Input.ImageCache = profile.ContainerAsset{
+						ImageRef: cmdFlags.ImageCache,
+					}
+				} else {
+					prof.Input.ImageCache = profile.ContainerAsset{
+						OCIPath: cmdFlags.ImageCache,
 					}
-
-					prof.Customization.EmbeddedMachineConfiguration = string(data)
 				}
 			}
 
-			if err := os.MkdirAll(cmdFlags.OutputPath, 0o755); err != nil {
-				return err
+			if cmdFlags.Insecure {
+				prof.Input.BaseInstaller.ForceInsecure = cmdFlags.Insecure
+				prof.Input.ImageCache.ForceInsecure = cmdFlags.Insecure
 			}
 
-			imager, err := imager.New(prof)
-			if err != nil {
-				return err
+			if cmdFlags.SecurebootIncludeWellKnownCerts {
+				if prof.Input.SecureBoot == nil {
+					prof.Input.SecureBoot = &profile.SecureBootAssets{}
+				}
+
+				prof.Input.SecureBoot.IncludeWellKnownCerts = true
 			}
 
-			if _, err = imager.Execute(ctx, cmdFlags.OutputPath, report); err != nil {
-				report.Report(reporter.Update{
-					Message: err.Error(),
-					Status:  reporter.StatusError,
-				})
+			if cmdFlags.EmbeddedConfigPath != "" {
+				data, err := os.ReadFile(cmdFlags.EmbeddedConfigPath)
+				if err != nil {
+					return fmt.Errorf("error reading embedded config file: %w", err)
+				}
 
-				return err
+				prof.Customization.EmbeddedMachineConfiguration = string(data)
 			}
+		}
 
-			if cmdFlags.TarToStdout {
-				return archiver.TarGz(ctx, cmdFlags.OutputPath, os.Stdout)
-			}
+		if err := os.MkdirAll(cmdFlags.OutputPath, 0o755); err != nil {
+			return err
+		}
 
-			return nil
-		})
+		imager, err := imager.New(prof)
+		if err != nil {
+			return err
+		}
+
+		if _, err = imager.Execute(ctx, cmdFlags.OutputPath, report); err != nil {
+			report.Report(reporter.Update{
+				Message: err.Error(),
+				Status:  reporter.StatusError,
+			})
+
+			return err
+		}
+
+		if cmdFlags.TarToStdout {
+			return archiver.TarGz(ctx, cmdFlags.OutputPath, os.Stdout)
+		}
+
+		return nil
 	},
 }
 
 // Execute adds all child commands to the root command and sets flags appropriately.
 // This is called by main.main(). It only needs to happen once to the rootCmd.
 func Execute() {
-	if err := rootCmd.Execute(); err != nil {
+	if _, err := cli.WithContextC(context.Background(), rootCmd.ExecuteContextC); err != nil {
 		os.Exit(1)
 	}
 }

cmd/talosctl/cmd/mgmt/cluster/create/cmd_dev.go

@@ -5,7 +5,6 @@
 package create
 
 import (
-	"context"
 	"fmt"
 	"runtime"
 	"slices"
@@ -268,55 +267,53 @@ func getCreateCmd(cmdName string, hidden bool) *cobra.Command {
 		Short: "Creates a local QEMU-based cluster for Talos development.",
 		Args:  cobra.NoArgs,
 		RunE: func(cmd *cobra.Command, args []string) error {
-			return cli.WithContext(context.Background(), func(ctx context.Context) error {
-				if cmdName == "create" {
-					cli.Warning("the developer oriented 'cluster create' command has been moved to 'cluster create dev'")
-				}
+			if cmdName == "create" {
+				cli.Warning("the developer oriented 'cluster create' command has been moved to 'cluster create dev'")
+			}
 
-				if err := validateQemuFlags(cmd.Flags(), unImplementedFlagsDarwin); err != nil {
-					return err
-				}
+			if err := validateQemuFlags(cmd.Flags(), unImplementedFlagsDarwin); err != nil {
+				return err
+			}
 
-				var disks strings.Builder
-				fmt.Fprintf(&disks, "virtio:%d", legacyOps.clusterDiskSize)
+			var disks strings.Builder
+			fmt.Fprintf(&disks, "virtio:%d", legacyOps.clusterDiskSize)
 
-				for i := range legacyOps.extraDisks {
-					driver := "ide"
-					tag := ""
-					serial := ""
+			for i := range legacyOps.extraDisks {
+				driver := "ide"
+				tag := ""
+				serial := ""
 
-					// ide driver is not supported on arm64
-					if qOps.TargetArch == "arm64" {
-						driver = "virtio"
-					}
+				// ide driver is not supported on arm64
+				if qOps.TargetArch == "arm64" {
+					driver = "virtio"
+				}
 
-					if i < len(legacyOps.extraDisksDrivers) {
-						driver = legacyOps.extraDisksDrivers[i]
-					}
+				if i < len(legacyOps.extraDisksDrivers) {
+					driver = legacyOps.extraDisksDrivers[i]
+				}
 
-					if i < len(legacyOps.extraDisksTags) {
-						if legacyOps.extraDisksTags[i] != "" {
-							tag = fmt.Sprintf(":tag=%s", legacyOps.extraDisksTags[i])
-						}
+				if i < len(legacyOps.extraDisksTags) {
+					if legacyOps.extraDisksTags[i] != "" {
+						tag = fmt.Sprintf(":tag=%s", legacyOps.extraDisksTags[i])
 					}
+				}
 
-					if i < len(legacyOps.extraDisksSerials) {
-						if legacyOps.extraDisksSerials[i] != "" {
-							serial = fmt.Sprintf(":serial=%s", legacyOps.extraDisksSerials[i])
-						}
+				if i < len(legacyOps.extraDisksSerials) {
+					if legacyOps.extraDisksSerials[i] != "" {
+						serial = fmt.Sprintf(":serial=%s", legacyOps.extraDisksSerials[i])
 					}
-
-					fmt.Fprintf(&disks, ",%s:%d%s%s", driver, legacyOps.extraDiskSize, tag, serial)
 				}
 
-				qOps.Disks = flags.Disks{}
+				fmt.Fprintf(&disks, ",%s:%d%s%s", driver, legacyOps.extraDiskSize, tag, serial)
+			}
 
-				if err := qOps.Disks.Set(disks.String()); err != nil {
-					return err
-				}
+			qOps.Disks = flags.Disks{}
+
+			if err := qOps.Disks.Set(disks.String()); err != nil {
+				return err
+			}
 
-				return createDevCluster(ctx, cOps, qOps)
-			})
+			return createDevCluster(cmd.Context(), cOps, qOps)
 		},
 	}
 	createCmd.Flags().IntVar(&legacyOps.clusterDiskSize, clusterDiskSizeFlag, 6*1024, "default limit on disk size in MB (each VM)")