diff --git a/cmd/air/logs.go b/cmd/air/logs.go index 9689e1e8a4a..1202978b998 100644 --- a/cmd/air/logs.go +++ b/cmd/air/logs.go @@ -5,12 +5,14 @@ import ( "errors" "fmt" "io" + "os" "slices" "strconv" "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/cmdctx" "github.com/databricks/cli/libs/flags" + "github.com/databricks/cli/libs/process" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/apierr" "github.com/databricks/databricks-sdk-go/service/iam" @@ -116,7 +118,7 @@ JOB_RUN_ID is the job run ID returned by "air run".`, streamCtx := ctx if downloadTo == "" && root.OutputType(cmd) != flags.OutputJSON { var stop context.CancelFunc - streamCtx, stop = notifyInterrupt(ctx) + streamCtx, stop = process.NotifyContext(ctx, os.Interrupt) defer stop() } diff --git a/cmd/air/run.go b/cmd/air/run.go index c8027b09e9a..a77c6a5ee3c 100644 --- a/cmd/air/run.go +++ b/cmd/air/run.go @@ -6,7 +6,6 @@ import ( "fmt" "io" "os" - "os/signal" "strconv" "strings" "time" @@ -16,6 +15,7 @@ import ( "github.com/databricks/cli/libs/cmdctx" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/flags" + "github.com/databricks/cli/libs/process" "github.com/databricks/cli/libs/shellquote" "github.com/databricks/databricks-sdk-go" "github.com/spf13/cobra" @@ -213,7 +213,9 @@ The path must be a separate argument: cobra reserves -h as a boolean, so retryTracker: newRetryTracker(logRunStatus{}), } - watchCtx, stop := notifyInterrupt(ctx) + // Catch the interrupt so the log stream unwinds and handleWatchResult + // prints resume guidance; the CLI root installs no signal handler of its own. + watchCtx, stop := process.NotifyContext(ctx, os.Interrupt) defer stop() if !jsonOut { @@ -288,34 +290,6 @@ func printPostSubmitGuidance(out io.Writer, profile, runID string) { fmt.Fprintln(out, " "+airLogsCommand(profile, runID)) } -// notifyInterrupt returns a context cancelled on the first interrupt (Ctrl-C), -// which lets the log stream unwind and print resume guidance via -// handleWatchResult; the CLI root installs no signal handler of its own. The -// caller must defer the returned stop. -// -// signal.Notify disables the default SIGINT disposition process-wide for as long -// as the channel stays registered, so a second Ctrl-C would merely be buffered -// and dropped, leaving no way to abort a hung teardown. Calling signal.Stop -// before cancel restores SIG_DFL first, so the cancellation is only observable -// once a second signal is guaranteed to terminate the process. -func notifyInterrupt(parent context.Context) (context.Context, context.CancelFunc) { - ctx, cancel := context.WithCancel(parent) - sigCh := make(chan os.Signal, 1) - signal.Notify(sigCh, os.Interrupt) - go func() { - // Selecting on ctx.Done() too lets the goroutine exit on the normal (no - // signal) path rather than parking on sigCh for the life of the process. - select { - case <-sigCh: - signal.Stop(sigCh) - cancel() - case <-ctx.Done(): - signal.Stop(sigCh) - } - }() - return ctx, cancel -} - func handleWatchResult(out io.Writer, profile, runID string, err error) error { if !errors.Is(err, context.Canceled) { return err diff --git a/cmd/air/run_signal_test.go b/cmd/air/run_signal_test.go deleted file mode 100644 index ca243ca0a06..00000000000 --- a/cmd/air/run_signal_test.go +++ /dev/null @@ -1,101 +0,0 @@ -//go:build unix - -// The second-signal escape hatch is a Unix signal-delivery behavior: the test -// re-execs itself and sends SIGINT, which os/signal cannot deliver to self on -// Windows (notifyInterrupt only watches os.Interrupt there anyway). -package aircmd - -import ( - "bufio" - "context" - "os" - "os/exec" - "strings" - "testing" - "time" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -// notifyInterrupt flips signal dispositions process-wide, so the second-signal -// behavior can only be observed in a real child process where the test binary's -// own handlers don't mask it. Go resets dispositions to SIG_DFL across exec. -const signalChildEnv = "TEST_AIR_SIGNAL_CHILD" - -// TestMain runs notifyInterrupt under test when re-executed as the child. -func TestMain(m *testing.M) { - if os.Getenv(signalChildEnv) == "" { - os.Exit(m.Run()) - } - - // TestMain has no *testing.T, so t.Context() is unavailable here. - //nolint:gocritic - ctx, stop := notifyInterrupt(context.Background()) - defer stop() - - // Stand in for a log stream still draining after the first Ctrl-C; the user - // needs the escape hatch during this window. - os.Stdout.WriteString("READY\n") - <-ctx.Done() - os.Stdout.WriteString("CANCELLED\n") - time.Sleep(30 * time.Second) - os.Stdout.WriteString("SURVIVED\n") - os.Exit(0) -} - -// TestNotifyInterruptSecondSignalStillKills is the regression test for the -// escape hatch: signal.NotifyContext disables the default SIGINT disposition -// process-wide, so relaying only the first signal would leave the user unable to -// abort. The first Ctrl-C must cancel the context (so handleWatchResult prints -// resume guidance), and the second must terminate the process outright. -func TestNotifyInterruptSecondSignalStillKills(t *testing.T) { - cmd := exec.Command(os.Args[0], "-test.run=TestNotifyInterruptSecondSignalStillKills") - cmd.Env = append(os.Environ(), signalChildEnv+"=1") - stdout, err := cmd.StdoutPipe() - require.NoError(t, err) - require.NoError(t, cmd.Start()) - t.Cleanup(func() { _ = cmd.Process.Kill() }) - - lines := make(chan string, 8) - go func() { - sc := bufio.NewScanner(stdout) - for sc.Scan() { - lines <- sc.Text() - } - close(lines) - }() - await := func(want string) bool { - deadline := time.After(30 * time.Second) - for { - select { - case l, ok := <-lines: - if !ok { - return false - } - if strings.Contains(l, want) { - return true - } - case <-deadline: - return false - } - } - } - - require.True(t, await("READY"), "child never started") - require.NoError(t, cmd.Process.Signal(os.Interrupt)) - require.True(t, await("CANCELLED"), "first interrupt did not cancel the context") - - // The second interrupt must reach the default disposition and kill the child - // rather than being buffered and dropped by the handler. - require.NoError(t, cmd.Process.Signal(os.Interrupt)) - waitErr := make(chan error, 1) - go func() { waitErr <- cmd.Wait() }() - select { - case err := <-waitErr: - // Killed by the signal, so a non-nil (non-exit-zero) error. - assert.Error(t, err, "child should die by signal, not exit cleanly") - case <-time.After(15 * time.Second): - t.Fatal("second interrupt was swallowed: the user has no way to abort a hung stream") - } -} diff --git a/cmd/environments/sync.go b/cmd/environments/sync.go index ae10ae08eb3..d5997a83969 100644 --- a/cmd/environments/sync.go +++ b/cmd/environments/sync.go @@ -1,9 +1,7 @@ package environments import ( - "context" "os" - "os/signal" "path/filepath" "syscall" @@ -13,6 +11,7 @@ import ( libslocalenv "github.com/databricks/cli/libs/localenv" "github.com/databricks/cli/libs/log" "github.com/databricks/cli/libs/logdiag" + "github.com/databricks/cli/libs/process" "github.com/spf13/cobra" ) @@ -81,55 +80,19 @@ func addComputeFlags(cmd *cobra.Command) { // consumer relies on, instead of a bare pre-RunE Cobra error. } -// watchInterruptSignals cancels ctx on the first SIGINT (Ctrl-C) or SIGTERM (how -// a supervisor, CI timeout, or VS Code stops the child), which propagates to the -// uv subprocesses the pipeline spawns so they are reaped instead of orphaned -// mid-provision. The CLI root installs no signal handler of its own. -// -// The returned stop function uninstalls the handler and joins the goroutine; the -// caller must defer it. -// -// The handler must give the *second* signal back to the OS. signal.Notify (like -// signal.NotifyContext, which wraps it) disables the default disposition for -// SIGINT/SIGTERM for as long as the channel stays registered, so without the -// signal.Stop below a second Ctrl-C is merely buffered and dropped: the user -// would have no way to abort during the process group's SIGKILL grace window. -// Stopping the relay as soon as the first signal lands restores SIG_DFL, so a -// second signal terminates the CLI immediately. That matters more here than in -// most commands because WithProcessGroup moves uv out of the foreground process -// group, so the tty no longer delivers Ctrl-C to it directly — this handler is -// the only delivery path. -func watchInterruptSignals(ctx context.Context, cancel context.CancelFunc) func() { - sigCh := make(chan os.Signal, 1) - signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM) - - done := make(chan struct{}) - go func() { - defer close(done) - // Selecting on ctx.Done() too lets the goroutine exit on the normal (no - // signal) path rather than blocking on sigCh for the rest of the process: - // signal.Stop unregisters the channel but never closes it. - select { - case <-sigCh: - signal.Stop(sigCh) - cancel() - case <-ctx.Done(): - } - }() - - return func() { - signal.Stop(sigCh) - // Wake the goroutine in case neither sigCh nor ctx.Done has fired. - cancel() - <-done - } -} - // runPipeline builds and runs the setup-local Pipeline. func runPipeline(cmd *cobra.Command) error { - ctx, cancel := context.WithCancel(cmd.Context()) - defer cancel() - defer watchInterruptSignals(ctx, cancel)() + // Cancel on the first SIGINT (Ctrl-C) or SIGTERM (how a supervisor, CI + // timeout, or VS Code stops the child), which propagates to the uv + // subprocesses the pipeline spawns so they are reaped instead of orphaned + // mid-provision. The CLI root installs no signal handler of its own. + // + // Letting the second signal through matters more here than in most commands + // because WithProcessGroup moves uv out of the foreground process group, so + // the tty no longer delivers Ctrl-C to it directly — this handler is the only + // delivery path. + ctx, stop := process.NotifyContext(cmd.Context(), os.Interrupt, syscall.SIGTERM) + defer stop() cluster, _ := cmd.Flags().GetString("cluster-id") clusterName, _ := cmd.Flags().GetString("cluster-name") diff --git a/cmd/environments/sync_signal_test.go b/cmd/environments/sync_signal_test.go deleted file mode 100644 index 59477bccd67..00000000000 --- a/cmd/environments/sync_signal_test.go +++ /dev/null @@ -1,131 +0,0 @@ -//go:build unix - -// The second-signal escape hatch is a Unix signal-delivery behavior: the test -// re-execs itself and sends SIGINT/SIGTERM, which os/signal does not support on -// Windows (the CLI's own signal handling there is likewise a no-op path). -package environments - -import ( - "bufio" - "context" - "os" - "os/exec" - "strings" - "syscall" - "testing" - "time" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -// The second-signal behavior can only be observed in a real process: signal -// dispositions are process-wide, and the test binary's own handlers would mask -// them. So the test re-executes itself as a child (Go resets dispositions to -// SIG_DFL across exec, unlike a shell background job, which inherits SIGINT as -// SIG_IGN and would silently invalidate the result). -const signalChildEnv = "TEST_ENVIRONMENTS_SIGNAL_CHILD" - -// TestMain runs the signal-handler-under-test when re-executed as the child. -func TestMain(m *testing.M) { - if os.Getenv(signalChildEnv) == "" { - os.Exit(m.Run()) - } - - // TestMain has no *testing.T, so t.Context() is unavailable here. - //nolint:gocritic - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - stop := watchInterruptSignals(ctx, cancel) - defer stop() - - // Stand in for the pipeline still draining the uv process group during the - // SIGKILL grace window — the interval in which the user needs the hatch. - os.Stdout.WriteString("READY\n") - <-ctx.Done() - os.Stdout.WriteString("CANCELLED\n") - time.Sleep(30 * time.Second) - os.Stdout.WriteString("SURVIVED\n") - os.Exit(0) -} - -// TestWatchInterruptSignalsSecondSignalStillKills is the regression test for the -// escape hatch: signal.Notify disables the default SIGINT/SIGTERM disposition -// process-wide, so a handler that only relays the first signal leaves the user -// unable to abort. The first signal must cancel the context, and the second must -// terminate the process outright. -func TestWatchInterruptSignalsSecondSignalStillKills(t *testing.T) { - for _, sig := range []syscall.Signal{syscall.SIGINT, syscall.SIGTERM} { - t.Run(sig.String(), func(t *testing.T) { - cmd := exec.Command(os.Args[0], "-test.run=TestWatchInterruptSignalsSecondSignalStillKills") - cmd.Env = append(os.Environ(), signalChildEnv+"=1") - stdout, err := cmd.StdoutPipe() - require.NoError(t, err) - require.NoError(t, cmd.Start()) - t.Cleanup(func() { _ = cmd.Process.Kill() }) - - lines := make(chan string, 8) - go func() { - sc := bufio.NewScanner(stdout) - for sc.Scan() { - lines <- sc.Text() - } - close(lines) - }() - await := func(want string) bool { - deadline := time.After(30 * time.Second) - for { - select { - case l, ok := <-lines: - if !ok { - return false - } - if strings.Contains(l, want) { - return true - } - case <-deadline: - return false - } - } - } - - require.True(t, await("READY"), "child never started") - require.NoError(t, cmd.Process.Signal(sig)) - require.True(t, await("CANCELLED"), "first %s did not cancel the context", sig) - - // The second signal must reach the default disposition and kill the - // child rather than being buffered and dropped by the handler. - require.NoError(t, cmd.Process.Signal(sig)) - waitErr := make(chan error, 1) - go func() { waitErr <- cmd.Wait() }() - select { - case err := <-waitErr: - // Killed by the signal, so a non-nil (non-exit-zero) error. - assert.Error(t, err, "child should die by signal, not exit cleanly") - case <-time.After(15 * time.Second): - t.Fatalf("second %s was swallowed: the user has no escape hatch "+ - "during the process group's SIGKILL grace window", sig) - } - }) - } -} - -// TestWatchInterruptSignalsStopsWithoutSignal covers the no-signal path: stop() -// must return (joining its goroutine) rather than leaving it parked on the -// channel for the life of the process. -func TestWatchInterruptSignalsStopsWithoutSignal(t *testing.T) { - ctx, cancel := context.WithCancel(t.Context()) - defer cancel() - stop := watchInterruptSignals(ctx, cancel) - - returned := make(chan struct{}) - go func() { - stop() - close(returned) - }() - select { - case <-returned: - case <-time.After(10 * time.Second): - t.Fatal("stop() blocked: the signal goroutine was never joined") - } -} diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index 81ee1439718..f093442521c 100644 --- a/experimental/ssh/internal/client/client.go +++ b/experimental/ssh/internal/client/client.go @@ -13,7 +13,6 @@ import ( "net/http" "os" "os/exec" - "os/signal" "path/filepath" "regexp" "strconv" @@ -30,6 +29,7 @@ import ( "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/log" + "github.com/databricks/cli/libs/process" "github.com/databricks/cli/libs/telemetry" "github.com/databricks/cli/libs/telemetry/protos" "github.com/databricks/databricks-sdk-go" @@ -304,15 +304,13 @@ func (o *ClientOptions) ToProxyCommand() (string, error) { } func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOptions) (retErr error) { - ctx, cancel := context.WithCancel(ctx) - defer cancel() - - sigCh := make(chan os.Signal, 1) - signal.Notify(sigCh, os.Interrupt, syscall.SIGHUP, syscall.SIGTERM) + ctx, stopSignals := process.NotifyContext(ctx, os.Interrupt, syscall.SIGHUP, syscall.SIGTERM) + defer stopSignals() go func() { - <-sigCh - cmdio.LogString(ctx, "Received termination signal, cleaning up...") - cancel() + <-ctx.Done() + if _, ok := errors.AsType[*process.SignalError](context.Cause(ctx)); ok { + cmdio.LogString(ctx, "Received termination signal, cleaning up...") + } }() // Report the outcome of every path below, so a failure is attributable to the step that @@ -326,7 +324,7 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt // A cancelled context is the only trace a Ctrl-C leaves: exec.CommandContext kills the // child and reports *exec.ExitError, which does not wrap context.Canceled, so a step // that shells out cannot recognise the interruption itself. This defer is registered - // after `defer cancel()` and so runs before it (LIFO), which means ctx is cancelled + // after `defer stopSignals()` and so runs before it (LIFO), which means ctx is cancelled // here only by the signal handler or the caller, never by Run's own cleanup. outcome.ctxErr = ctx.Err() logSshTunnelEvent(ctx, opts, outcome)