From 650a7af0ae5fe96b37166bfa81ed5cd8cf4df287 Mon Sep 17 00:00:00 2001 From: Laura Brehm Date: Fri, 5 Jan 2024 13:42:06 +0000 Subject: [PATCH] cobra/commands: cancel command context on signal See https://github.com/docker/cli/pull/4599 and https://github.com/docker/cli/pull/4769. Since we switched to using the cobra command context instead of `appcontext`, we need to set up the signal handling that was being provided by `appcontext`, as well as configuring the context with the OTEL tracing utilities also used by `appcontext`. This commit introduces `cobrautil.ConfigureContext` which implements the pre-existing signal handling logic from `appcontext` and cancels the command's context when a signal is received, as well as doing the relevant OTEL config. Signed-off-by: Laura Brehm --- commands/bake.go | 5 ++- commands/build.go | 4 +- commands/create.go | 4 +- commands/diskusage.go | 5 ++- commands/imagetools/create.go | 5 ++- commands/imagetools/inspect.go | 5 ++- commands/inspect.go | 5 ++- commands/ls.go | 4 +- commands/prune.go | 5 ++- commands/rm.go | 5 ++- commands/stop.go | 5 ++- util/cobrautil/cobrautil.go | 39 +++++++++++++++++++ util/cobrautil/signals_unix.go | 10 +++++ util/cobrautil/signals_windows.go | 9 +++++ .../buildkit/util/tracing/env/traceenv.go | 4 +- 15 files changed, 90 insertions(+), 24 deletions(-) create mode 100644 util/cobrautil/signals_unix.go create mode 100644 util/cobrautil/signals_windows.go diff --git a/commands/bake.go b/commands/bake.go index e17a8ebc7..b07df5fa8 100644 --- a/commands/bake.go +++ b/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" @@ -261,7 +262,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.ConfigureContext(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 @@ -273,7 +274,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 index e142befd9..555fab595 100644 --- 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: func(cmd *cobra.Command, args []string) error { + RunE: cobrautil.ConfigureContext(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 index 9362aebba..a58029ce5 100644 --- 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: func(cmd *cobra.Command, args []string) error { + RunE: cobrautil.ConfigureContext(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 index 4460c0fee..2994ae09e 100644 --- a/commands/diskusage.go +++ b/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.ConfigureContext(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 index f05571b03..689e4d580 100644 --- a/commands/imagetools/create.go +++ b/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.ConfigureContext(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 index 5b06f52fc..e5206b711 100644 --- a/commands/imagetools/inspect.go +++ b/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.ConfigureContext(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 index e341691c6..90226cdca 100644 --- a/commands/inspect.go +++ b/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.ConfigureContext(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 index 3b18e42cc..c90c5004f 100644 --- 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: func(cmd *cobra.Command, args []string) error { + RunE: cobrautil.ConfigureContext(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 index 1a2a5220c..4dc03e6db 100644 --- a/commands/prune.go +++ b/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.ConfigureContext(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 index af6ffd59f..248fee8eb 100644 --- a/commands/rm.go +++ b/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/command" "github.com/pkg/errors" @@ -97,7 +98,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: func(cmd *cobra.Command, args []string) error { + RunE: cobrautil.ConfigureContext(func(cmd *cobra.Command, args []string) error { options.builders = []string{rootOpts.builder} if len(args) > 0 { if options.allInactive { @@ -106,7 +107,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/stop.go b/commands/stop.go index b27985af8..865ecd93c 100644 --- a/commands/stop.go +++ b/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.ConfigureContext(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/util/cobrautil/cobrautil.go b/util/cobrautil/cobrautil.go index 71633bb1d..97c1aa04a 100644 --- a/util/cobrautil/cobrautil.go +++ b/util/cobrautil/cobrautil.go @@ -1,6 +1,13 @@ 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" @@ -51,3 +58,35 @@ 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/util/cobrautil/signals_unix.go b/util/cobrautil/signals_unix.go new file mode 100644 index 000000000..8e6899d8c --- /dev/null +++ b/util/cobrautil/signals_unix.go @@ -0,0 +1,10 @@ +//go:build !windows + +package cobrautil + +import ( + "golang.org/x/sys/unix" + "os" +) + +var interruptSignals = []os.Signal{unix.SIGTERM, unix.SIGINT} diff --git a/util/cobrautil/signals_windows.go b/util/cobrautil/signals_windows.go new file mode 100644 index 000000000..bab02ac8e --- /dev/null +++ b/util/cobrautil/signals_windows.go @@ -0,0 +1,9 @@ +//go:build windows + +package cobrautil + +import ( + "os" +) + +var interruptSignals = []os.Signal{os.Interrupt} diff --git a/vendor/github.com/moby/buildkit/util/tracing/env/traceenv.go b/vendor/github.com/moby/buildkit/util/tracing/env/traceenv.go index 0a95619d1..711d7cab3 100644 --- a/vendor/github.com/moby/buildkit/util/tracing/env/traceenv.go +++ b/vendor/github.com/moby/buildkit/util/tracing/env/traceenv.go @@ -14,10 +14,10 @@ const ( ) func init() { - appcontext.Register(initContext) + appcontext.Register(InitContext) } -func initContext(ctx context.Context) context.Context { +func InitContext(ctx context.Context) context.Context { // open-telemetry/opentelemetry-specification#740 parent := os.Getenv("TRACEPARENT") state := os.Getenv("TRACESTATE")