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 <laurabrehm@hey.com>
This commit is contained in:
Laura Brehm 2024-01-05 13:42:06 +00:00 committed by Tonis Tiigi
parent 4f738020fd
commit 650a7af0ae
No known key found for this signature in database
GPG Key ID: AFA9DE5F8AB7AF39
15 changed files with 90 additions and 24 deletions

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@ -0,0 +1,10 @@
//go:build !windows
package cobrautil
import (
"golang.org/x/sys/unix"
"os"
)
var interruptSignals = []os.Signal{unix.SIGTERM, unix.SIGINT}

View File

@ -0,0 +1,9 @@
//go:build windows
package cobrautil
import (
"os"
)
var interruptSignals = []os.Signal{os.Interrupt}

View File

@ -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")