Clean up Windows action process trees

pkg/runner/BUILD.bazel

@@ -7,6 +7,8 @@ go_library(
         "clean_runner.go",
         "local_runner.go",
         "local_runner_darwin.go",
+        "local_runner_process.go",
+        "local_runner_process_windows.go",
         "local_runner_rss_bytes.go",
         "local_runner_rss_kibibytes.go",
         "local_runner_unix.go",
@@ -69,12 +71,13 @@ go_test(
     srcs = [
         "apple_xcode_resolving_runner_test.go",
         "clean_runner_test.go",
+        "local_runner_process_windows_test.go",
         "local_runner_test.go",
         "path_existence_checking_runner_test.go",
         "temporary_directory_symlinking_runner_test.go",
     ],
+    embed = [":runner"],
     deps = [
-        ":runner",
         "//internal/mock",
         "//pkg/cleaner",
         "//pkg/proto/resourceusage",

pkg/runner/local_runner.go

@@ -125,7 +125,9 @@ func NewPlainCommandCreator(sysProcAttr *syscall.SysProcAttr) CommandCreator {
 }
 
 // NewLocalRunner returns a Runner capable of running commands on the
-// local system directly.
+// local system directly. On Windows, commands are placed in a
+// non-breakaway job object, and any surviving descendants are terminated
+// and waited for when the root process exits.
 func NewLocalRunner(buildDirectory filesystem.Directory, buildDirectoryPath *path.Builder, commandCreator CommandCreator, setTmpdirEnvironmentVariable bool) runner.RunnerServer {
 	return &localRunner{
 		buildDirectory:               buildDirectory,
@@ -185,10 +187,17 @@ func (r *localRunner) Run(ctx context.Context, request *runner.RunRequest) (*run
 
 	// Start the subprocess. We can already close the output files
 	// while the process is running.
+	commandProcess, err := prepareCommandForStart(cmd)
+	if err != nil {
+		stdout.Close()
+		stderr.Close()
+		return nil, util.StatusWrap(err, "Failed to prepare process")
+	}
 	err = cmd.Start()
 	stdout.Close()
 	stderr.Close()
 	if err != nil {
+		commandProcess.Close()
 		code := codes.Internal
 		for _, invalidArgumentErr := range invalidArgumentErrs {
 			if errors.Is(err, invalidArgumentErr) {
@@ -198,13 +207,20 @@ func (r *localRunner) Run(ctx context.Context, request *runner.RunRequest) (*run
 		}
 		return nil, util.StatusWrapWithCode(err, code, "Failed to start process")
 	}
-
 	// Wait for execution to complete. Permit non-zero exit codes.
-	if err := cmd.Wait(); err != nil {
-		if _, ok := err.(*exec.ExitError); !ok {
-			return nil, err
+	waitErr := cmd.Wait()
+	afterWaitErr := commandProcess.AfterWait(cmd)
+	if waitErr != nil {
+		if afterWaitErr != nil {
+			return nil, afterWaitErr
+		}
+		if _, ok := waitErr.(*exec.ExitError); !ok {
+			return nil, waitErr
 		}
 	}
+	if afterWaitErr != nil {
+		return nil, afterWaitErr
+	}
 
 	// Attach rusage information to the response.
 	posixResourceUsage, err := anypb.New(getPOSIXResourceUsage(cmd))

pkg/runner/local_runner_process.go

@@ -0,0 +1,24 @@
+//go:build !windows
+// +build !windows
+
+package runner
+
+import "os/exec"
+
+type commandProcess struct{}
+
+// prepareCommandForStart is called after the command is fully configured and
+// before Start. Non-Windows platforms do not need extra process-tree state.
+func prepareCommandForStart(cmd *exec.Cmd) (*commandProcess, error) {
+	return &commandProcess{}, nil
+}
+
+// AfterWait is called after Wait returns, before build directory cleanup can
+// proceed.
+func (commandProcess) AfterWait(cmd *exec.Cmd) error {
+	return nil
+}
+
+// Close releases any resources allocated by prepareCommandForStart. It must be
+// safe to call if Start fails.
+func (commandProcess) Close() {}

pkg/runner/local_runner_process_windows.go

@@ -0,0 +1,194 @@
+//go:build windows
+// +build windows
+
+// This file intentionally mirrors the native Bazel Windows launcher:
+// https://github.com/bazelbuild/bazel/blob/master/src/main/native/windows/process.cc
+//
+// The root process is assigned to a non-breakaway job as part of process
+// creation. This prevents it from spawning descendants before the runner has
+// contained it, which could leave those descendants alive to keep the input
+// root undeletable.
+
+package runner
+
+import (
+	"errors"
+	"os"
+	"os/exec"
+	"sync"
+	"syscall"
+	"unsafe"
+
+	"golang.org/x/sys/windows"
+)
+
+const (
+	// jobObjectMsgActiveProcessZero is JOB_OBJECT_MSG_ACTIVE_PROCESS_ZERO.
+	// The Windows SDK's winnt.h defines this message value as 4, but
+	// x/sys/windows does not expose it. The official job completion-port docs
+	// describe the message:
+	// https://learn.microsoft.com/en-us/windows/win32/api/winnt/ns-winnt-jobobject_associate_completion_port
+	jobObjectMsgActiveProcessZero = 4
+
+	windowsTerminatedExitCode = 130
+)
+
+type jobObjectAssociateCompletionPort struct {
+	CompletionKey  uintptr
+	CompletionPort windows.Handle
+}
+
+type commandProcess struct {
+	lock   sync.Mutex
+	job    windows.Handle
+	ioport windows.Handle
+	closed bool
+}
+
+// prepareCommandForStart creates the job object and completion port before the
+// process exists. It also amends cmd so Start creates the root process in the
+// job and in a new process group. The returned object owns those handles until
+// Close or AfterWait.
+func prepareCommandForStart(cmd *exec.Cmd) (*commandProcess, error) {
+	job, err := windows.CreateJobObject(nil, nil)
+	if err != nil {
+		return nil, err
+	}
+	p := &commandProcess{
+		job: job,
+	}
+	success := false
+	defer func() {
+		if !success {
+			p.Close()
+		}
+	}()
+
+	jobInfo := windows.JOBOBJECT_EXTENDED_LIMIT_INFORMATION{}
+	jobInfo.BasicLimitInformation.LimitFlags = windows.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE
+	if _, err := windows.SetInformationJobObject(
+		job,
+		windows.JobObjectExtendedLimitInformation,
+		uintptr(unsafe.Pointer(&jobInfo)),
+		uint32(unsafe.Sizeof(jobInfo)),
+	); err != nil {
+		return nil, err
+	}
+
+	ioport, err := windows.CreateIoCompletionPort(windows.InvalidHandle, 0, 0, 1)
+	if err != nil {
+		return nil, err
+	}
+	p.ioport = ioport
+
+	port := jobObjectAssociateCompletionPort{
+		CompletionKey:  uintptr(job),
+		CompletionPort: ioport,
+	}
+	if _, err := windows.SetInformationJobObject(
+		job,
+		windows.JobObjectAssociateCompletionPortInformation,
+		uintptr(unsafe.Pointer(&port)),
+		uint32(unsafe.Sizeof(port)),
+	); err != nil {
+		return nil, err
+	}
+
+	var sysProcAttr syscall.SysProcAttr
+	if cmd.SysProcAttr != nil {
+		sysProcAttr = *cmd.SysProcAttr
+	}
+	// Clone Jobs before appending, so preparing the command does not mutate the
+	// backing array of a caller-provided SysProcAttr.
+	sysProcAttr.Jobs = append([]syscall.Handle(nil), sysProcAttr.Jobs...)
+	sysProcAttr.Jobs = append(sysProcAttr.Jobs, syscall.Handle(job))
+	sysProcAttr.CreationFlags |= windows.CREATE_NEW_PROCESS_GROUP
+	cmd.SysProcAttr = &sysProcAttr
+	if cmd.Cancel != nil {
+		cmd.Cancel = p.Cancel
+	}
+
+	success = true
+	return p, nil
+}
+
+// Cancel terminates the root process and all of its descendants through the job
+// to which Start assigned the process atomically.
+func (p *commandProcess) Cancel() error {
+	return p.terminateJob()
+}
+
+// AfterWait runs after the root process has exited. It terminates anything that
+// remains in the job, waits for the active-process-zero notification, and only
+// then releases the job handles so build-directory cleanup cannot race lingering
+// descendants.
+func (p *commandProcess) AfterWait(cmd *exec.Cmd) error {
+	defer p.Close()
+	if err := p.terminateJob(); err != nil && !errors.Is(err, os.ErrProcessDone) {
+		return err
+	}
+	return p.waitForActiveProcessZero()
+}
+
+// Close releases job resources. Once a process has been assigned to the job,
+// callers should prefer AfterWait so descendants are terminated and waited for
+// before these handles are closed.
+func (p *commandProcess) Close() {
+	p.lock.Lock()
+	if p.closed {
+		p.lock.Unlock()
+		return
+	}
+	p.closed = true
+	job := p.job
+	ioport := p.ioport
+	p.job = 0
+	p.ioport = 0
+	p.lock.Unlock()
+
+	if job != 0 {
+		windows.CloseHandle(job)
+	}
+	if ioport != 0 {
+		windows.CloseHandle(ioport)
+	}
+}
+
+func (p *commandProcess) terminateJob() error {
+	p.lock.Lock()
+	closed := p.closed
+	job := p.job
+	p.lock.Unlock()
+	if closed || job == 0 {
+		return os.ErrProcessDone
+	}
+	return windows.TerminateJobObject(job, windowsTerminatedExitCode)
+}
+
+func (p *commandProcess) waitForActiveProcessZero() error {
+	p.lock.Lock()
+	closed := p.closed
+	job := p.job
+	ioport := p.ioport
+	p.lock.Unlock()
+	if closed || job == 0 || ioport == 0 {
+		return nil
+	}
+	for {
+		var completionCode uint32
+		var completionKey uintptr
+		var overlapped *windows.Overlapped
+		if err := windows.GetQueuedCompletionStatus(
+			ioport,
+			&completionCode,
+			&completionKey,
+			&overlapped,
+			windows.INFINITE,
+		); err != nil {
+			return err
+		}
+		if windows.Handle(completionKey) == job && completionCode == jobObjectMsgActiveProcessZero {
+			return nil
+		}
+	}
+}

pkg/runner/local_runner_process_windows_test.go

@@ -0,0 +1,43 @@
+package runner
+
+import (
+	"os/exec"
+	"syscall"
+	"testing"
+
+	"golang.org/x/sys/windows"
+)
+
+func TestPrepareCommandForStartAppendsJob(t *testing.T) {
+	const existingJob syscall.Handle = 123
+	jobs := make([]syscall.Handle, 1, 2)
+	jobs[0] = existingJob
+	cmd := exec.Command("does-not-need-to-exist")
+	originalSysProcAttr := &syscall.SysProcAttr{
+		CreationFlags: windows.CREATE_NO_WINDOW,
+		Jobs:          jobs,
+	}
+	cmd.SysProcAttr = originalSysProcAttr
+
+	commandProcess, err := prepareCommandForStart(cmd)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer commandProcess.Close()
+
+	if cmd.SysProcAttr == originalSysProcAttr {
+		t.Fatal("prepareCommandForStart() reused the caller's SysProcAttr")
+	}
+	if got, want := cmd.SysProcAttr.CreationFlags, uint32(windows.CREATE_NO_WINDOW|windows.CREATE_NEW_PROCESS_GROUP); got != want {
+		t.Errorf("CreationFlags = %#x, want %#x", got, want)
+	}
+	if got, want := cmd.SysProcAttr.Jobs, []syscall.Handle{existingJob, syscall.Handle(commandProcess.job)}; len(got) != len(want) || got[0] != want[0] || got[1] != want[1] {
+		t.Errorf("Jobs = %v, want %v", got, want)
+	}
+	if got, want := originalSysProcAttr.Jobs, []syscall.Handle{existingJob}; len(got) != len(want) || got[0] != want[0] {
+		t.Errorf("original Jobs = %v, want %v", got, want)
+	}
+	if got := jobs[:cap(jobs)][1]; got != 0 {
+		t.Errorf("caller-owned Jobs backing array was modified: jobs[1] = %v", got)
+	}
+}