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>

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

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}

util/cobrautil/signals_windows.go

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