cobra/commands: handle signals when standalone

See docker/cli#4599 and
docker/cli#4769.

When running through the CLI, signal handling + the "3
SIGINTS = exit" behavior is handled by the CLI, and the
CLI will signal buildx through the plugin socket to cancel
it's context.

To deal with the standalone case, this commit introduces
`cobrautil.HandleContextCancellation` which checks if the
context being executed is already cancellable/plugin is
running through the CLI and sets up signal handling if
necessary.

Signed-off-by: Laura Brehm <laurabrehm@hey.com>

commands/bake.go

@@ -15,6 +15,7 @@ import (
 	"github.com/docker/buildx/builder"
 	"github.com/docker/buildx/localstate"
 	"github.com/docker/buildx/util/buildflags"
+	"github.com/docker/buildx/util/cobrautil"
 	"github.com/docker/buildx/util/cobrautil/completion"
 	"github.com/docker/buildx/util/confutil"
 	"github.com/docker/buildx/util/desktop"
@@ -252,7 +253,7 @@ func bakeCmd(dockerCli command.Cli, rootOpts *rootOptions) *cobra.Command {
 		Use:     "bake [OPTIONS] [TARGET...]",
 		Aliases: []string{"f"},
 		Short:   "Build from a file",
-		RunE: func(cmd *cobra.Command, args []string) error {
+		RunE: cobrautil.HandleContextCancellation(func(cmd *cobra.Command, args []string) error {
 			// reset to nil to avoid override is unset
 			if !cmd.Flags().Lookup("no-cache").Changed {
 				cFlags.noCache = nil
@@ -264,7 +265,7 @@ func bakeCmd(dockerCli command.Cli, rootOpts *rootOptions) *cobra.Command {
 			options.metadataFile = cFlags.metadataFile
 			// Other common flags (noCache, pull and progress) are processed in runBake function.
 			return runBake(cmd.Context(), dockerCli, args, options, cFlags)
-		},
+		}),
 		ValidArgsFunction: completion.BakeTargets(options.files),
 	}
 

commands/build.go

@@ -27,6 +27,7 @@ import (
 	"github.com/docker/buildx/store"
 	"github.com/docker/buildx/store/storeutil"
 	"github.com/docker/buildx/util/buildflags"
+	"github.com/docker/buildx/util/cobrautil"
 	"github.com/docker/buildx/util/desktop"
 	"github.com/docker/buildx/util/ioset"
 	"github.com/docker/buildx/util/progress"
@@ -447,7 +448,7 @@ func buildCmd(dockerCli command.Cli, rootOpts *rootOptions, debugConfig *debug.D
 		Aliases: []string{"b"},
 		Short:   "Start a build",
 		Args:    cli.ExactArgs(1),
-		RunE: func(cmd *cobra.Command, args []string) error {
+		RunE: cobrautil.HandleContextCancellation(func(cmd *cobra.Command, args []string) error {
 			options.contextPath = args[0]
 			options.builder = rootOpts.builder
 			options.metadataFile = cFlags.metadataFile
@@ -471,7 +472,7 @@ func buildCmd(dockerCli command.Cli, rootOpts *rootOptions, debugConfig *debug.D
 			}
 
 			return runBuild(cmd.Context(), dockerCli, *options)
-		},
+		}),
 		ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
 			return nil, cobra.ShellCompDirectiveFilterDirs
 		},

commands/create.go

@@ -95,9 +95,9 @@ func createCmd(dockerCli command.Cli) *cobra.Command {
 		Use:   "create [OPTIONS] [CONTEXT|ENDPOINT]",
 		Short: "Create a new builder instance",
 		Args:  cli.RequiresMaxArgs(1),
-		RunE: func(cmd *cobra.Command, args []string) error {
+		RunE: cobrautil.HandleContextCancellation(func(cmd *cobra.Command, args []string) error {
 			return runCreate(cmd.Context(), dockerCli, options, args)
-		},
+		}),
 		ValidArgsFunction: completion.Disable,
 	}
 

commands/diskusage.go

@@ -10,6 +10,7 @@ import (
 	"time"
 
 	"github.com/docker/buildx/builder"
+	"github.com/docker/buildx/util/cobrautil"
 	"github.com/docker/buildx/util/cobrautil/completion"
 	"github.com/docker/cli/cli"
 	"github.com/docker/cli/cli/command"
@@ -110,10 +111,10 @@ func duCmd(dockerCli command.Cli, rootOpts *rootOptions) *cobra.Command {
 		Use:   "du",
 		Short: "Disk usage",
 		Args:  cli.NoArgs,
-		RunE: func(cmd *cobra.Command, args []string) error {
+		RunE: cobrautil.HandleContextCancellation(func(cmd *cobra.Command, args []string) error {
 			options.builder = rootOpts.builder
 			return runDiskUsage(cmd.Context(), dockerCli, options)
-		},
+		}),
 		ValidArgsFunction: completion.Disable,
 	}
 

commands/imagetools/create.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/distribution/reference"
 	"github.com/docker/buildx/builder"
+	"github.com/docker/buildx/util/cobrautil"
 	"github.com/docker/buildx/util/cobrautil/completion"
 	"github.com/docker/buildx/util/imagetools"
 	"github.com/docker/buildx/util/progress"
@@ -269,10 +270,10 @@ func createCmd(dockerCli command.Cli, opts RootOptions) *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "create [OPTIONS] [SOURCE] [SOURCE...]",
 		Short: "Create a new image based on source images",
-		RunE: func(cmd *cobra.Command, args []string) error {
+		RunE: cobrautil.HandleContextCancellation(func(cmd *cobra.Command, args []string) error {
 			options.builder = *opts.Builder
 			return runCreate(cmd.Context(), dockerCli, options, args)
-		},
+		}),
 		ValidArgsFunction: completion.Disable,
 	}
 

commands/imagetools/inspect.go

@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"github.com/docker/buildx/builder"
+	"github.com/docker/buildx/util/cobrautil"
 	"github.com/docker/buildx/util/cobrautil/completion"
 	"github.com/docker/buildx/util/imagetools"
 	"github.com/docker/cli-docs-tool/annotation"
@@ -48,10 +49,10 @@ func inspectCmd(dockerCli command.Cli, rootOpts RootOptions) *cobra.Command {
 		Use:   "inspect [OPTIONS] NAME",
 		Short: "Show details of an image in the registry",
 		Args:  cli.ExactArgs(1),
-		RunE: func(cmd *cobra.Command, args []string) error {
+		RunE: cobrautil.HandleContextCancellation(func(cmd *cobra.Command, args []string) error {
 			options.builder = *rootOpts.Builder
 			return runInspect(cmd.Context(), dockerCli, options, args[0])
-		},
+		}),
 		ValidArgsFunction: completion.Disable,
 	}
 

commands/inspect.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/docker/buildx/builder"
 	"github.com/docker/buildx/driver"
+	"github.com/docker/buildx/util/cobrautil"
 	"github.com/docker/buildx/util/cobrautil/completion"
 	"github.com/docker/buildx/util/platformutil"
 	"github.com/docker/cli/cli"
@@ -142,13 +143,13 @@ func inspectCmd(dockerCli command.Cli, rootOpts *rootOptions) *cobra.Command {
 		Use:   "inspect [NAME]",
 		Short: "Inspect current builder instance",
 		Args:  cli.RequiresMaxArgs(1),
-		RunE: func(cmd *cobra.Command, args []string) error {
+		RunE: cobrautil.HandleContextCancellation(func(cmd *cobra.Command, args []string) error {
 			options.builder = rootOpts.builder
 			if len(args) > 0 {
 				options.builder = args[0]
 			}
 			return runInspect(cmd.Context(), dockerCli, options)
-		},
+		}),
 		ValidArgsFunction: completion.BuilderNames(dockerCli),
 	}
 

commands/ls.go

@@ -99,9 +99,9 @@ func lsCmd(dockerCli command.Cli) *cobra.Command {
 		Use:   "ls",
 		Short: "List builder instances",
 		Args:  cli.ExactArgs(0),
-		RunE: func(cmd *cobra.Command, args []string) error {
+		RunE: cobrautil.HandleContextCancellation(func(cmd *cobra.Command, args []string) error {
 			return runLs(cmd.Context(), dockerCli, options)
-		},
+		}),
 		ValidArgsFunction: completion.Disable,
 	}
 

commands/prune.go

@@ -9,6 +9,7 @@ import (
 	"time"
 
 	"github.com/docker/buildx/builder"
+	"github.com/docker/buildx/util/cobrautil"
 	"github.com/docker/buildx/util/cobrautil/completion"
 	"github.com/docker/cli/cli"
 	"github.com/docker/cli/cli/command"
@@ -134,10 +135,10 @@ func pruneCmd(dockerCli command.Cli, rootOpts *rootOptions) *cobra.Command {
 		Use:   "prune",
 		Short: "Remove build cache",
 		Args:  cli.NoArgs,
-		RunE: func(cmd *cobra.Command, args []string) error {
+		RunE: cobrautil.HandleContextCancellation(func(cmd *cobra.Command, args []string) error {
 			options.builder = rootOpts.builder
 			return runPrune(cmd.Context(), dockerCli, options)
-		},
+		}),
 		ValidArgsFunction: completion.Disable,
 	}
 

commands/rm.go

@@ -8,6 +8,7 @@ import (
 	"github.com/docker/buildx/builder"
 	"github.com/docker/buildx/store"
 	"github.com/docker/buildx/store/storeutil"
+	"github.com/docker/buildx/util/cobrautil"
 	"github.com/docker/buildx/util/cobrautil/completion"
 	"github.com/docker/cli/cli"
 	"github.com/docker/cli/cli/command"
@@ -80,7 +81,7 @@ func rmCmd(dockerCli command.Cli, rootOpts *rootOptions) *cobra.Command {
 		Use:   "rm [NAME]",
 		Short: "Remove a builder instance",
 		Args:  cli.RequiresMaxArgs(1),
-		RunE: func(cmd *cobra.Command, args []string) error {
+		RunE: cobrautil.HandleContextCancellation(func(cmd *cobra.Command, args []string) error {
 			options.builder = rootOpts.builder
 			if len(args) > 0 {
 				if options.allInactive {
@@ -89,7 +90,7 @@ func rmCmd(dockerCli command.Cli, rootOpts *rootOptions) *cobra.Command {
 				options.builder = args[0]
 			}
 			return runRm(cmd.Context(), dockerCli, options)
-		},
+		}),
 		ValidArgsFunction: completion.BuilderNames(dockerCli),
 	}
 

commands/stop.go

@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"github.com/docker/buildx/builder"
+	"github.com/docker/buildx/util/cobrautil"
 	"github.com/docker/buildx/util/cobrautil/completion"
 	"github.com/docker/cli/cli"
 	"github.com/docker/cli/cli/command"
@@ -37,13 +38,13 @@ func stopCmd(dockerCli command.Cli, rootOpts *rootOptions) *cobra.Command {
 		Use:   "stop [NAME]",
 		Short: "Stop builder instance",
 		Args:  cli.RequiresMaxArgs(1),
-		RunE: func(cmd *cobra.Command, args []string) error {
+		RunE: cobrautil.HandleContextCancellation(func(cmd *cobra.Command, args []string) error {
 			options.builder = rootOpts.builder
 			if len(args) > 0 {
 				options.builder = args[0]
 			}
 			return runStop(cmd.Context(), dockerCli, options)
-		},
+		}),
 		ValidArgsFunction: completion.BuilderNames(dockerCli),
 	}
 

util/cobrautil/cobrautil.go

@@ -1,6 +1,16 @@
 package cobrautil
 
-import "github.com/spf13/cobra"
+import (
+	"context"
+	"fmt"
+	"os"
+	"os/signal"
+	"strings"
+	"syscall"
+
+	"github.com/docker/cli/cli-plugins/plugin"
+	"github.com/spf13/cobra"
+)
 
 // HideInheritedFlags hides inherited flags
 func HideInheritedFlags(cmd *cobra.Command, hidden ...string) {
@@ -12,3 +22,39 @@ func HideInheritedFlags(cmd *cobra.Command, hidden ...string) {
 		_ = cmd.Flags().MarkHidden(h)
 	}
 }
+
+// HandleContextCancellation checks and, if necessary, sets up signal
+// handling and hooks it up the cobra command's command.Context. This
+// is necessary for standalone plugin scenarios where the CLI isn't
+// already doing the signal/context handling.
+func HandleContextCancellation(fn func(*cobra.Command, []string) error) func(cmd *cobra.Command, args []string) error {
+	return func(cmd *cobra.Command, args []string) error {
+		ctx := cmd.Context()
+		contextString := fmt.Sprintf("%s", ctx)
+		if !strings.Contains(contextString, ".WithCancel") || plugin.RunningStandalone() {
+			// context isn't cancellable/plugin is running standalone
+			// so we need to do signal handling ourselves
+			cancellableCtx, cancel := context.WithCancel(cmd.Context())
+			ctx = cancellableCtx
+
+			signalLimit := 3
+			s := make(chan os.Signal, signalLimit)
+			signal.Notify(s, syscall.SIGTERM, syscall.SIGINT)
+			go func() {
+				retries := 0
+				for range s {
+					if retries == 0 {
+						cancel()
+					}
+					retries++
+					if retries >= signalLimit {
+						_, _ = fmt.Fprintf(os.Stderr, "got %d SIGTERM/SIGINTs, forcefully exiting\n", retries)
+						os.Exit(1)
+					}
+				}
+			}()
+		}
+		cmd.SetContext(ctx)
+		return fn(cmd, args)
+	}
+}