simplify signal handling for cobra context

Signed-off-by: Tonis Tiigi <[email protected]>

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),
  }
 

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
  },

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,
  }
 

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,
  }
 

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,
  }
 

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,
  }
 

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),
  }
 

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,
  }
 

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,
  }
 

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),
  }
 

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:/docker/cli/blob/6c9eb708fa6d17765d71965f90e1c59cea686ee9/cli-plugins/plugin/plugin.go#L117-L127
  cmd.SilenceUsage = true

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),
  }
 

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)
- }
-}