diff --git a/agent/reaper/reaper.go b/agent/reaper/reaper.go index d968937a3a..58e2196dc3 100644 --- a/agent/reaper/reaper.go +++ b/agent/reaper/reaper.go @@ -42,9 +42,20 @@ func WithLogger(logger slog.Logger) Option { } } +// WithDone sets a channel that, when closed, stops the reaper +// goroutine. Callers that invoke ForkReap more than once in the +// same process (e.g. tests) should use this to prevent goroutine +// accumulation. +func WithDone(ch chan struct{}) Option { + return func(o *options) { + o.Done = ch + } +} + type options struct { ExecArgs []string PIDs reap.PidCh CatchSignals []os.Signal Logger slog.Logger + Done chan struct{} } diff --git a/agent/reaper/reaper_test.go b/agent/reaper/reaper_test.go index 7ef3f0a50b..f8d4b32a9c 100644 --- a/agent/reaper/reaper_test.go +++ b/agent/reaper/reaper_test.go @@ -18,6 +18,15 @@ import ( "github.com/coder/coder/v2/testutil" ) +// withDone returns an option that stops the reaper goroutine when t +// completes, preventing goroutine accumulation across subtests. +func withDone(t *testing.T) reaper.Option { + t.Helper() + done := make(chan struct{}) + t.Cleanup(func() { close(done) }) + return reaper.WithDone(done) +} + // TestReap checks that's the reaper is successfully reaping // exited processes and passing the PIDs through the shared // channel. @@ -36,6 +45,7 @@ func TestReap(t *testing.T) { reaper.WithPIDCallback(pids), // Provide some argument that immediately exits. reaper.WithExecArgs("/bin/sh", "-c", "exit 0"), + withDone(t), ) require.NoError(t, err) require.Equal(t, 0, exitCode) @@ -89,6 +99,7 @@ func TestForkReapExitCodes(t *testing.T) { t.Run(tt.name, func(t *testing.T) { exitCode, err := reaper.ForkReap( reaper.WithExecArgs("/bin/sh", "-c", tt.command), + withDone(t), ) require.NoError(t, err) require.Equal(t, tt.expectedCode, exitCode, "exit code mismatch for %q", tt.command) @@ -118,6 +129,7 @@ func TestReapInterrupt(t *testing.T) { exitCode, err := reaper.ForkReap( reaper.WithPIDCallback(pids), reaper.WithCatchSignals(os.Interrupt), + withDone(t), // Signal propagation does not extend to children of children, so // we create a little bash script to ensure sleep is interrupted. reaper.WithExecArgs("/bin/sh", "-c", fmt.Sprintf("pid=0; trap 'kill -USR2 %d; kill -TERM $pid' INT; sleep 10 &\npid=$!; kill -USR1 %d; wait", os.Getpid(), os.Getpid())), diff --git a/agent/reaper/reaper_unix.go b/agent/reaper/reaper_unix.go index b095c5a7f9..dd73fc5537 100644 --- a/agent/reaper/reaper_unix.go +++ b/agent/reaper/reaper_unix.go @@ -64,7 +64,7 @@ func ForkReap(opt ...Option) (int, error) { o(opts) } - go reap.ReapChildren(opts.PIDs, nil, nil, nil) + go reap.ReapChildren(opts.PIDs, nil, opts.Done, nil) pwd, err := os.Getwd() if err != nil {