fix: cleanup reaper implementation (#2563)

sreya · web-flow · commit ee5918217b0e · 2022-06-21T18:01:34.000-05:00
- Clean up the agent/reaper API to be a more isolated and reusable package.
diff --git a/agent/reaper/reaper.go b/agent/reaper/reaper.go
@@ -0,0 +1,28 @@
+package reaper
+
+import "github.com/hashicorp/go-reap"
+
+type Option func(o *options)
+
+// WithExecArgs specifies the exec arguments for the fork exec call.
+// By default the same arguments as the parent are used as dictated by
+// os.Args. Since ForkReap calls a fork-exec it is the responsibility of
+// the caller to avoid fork-bombing oneself.
+func WithExecArgs(args ...string) Option {
+	return func(o *options) {
+		o.ExecArgs = args
+	}
+}
+
+// WithPIDCallback sets the channel that reaped child process PIDs are pushed
+// onto.
+func WithPIDCallback(ch reap.PidCh) Option {
+	return func(o *options) {
+		o.PIDs = ch
+	}
+}
+
+type options struct {
+	ExecArgs []string
+	PIDs     reap.PidCh
+}
diff --git a/agent/reaper/reaper_stub.go b/agent/reaper/reaper_stub.go
@@ -2,18 +2,11 @@
 
 package reaper
 
-import "github.com/hashicorp/go-reap"
-
-// IsChild returns true if we're the forked process.
-func IsChild() bool {
-	return false
-}
-
 // IsInitProcess returns true if the current process's PID is 1.
 func IsInitProcess() bool {
 	return false
 }
 
-func ForkReap(_ reap.PidCh) error {
+func ForkReap(opt ...Option) error {
 	return nil
 }
diff --git a/agent/reaper/reaper_test.go b/agent/reaper/reaper_test.go
@@ -24,17 +24,16 @@ func TestReap(t *testing.T) {
 		t.Skip("Detected CI, skipping reaper tests")
 	}
 
-	// Because we're forkexecing these tests will try to run twice...
-	if reaper.IsChild() {
-		t.Skip("I'm a child!")
-	}
-
 	// OK checks that's the reaper is successfully reaping
 	// exited processes and passing the PIDs through the shared
 	// channel.
 	t.Run("OK", func(t *testing.T) {
 		pids := make(reap.PidCh, 1)
-		err := reaper.ForkReap(pids)
+		err := reaper.ForkReap(
+			reaper.WithPIDCallback(pids),
+			// Provide some argument that immediately exits.
+			reaper.WithExecArgs("/bin/sh", "-c", "exit 0"),
+		)
 		require.NoError(t, err)
 
 		cmd := exec.Command("tail", "-f", "/dev/null")
diff --git a/agent/reaper/reaper_unix.go b/agent/reaper/reaper_unix.go
@@ -3,25 +3,13 @@
 package reaper
 
 import (
-	"fmt"
 	"os"
 	"syscall"
 
 	"github.com/hashicorp/go-reap"
 	"golang.org/x/xerrors"
 )
 
-// agentEnvMark is a simple environment variable that we use as a marker
-// to indicated that the process is a child as opposed to the reaper.
-// Since we are forkexec'ing we need to be able to differentiate between
-// the two to avoid fork bombing ourselves.
-const agentEnvMark = "CODER_DO_NOT_REAP"
-
-// IsChild returns true if we're the forked process.
-func IsChild() bool {
-	return os.Getenv(agentEnvMark) != ""
-}
-
 // IsInitProcess returns true if the current process's PID is 1.
 func IsInitProcess() bool {
 	return os.Getpid() == 1
@@ -33,19 +21,16 @@ func IsInitProcess() bool {
 // the reaper and an exec.Command waiting for its process to complete.
 // The provided 'pids' channel may be nil if the caller does not care about the
 // reaped children PIDs.
-func ForkReap(pids reap.PidCh) error {
-	// Check if the process is the parent or the child.
-	// If it's the child we want to skip attempting to reap.
-	if IsChild() {
-		return nil
+func ForkReap(opt ...Option) error {
+	opts := &options{
+		ExecArgs: os.Args,
 	}
 
-	go reap.ReapChildren(pids, nil, nil, nil)
+	for _, o := range opt {
+		o(opts)
+	}
 
-	args := os.Args
-	// This is simply done to help identify the real agent process
-	// when viewing in something like 'ps'.
-	args = append(args, "#Agent")
+	go reap.ReapChildren(opts.PIDs, nil, nil, nil)
 
 	pwd, err := os.Getwd()
 	if err != nil {
@@ -54,8 +39,7 @@ func ForkReap(pids reap.PidCh) error {
 
 	pattrs := &syscall.ProcAttr{
 		Dir: pwd,
-		// Add our marker for identifying the child process.
-		Env: append(os.Environ(), fmt.Sprintf("%s=true", agentEnvMark)),
+		Env: os.Environ(),
 		Sys: &syscall.SysProcAttr{
 			Setsid: true,
 		},
@@ -67,7 +51,7 @@ func ForkReap(pids reap.PidCh) error {
 	}
 
 	//#nosec G204
-	pid, _ := syscall.ForkExec(args[0], args, pattrs)
+	pid, _ := syscall.ForkExec(opts.ExecArgs[0], opts.ExecArgs, pattrs)
 
 	var wstatus syscall.WaitStatus
 	_, err = syscall.Wait4(pid, &wstatus, 0, nil)
diff --git a/cli/agent.go b/cli/agent.go
@@ -32,6 +32,7 @@ func workspaceAgent() *cobra.Command {
 		auth         string
 		pprofEnabled bool
 		pprofAddress string
+		noReap       bool
 	)
 	cmd := &cobra.Command{
 		Use: "agent",
@@ -58,9 +59,12 @@ func workspaceAgent() *cobra.Command {
 
 			// Spawn a reaper so that we don't accumulate a ton
 			// of zombie processes.
-			if reaper.IsInitProcess() && !reaper.IsChild() && isLinux {
+			if reaper.IsInitProcess() && !noReap && isLinux {
 				logger.Info(cmd.Context(), "spawning reaper process")
-				err := reaper.ForkReap(nil)
+				// Do not start a reaper on the child process. It's important
+				// to do this else we fork bomb ourselves.
+				args := append(os.Args, "--no-reap")
+				err := reaper.ForkReap(reaper.WithExecArgs(args...))
 				if err != nil {
 					logger.Error(cmd.Context(), "failed to reap", slog.Error(err))
 					return xerrors.Errorf("fork reap: %w", err)
@@ -182,6 +186,7 @@ func workspaceAgent() *cobra.Command {
 
 	cliflag.StringVarP(cmd.Flags(), &auth, "auth", "", "CODER_AGENT_AUTH", "token", "Specify the authentication type to use for the agent")
 	cliflag.BoolVarP(cmd.Flags(), &pprofEnabled, "pprof-enable", "", "CODER_AGENT_PPROF_ENABLE", false, "Enable serving pprof metrics on the address defined by --pprof-address.")
+	cliflag.BoolVarP(cmd.Flags(), &noReap, "no-reap", "", "", false, "Do not start a process reaper.")
 	cliflag.StringVarP(cmd.Flags(), &pprofAddress, "pprof-address", "", "CODER_AGENT_PPROF_ADDRESS", "127.0.0.1:6060", "The address to serve pprof.")
 	return cmd
 }
