simplify signal handling for cobra context

tonistiigi · tonistiigi · commit 5b2ecca2b79e · 2024-01-19T16:59:56.000-08:00
Signed-off-by: Tonis Tiigi &lt;tonistiigi@gmail.com&gt;
diff --git a/commands/bake.go b/commands/bake.go
@@ -15,7 +15,6 @@ 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"
@@ -262,7 +261,7 @@ func bakeCmd(dockerCli command.Cli, rootOpts *rootOptions) *cobra.Command {
 		Use:     "bake [OPTIONS] [TARGET...]",
 		Aliases: []string{"f"},
 		Short:   "Build from a file",
-		RunE: cobrautil.ConfigureContext(func(cmd *cobra.Command, args []string) error {
+		RunE: 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
@@ -274,7 +273,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),
 	}
 
diff --git a/commands/build.go b/commands/build.go
@@ -461,7 +461,7 @@ func buildCmd(dockerCli command.Cli, rootOpts *rootOptions, debugConfig *debug.D
 		Aliases: []string{"b"},
 		Short:   "Start a build",
 		Args:    cli.ExactArgs(1),
-		RunE: cobrautil.ConfigureContext(func(cmd *cobra.Command, args []string) error {
+		RunE: func(cmd *cobra.Command, args []string) error {
 			options.contextPath = args[0]
 			options.builder = rootOpts.builder
 			options.metadataFile = cFlags.metadataFile
@@ -485,7 +485,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
 		},
diff --git a/commands/create.go b/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: cobrautil.ConfigureContext(func(cmd *cobra.Command, args []string) error {
+		RunE: func(cmd *cobra.Command, args []string) error {
 			return runCreate(cmd.Context(), dockerCli, options, args)
-		}),
+		},
 		ValidArgsFunction: completion.Disable,
 	}
 
diff --git a/commands/diskusage.go b/commands/diskusage.go
@@ -10,7 +10,6 @@ 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"
@@ -111,10 +110,10 @@ func duCmd(dockerCli command.Cli, rootOpts *rootOptions) *cobra.Command {
 		Use:   "du",
 		Short: "Disk usage",
 		Args:  cli.NoArgs,
-		RunE: cobrautil.ConfigureContext(func(cmd *cobra.Command, args []string) error {
+		RunE: func(cmd *cobra.Command, args []string) error {
 			options.builder = rootOpts.builder
 			return runDiskUsage(cmd.Context(), dockerCli, options)
-		}),
+		},
 		ValidArgsFunction: completion.Disable,
 	}
 
diff --git a/commands/imagetools/create.go b/commands/imagetools/create.go
@@ -9,7 +9,6 @@ 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"
@@ -270,10 +269,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: cobrautil.ConfigureContext(func(cmd *cobra.Command, args []string) error {
+		RunE: func(cmd *cobra.Command, args []string) error {
 			options.builder = *opts.Builder
 			return runCreate(cmd.Context(), dockerCli, options, args)
-		}),
+		},
 		ValidArgsFunction: completion.Disable,
 	}
 
diff --git a/commands/imagetools/inspect.go b/commands/imagetools/inspect.go
@@ -4,7 +4,6 @@ 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"
@@ -49,10 +48,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: cobrautil.ConfigureContext(func(cmd *cobra.Command, args []string) error {
+		RunE: func(cmd *cobra.Command, args []string) error {
 			options.builder = *rootOpts.Builder
 			return runInspect(cmd.Context(), dockerCli, options, args[0])
-		}),
+		},
 		ValidArgsFunction: completion.Disable,
 	}
 
diff --git a/commands/inspect.go b/commands/inspect.go
@@ -11,7 +11,6 @@ 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"
@@ -143,13 +142,13 @@ func inspectCmd(dockerCli command.Cli, rootOpts *rootOptions) *cobra.Command {
 		Use:   "inspect [NAME]",
 		Short: "Inspect current builder instance",
 		Args:  cli.RequiresMaxArgs(1),
-		RunE: cobrautil.ConfigureContext(func(cmd *cobra.Command, args []string) error {
+		RunE: 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),
 	}
 
diff --git a/commands/ls.go b/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: cobrautil.ConfigureContext(func(cmd *cobra.Command, args []string) error {
+		RunE: func(cmd *cobra.Command, args []string) error {
 			return runLs(cmd.Context(), dockerCli, options)
-		}),
+		},
 		ValidArgsFunction: completion.Disable,
 	}
 
diff --git a/commands/prune.go b/commands/prune.go
@@ -9,7 +9,6 @@ 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"
@@ -135,10 +134,10 @@ func pruneCmd(dockerCli command.Cli, rootOpts *rootOptions) *cobra.Command {
 		Use:   "prune",
 		Short: "Remove build cache",
 		Args:  cli.NoArgs,
-		RunE: cobrautil.ConfigureContext(func(cmd *cobra.Command, args []string) error {
+		RunE: func(cmd *cobra.Command, args []string) error {
 			options.builder = rootOpts.builder
 			return runPrune(cmd.Context(), dockerCli, options)
-		}),
+		},
 		ValidArgsFunction: completion.Disable,
 	}
 
diff --git a/commands/rm.go b/commands/rm.go
@@ -8,7 +8,6 @@ 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/command"
 	"github.com/pkg/errors"
@@ -98,7 +97,7 @@ func rmCmd(dockerCli command.Cli, rootOpts *rootOptions) *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "rm [OPTIONS] [NAME] [NAME...]",
 		Short: "Remove one or more builder instances",
-		RunE: cobrautil.ConfigureContext(func(cmd *cobra.Command, args []string) error {
+		RunE: func(cmd *cobra.Command, args []string) error {
 			options.builders = []string{rootOpts.builder}
 			if len(args) > 0 {
 				if options.allInactive {
@@ -107,7 +106,7 @@ func rmCmd(dockerCli command.Cli, rootOpts *rootOptions) *cobra.Command {
 				options.builders = args
 			}
 			return runRm(cmd.Context(), dockerCli, options)
-		}),
+		},
 		ValidArgsFunction: completion.BuilderNames(dockerCli),
 	}
 
diff --git a/commands/root.go b/commands/root.go
@@ -13,6 +13,7 @@ import (
 	"github.com/docker/cli/cli-plugins/plugin"
 	"github.com/docker/cli/cli/command"
 	"github.com/docker/cli/cli/debug"
+	"github.com/moby/buildkit/util/appcontext"
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
@@ -29,12 +30,15 @@ func NewRootCmd(name string, isPlugin bool, dockerCli command.Cli) *cobra.Comman
 		CompletionOptions: cobra.CompletionOptions{
 			HiddenDefaultCmd: true,
 		},
-	}
-	if isPlugin {
-		cmd.PersistentPreRunE = func(cmd *cobra.Command, args []string) error {
+		PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
+			cmd.SetContext(appcontext.Context())
+			if !isPlugin {
+				return nil
+			}
 			return plugin.PersistentPreRunE(cmd, args)
-		}
-	} else {
+		},
+	}
+	if !isPlugin {
 		// match plugin behavior for standalone mode
 		// https://github.com/docker/cli/blob/6c9eb708fa6d17765d71965f90e1c59cea686ee9/cli-plugins/plugin/plugin.go#L117-L127
 		cmd.SilenceUsage = true
diff --git a/commands/stop.go b/commands/stop.go
@@ -4,7 +4,6 @@ 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"
@@ -38,13 +37,13 @@ func stopCmd(dockerCli command.Cli, rootOpts *rootOptions) *cobra.Command {
 		Use:   "stop [NAME]",
 		Short: "Stop builder instance",
 		Args:  cli.RequiresMaxArgs(1),
-		RunE: cobrautil.ConfigureContext(func(cmd *cobra.Command, args []string) error {
+		RunE: 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),
 	}
 
diff --git a/go.mod b/go.mod
@@ -55,6 +55,8 @@ require (
 	k8s.io/client-go v0.26.7
 )
 
+require golang.org/x/sys v0.15.0
+
 require (
 	github.com/AdaLogics/go-fuzz-headers v0.0.0-20230811130428-ced1acdcaa24 // indirect
 	github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1 // indirect
@@ -152,7 +154,6 @@ require (
 	golang.org/x/exp v0.0.0-20230713183714-613f0c0eb8a1 // indirect
 	golang.org/x/net v0.17.0 // indirect
 	golang.org/x/oauth2 v0.10.0 // indirect
-	golang.org/x/sys v0.15.0 // indirect
 	golang.org/x/text v0.14.0 // indirect
 	golang.org/x/time v0.3.0 // indirect
 	golang.org/x/tools v0.10.0 // indirect
diff --git a/util/cobrautil/cobrautil.go b/util/cobrautil/cobrautil.go
@@ -1,13 +1,6 @@
 package cobrautil
 
 import (
-	"context"
-	"os"
-	"os/signal"
-
-	"github.com/moby/buildkit/util/bklog"
-	detect "github.com/moby/buildkit/util/tracing/env"
-	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
@@ -58,35 +51,3 @@ func MarkCommandExperimental(c *cobra.Command) {
 	c.Annotations[annotationExperimentalCLI] = ""
 	c.Short += " (EXPERIMENTAL)"
 }
-
-// ConfigureContext sets up signal handling and hooks into the command's
-// context so that it will be cancelled when signalled, as well as implementing
-// the "hard exit after 3 signals" logic. It also configures OTEL tracing
-// for the relevant context.
-func ConfigureContext(fn func(*cobra.Command, []string) error) func(cmd *cobra.Command, args []string) error {
-	return func(cmd *cobra.Command, args []string) error {
-		ctx := detect.InitContext(cmd.Context())
-		cancellableCtx, cancel := context.WithCancelCause(ctx)
-		ctx = cancellableCtx
-
-		signalLimit := 3
-		s := make(chan os.Signal, signalLimit)
-		signal.Notify(s, interruptSignals...)
-		go func() {
-			retries := 0
-			for {
-				<-s
-				retries++
-				err := errors.Errorf("got %d SIGTERM/SIGINTs, forcing shutdown", retries)
-				cancel(err)
-				if retries >= signalLimit {
-					bklog.G(ctx).Errorf(err.Error())
-					os.Exit(1)
-				}
-			}
-		}()
-
-		cmd.SetContext(ctx)
-		return fn(cmd, args)
-	}
-}
diff --git a/vendor/github.com/moby/buildkit/util/tracing/env/traceenv.go b/vendor/github.com/moby/buildkit/util/tracing/env/traceenv.go
