diff --git a/cmd/environments/sync.go b/cmd/environments/sync.go index 5f534f8605a..acce4eb4c57 100644 --- a/cmd/environments/sync.go +++ b/cmd/environments/sync.go @@ -3,7 +3,9 @@ package environments import ( "context" "os" + "os/signal" "path/filepath" + "syscall" "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/cmdctx" @@ -64,7 +66,12 @@ func addComputeFlags(cmd *cobra.Command) { // runPipeline builds and runs the setup-local Pipeline. func runPipeline(cmd *cobra.Command) error { - ctx := cmd.Context() + // The CLI root doesn't cancel ctx on signals, so handle them here: SIGINT + // (Ctrl-C) and SIGTERM (how a supervisor, CI timeout, or VS Code stops the + // child) cancel ctx, which propagates to the uv subprocesses the pipeline + // spawns so they are reaped instead of orphaned mid-provision. + ctx, stop := signal.NotifyContext(cmd.Context(), os.Interrupt, syscall.SIGTERM) + defer stop() cluster, _ := cmd.Flags().GetString("cluster-id") clusterName, _ := cmd.Flags().GetString("cluster-name") diff --git a/libs/localenv/uv.go b/libs/localenv/uv.go index 81fbc8802c8..dc8c5bb35e0 100644 --- a/libs/localenv/uv.go +++ b/libs/localenv/uv.go @@ -63,7 +63,7 @@ func (m *uvManager) EnsureAvailable(ctx context.Context) (string, error) { m.bin = bin // Use --version (not "version") to avoid project-scoped sub-command that requires pyproject.toml. - version, err := process.Background(ctx, []string{m.bin, "--version"}) + version, err := process.Background(ctx, []string{m.bin, "--version"}, process.WithProcessGroup()) if err != nil { return "", uvFailure(ErrUvMissing, err, "uv version check") } @@ -75,12 +75,15 @@ func (m *uvManager) EnsureAvailable(ctx context.Context) (string, error) { // (process.WithDir("") is a no-op). The index-url is injected only when // resolveIndexURL returns non-empty; it returns "" when UV_INDEX_URL is already // set, so an explicit value in the environment is never clobbered. +// WithProcessGroup is applied because uv fans out to its own subprocesses +// (Python, build backends); on SIGINT/SIGTERM they must be reaped as a group +// rather than left as orphans holding locks over a half-written .venv. func (m *uvManager) runUv(ctx context.Context, args []string, dir string) error { if indexURL := m.resolveIndexURL(ctx); indexURL != "" { - _, err := process.Background(ctx, args, process.WithDir(dir), process.WithEnv("UV_INDEX_URL", indexURL)) + _, err := process.Background(ctx, args, process.WithDir(dir), process.WithEnv("UV_INDEX_URL", indexURL), process.WithProcessGroup()) return err } - _, err := process.Background(ctx, args, process.WithDir(dir)) + _, err := process.Background(ctx, args, process.WithDir(dir), process.WithProcessGroup()) return err } @@ -155,6 +158,7 @@ except importlib.metadata.PackageNotFoundError: out, err := process.Background(ctx, []string{venvPython(projectDir), "-c", pyCode}, process.WithDir(projectDir), + process.WithProcessGroup(), ) if err != nil { return "", "", uvFailure(ErrValidate, err, "venv python validation") @@ -361,7 +365,10 @@ func installUv(ctx context.Context) error { // (~/.local/bin), so record exactly what ran before it fires — visible under // --debug for anyone auditing where uv came from. log.Debugf(ctx, "uv: not found; running installer: %s", strings.Join(cmd, " ")) - _, err := process.Background(ctx, cmd) + // The installer is a shell/PowerShell pipeline that spawns curl and the + // downloaded script; reap the whole group on cancellation so an interrupted + // install leaves no orphaned downloader behind. + _, err := process.Background(ctx, cmd, process.WithProcessGroup()) return err } diff --git a/libs/process/group.go b/libs/process/group.go new file mode 100644 index 00000000000..6da5309075a --- /dev/null +++ b/libs/process/group.go @@ -0,0 +1,8 @@ +package process + +import "time" + +// processGroupGracePeriod bounds how long WithProcessGroup waits after the +// context is cancelled before escalating to SIGKILL. It mirrors the 10s grace +// period used elsewhere for subprocess termination (see experimental/ssh). +const processGroupGracePeriod = 10 * time.Second diff --git a/libs/process/group_other.go b/libs/process/group_other.go new file mode 100644 index 00000000000..7f4750aa3ff --- /dev/null +++ b/libs/process/group_other.go @@ -0,0 +1,24 @@ +//go:build !unix + +package process + +import ( + "context" + "os/exec" +) + +// WithProcessGroup sets a WaitDelay so a cancelled command does not block +// indefinitely on a stuck child. +// +// Unlike the Unix build, this does not reap the child's descendants: killing a +// whole process tree on Windows requires a Job Object (CreateJobObject + +// AssignProcessToJobObject with JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE), which is +// out of scope here. The direct child is still terminated by the default +// cancellation, and the caller's signal handler still fires; only grandchildren +// spawned by the child may outlive it. +func WithProcessGroup() execOption { + return func(_ context.Context, c *exec.Cmd) error { + c.WaitDelay = processGroupGracePeriod + return nil + } +} diff --git a/libs/process/group_unix.go b/libs/process/group_unix.go new file mode 100644 index 00000000000..31c2fdf6dce --- /dev/null +++ b/libs/process/group_unix.go @@ -0,0 +1,44 @@ +//go:build unix + +package process + +import ( + "context" + "errors" + "os" + "os/exec" + "syscall" +) + +// WithProcessGroup makes the child the leader of a new process group and, when +// the context is cancelled, signals the entire group rather than just the child. +// +// exec.CommandContext's default cancellation only SIGKILLs the direct child, so +// a tool that fans out to its own subprocesses (e.g. `uv sync` spawning Python +// and build backends) leaves those grandchildren running as orphans when the CLI +// receives SIGINT/SIGTERM. Putting the child in its own group and signalling the +// group (negative PID) delivers SIGTERM to every descendant at once, giving them +// a chance to exit cleanly; WaitDelay then escalates to SIGKILL if the leader +// hangs past the grace period. +func WithProcessGroup() execOption { + return func(_ context.Context, c *exec.Cmd) error { + if c.SysProcAttr == nil { + c.SysProcAttr = &syscall.SysProcAttr{} + } + c.SysProcAttr.Setpgid = true + + c.WaitDelay = processGroupGracePeriod + c.Cancel = func() error { + // With Setpgid and Pgid unset, the child's group ID equals its PID; + // a negative PID targets the whole group. Map "no such process" to + // os.ErrProcessDone so a benign exit/cancel race is not surfaced as a + // Wait error. + err := syscall.Kill(-c.Process.Pid, syscall.SIGTERM) + if errors.Is(err, syscall.ESRCH) { + return os.ErrProcessDone + } + return err + } + return nil + } +} diff --git a/libs/process/group_unix_test.go b/libs/process/group_unix_test.go new file mode 100644 index 00000000000..837db600f42 --- /dev/null +++ b/libs/process/group_unix_test.go @@ -0,0 +1,75 @@ +//go:build unix + +package process + +import ( + "context" + "errors" + "os" + "path/filepath" + "strconv" + "strings" + "syscall" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestWithProcessGroupReapsGrandchild verifies that cancelling the context kills +// the whole process group, not just the direct child. The shell (the group +// leader) backgrounds a long sleep — a grandchild of the test process — and +// records its PID; after cancellation that grandchild must be gone. +func TestWithProcessGroupReapsGrandchild(t *testing.T) { + ctx, cancel := context.WithCancel(t.Context()) + + pidFile := filepath.Join(t.TempDir(), "grandchild.pid") + // $1 is pidFile: background a sleep, record its PID, then wait on it so the + // shell stays alive as the group leader until the group is signalled. + script := []string{"sh", "-c", `sleep 300 & echo $! > "$1"; wait`, "sh", pidFile} + + done := make(chan struct{}) + go func() { + defer close(done) + _, _ = Background(ctx, script, WithProcessGroup()) + }() + + grandchildPid := waitForPid(t, pidFile) + + cancel() + + select { + case <-done: + case <-time.After(processGroupGracePeriod + 5*time.Second): + t.Fatal("Background did not return after context cancellation") + } + + // The grandchild inherited the leader's group, so the group SIGTERM reaches + // it directly. Poll briefly to let the kernel deliver the signal and reap it. + assert.Eventually(t, func() bool { + return errIsProcessGone(syscall.Kill(grandchildPid, 0)) + }, 5*time.Second, 20*time.Millisecond, "grandchild %d was orphaned, not reaped", grandchildPid) +} + +// waitForPid waits for the shell to write the grandchild PID and returns it. +func waitForPid(t *testing.T, pidFile string) int { + t.Helper() + var pid int + require.Eventually(t, func() bool { + b, err := os.ReadFile(pidFile) + if err != nil { + return false + } + pid, err = strconv.Atoi(strings.TrimSpace(string(b))) + return err == nil && pid > 0 + }, 5*time.Second, 20*time.Millisecond, "grandchild PID was never recorded") + return pid +} + +// errIsProcessGone reports whether a signal-0 probe means the process no longer +// exists (ESRCH). A dead-but-unreaped zombie would answer differently, but the +// grandchild reparents to init on the leader's death and is reaped there. +func errIsProcessGone(err error) bool { + return errors.Is(err, syscall.ESRCH) +}