Merge pull request #212 from circleci/cs/flaky-zombie-test

[ONPREM-2536] Fix child process cleanup on Windows

Author: christian-stephen

CHANGELOG.md

@@ -10,6 +10,7 @@ By following these guidelines, we can easily determine which changes should be i
 
 ## Edge
 
+- [#212](https://github.com/circleci/runner-init/pull/212) Fix child process cleanup on Windows using job objects. This ensures that child processes are destroyed when the parent process (task-agent) terminates.
 - [#197](https://github.com/circleci/runner-init/pull/197) Fix `%PATH%` on Windows by using the OS-specific path list separator.
 - [#133](https://github.com/circleci/runner-init/pull/133) Don't re-handle task errors. If GOAT handles a task error (either with an infra-fail or retry), don't exit with a nonzero status code. Doing so causes container agent to overwrite the original error message in the UI.
 - [#98](https://github.com/circleci/runner-init/pull/98) [INTERNAL] A small refactor to the builds and Dockerfiles in preparation for adding Windows support.

go.mod

@@ -623,7 +623,7 @@ require (
 	go.opentelemetry.io/proto/otlp v1.7.1 // indirect
 	golang.org/x/net v0.43.0 // indirect
 	golang.org/x/sync v0.16.0 // indirect
-	golang.org/x/sys v0.35.0 // indirect
+	golang.org/x/sys v0.35.0
 	golang.org/x/text v0.28.0 // indirect
 	google.golang.org/genproto/googleapis/api v0.0.0-20250728155136-f173205681a0 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20250728155136-f173205681a0 // indirect

task/cmd/cmd.go

@@ -33,7 +33,7 @@ func New(ctx context.Context, cmd []string, forwardSignals bool, user string, en
 func (c *Command) Start() error {
 	cmd := c.cmd
 
-	if err := cmd.Start(); err != nil {
+	if err := c.start(); err != nil {
 		return err
 	}
 

task/cmd/cmd_unix.go

@@ -49,3 +49,7 @@ func additionalSetup(_ context.Context, cmd *exec.Cmd) {
 		return syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
 	}
 }
+
+func (c *Command) start() error {
+	return c.cmd.Start()
+}

task/cmd/cmd_windows.go

@@ -5,15 +5,20 @@ import (
 	"os"
 	"os/exec"
 	"os/signal"
+	"sync"
+	"sync/atomic"
+	"unsafe"
 
 	"github.com/circleci/ex/o11y"
+	"golang.org/x/sys/windows"
 )
 
 func forwardSignals(cmd *exec.Cmd) {
 	ch := make(chan os.Signal, 1)
 	signal.Notify(ch, os.Interrupt)
 
 	go func() {
+		defer signal.Stop(ch)
 		for range ch {
 			_ = cmd.Process.Kill()
 		}
@@ -24,4 +29,67 @@ func switchUser(ctx context.Context, _ *exec.Cmd, user string) {
 	o11y.Log(ctx, "switching users is unsupported on windows", o11y.Field("username", user))
 }
 
-func additionalSetup(_ context.Context, _ *exec.Cmd) {}
+func additionalSetup(_ context.Context, cmd *exec.Cmd) {
+	cmd.SysProcAttr.CreationFlags = windows.CREATE_NEW_PROCESS_GROUP
+}
+
+func (c *Command) start() error {
+	g, err := newProcessExitGroup()
+	if err != nil {
+		return err
+	}
+
+	c.cmd.Cancel = func() error {
+		return g.Dispose()
+	}
+
+	if err := c.cmd.Start(); err != nil {
+		return err
+	}
+
+	return g.AddProcess(c.cmd.Process)
+}
+
+// Process struct matches os.Process layout to extract handle via unsafe pointer
+// https://stackoverflow.com/a/56739249
+type process struct {
+	Pid    int
+	_      uint8
+	_      atomic.Uint64
+	_      sync.RWMutex
+	Handle uintptr
+}
+
+type processExitGroup windows.Handle
+
+func newProcessExitGroup() (processExitGroup, error) {
+	handle, err := windows.CreateJobObject(nil, nil)
+	if err != nil {
+		return 0, err
+	}
+
+	info := windows.JOBOBJECT_EXTENDED_LIMIT_INFORMATION{
+		BasicLimitInformation: windows.JOBOBJECT_BASIC_LIMIT_INFORMATION{
+			LimitFlags: windows.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE,
+		},
+	}
+	if _, err := windows.SetInformationJobObject(
+		handle,
+		windows.JobObjectExtendedLimitInformation,
+		uintptr(unsafe.Pointer(&info)),
+		uint32(unsafe.Sizeof(info))); err != nil {
+		return 0, err
+	}
+
+	return processExitGroup(handle), nil
+}
+
+func (g processExitGroup) Dispose() error {
+	return windows.CloseHandle(windows.Handle(g))
+}
+
+func (g processExitGroup) AddProcess(p *os.Process) error {
+	return windows.AssignProcessToJobObject(
+		windows.Handle(g),
+		windows.Handle((*process)(unsafe.Pointer(p)).Handle))
+}

task/orchestrator_test.go

@@ -2,6 +2,7 @@ package task
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"net/http/httptest"
@@ -111,23 +112,26 @@ func TestOrchestrator(t *testing.T) {
 					pid, err := strconv.Atoi(strings.TrimSpace(string(b)))
 					assert.NilError(t, err)
 
+					time.Sleep(20 * time.Second)
+
 					check := func(poll.LogT) poll.Result {
+						t.Logf("Checking for PID: %d", pid)
 						p, err := os.FindProcess(pid)
 						if runtime.GOOS == "windows" {
-							assert.Check(t, cmp.Nil(p))
-						} else {
-							assert.NilError(t, err)
-
+							if err != nil {
+								assert.Check(t, cmp.ErrorContains(err, "OpenProcess: The parameter is incorrect."),
+									"an error implies the process is dead")
+								return poll.Success()
+							}
+						} else if err == nil {
 							err = p.Signal(syscall.Signal(0))
-							assert.Check(t, cmp.ErrorIs(err, os.ErrProcessDone))
-						}
-
-						if t.Failed() {
-							return poll.Continue("process checks failed")
+							if errors.Is(err, os.ErrProcessDone) {
+								return poll.Success()
+							}
 						}
-						return poll.Success()
+						return poll.Continue("process checks failed")
 					}
-					poll.WaitOn(t, check, poll.WithTimeout(1*time.Minute))
+					poll.WaitOn(t, check, poll.WithTimeout(2*time.Minute))
 				},
 			},
 		},
@@ -402,7 +406,13 @@ func beFakeTaskAgent(t *testing.T) {
 	}
 
 	if pidfile := os.Getenv("SIMULATE_A_ZOMBIE_PROCESS"); pidfile != "" {
-		c := exec.Command(shell(t), "-c", "echo $$ >"+pidfile+" && sleep 300") //nolint:gosec // this is a test
+		pidCmd := "echo $$ >"
+		if runtime.GOOS == "windows" {
+			// https://stackoverflow.com/a/77115597
+			pidCmd = "ps -p $$ | awk 'NR ==2{print $4}' >"
+		}
+		pidCmd += pidfile + " && sleep 300"
+		c := exec.Command(shell(t), "-c", pidCmd) //nolint:gosec // this is a test
 		assert.NilError(t, c.Start())
 	}