fix: Add reaper to coder agent

sreya · sreya · commit 99ab87e1b72c · 2022-06-16T20:27:32.000Z
- The coder agent runs as PID 1 in some of our Docker workspaces.
  In such cases it is the responsibility of the init process to
  reap dead processes. Failing to do so can result in an inability
  to create new processes by running out of PIDs.

  This PR adds a reaper to our agent that is only spawned if it
  detects that it is PID1.
diff --git a/agent/reaper/reaper.go b/agent/reaper/reaper.go
@@ -0,0 +1,68 @@
+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_AGENT"
+
+// IsChild returns true if we're the forked process.
+func IsChild() bool {
+	return os.Getenv(agentEnvMark) != ""
+}
+
+// ForkReap spawns a goroutine that reaps children. In order to avoid
+// complications with spawning `exec.Commands` in the same process that
+// is reaping, we forkexec a child process. This prevents a race between
+// the reaper and an exec.Command waiting for its process to complete.
+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() {
+		go reap.ReapChildren(pids, nil, nil, nil)
+
+		args := os.Args
+		// This is simply done to help identify the real agent process
+		// when viewing in something like 'ps'.
+		args = append(args, "#Agent")
+
+		pwd, err := os.Getwd()
+		if err != nil {
+			return xerrors.Errorf("get wd: %w", err)
+		}
+
+		pattrs := &syscall.ProcAttr{
+			Dir: pwd,
+			// Add our marker for identifying the child process.
+			Env: append(os.Environ(), fmt.Sprintf("%s=true", agentEnvMark)),
+			Sys: &syscall.SysProcAttr{
+				Setsid: true,
+			},
+			Files: []uintptr{
+				uintptr(syscall.Stdin),
+				uintptr(syscall.Stdout),
+				uintptr(syscall.Stderr),
+			},
+		}
+
+		pid, _ := syscall.ForkExec(args[0], args, pattrs)
+
+		var wstatus syscall.WaitStatus
+		_, err = syscall.Wait4(pid, &wstatus, 0, nil)
+		for xerrors.Is(err, syscall.EINTR) {
+			_, err = syscall.Wait4(pid, &wstatus, 0, nil)
+		}
+		return nil
+	}
+
+	return nil
+}
diff --git a/agent/reaper/reaper_test.go b/agent/reaper/reaper_test.go
@@ -0,0 +1,51 @@
+package reaper_test
+
+import (
+	"os/exec"
+	"testing"
+	"time"
+
+	"github.com/hashicorp/go-reap"
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/agent/reaper"
+)
+
+func TestReap(t *testing.T) {
+	t.Parallel()
+
+	// Because we're forkexecing these tests will try to run twice...
+	if reaper.IsChild() {
+		t.Skip("I'm a child!")
+	}
+
+	t.Run("OK", func(t *testing.T) {
+		pids := make(reap.PidCh, 1)
+		err := reaper.ForkReap(pids)
+		require.NoError(t, err)
+
+		cmd := exec.Command("sleep", "5")
+		err = cmd.Start()
+		require.NoError(t, err)
+
+		cmd2 := exec.Command("sleep", "5")
+		err = cmd2.Start()
+		require.NoError(t, err)
+
+		err = cmd.Process.Kill()
+		require.NoError(t, err)
+
+		err = cmd2.Process.Kill()
+		require.NoError(t, err)
+
+		expectedPIDs := []int{cmd.Process.Pid, cmd2.Process.Pid}
+
+		deadline := time.NewTimer(time.Second * 5)
+		select {
+		case <-deadline.C:
+			t.Fatalf("Timed out waiting for process")
+		case pid := <-pids:
+			require.Contains(t, expectedPIDs, pid)
+		}
+	})
+}
diff --git a/cli/agent.go b/cli/agent.go
@@ -17,6 +17,7 @@ import (
 	"cdr.dev/slog/sloggers/sloghuman"
 
 	"github.com/coder/coder/agent"
+	"github.com/coder/coder/agent/reaper"
 	"github.com/coder/coder/cli/cliflag"
 	"github.com/coder/coder/codersdk"
 	"github.com/coder/retry"
@@ -35,6 +36,16 @@ func workspaceAgent() *cobra.Command {
 		// This command isn't useful to manually execute.
 		Hidden: true,
 		RunE: func(cmd *cobra.Command, args []string) error {
+			// Spawn a reaper so that we don't accumulate a ton
+			// of zombie processes.
+			if !reaper.IsChild() {
+				err := reaper.ForkReap(nil)
+				if err != nil {
+					return xerrors.Errorf("fork reap: %w", err)
+				}
+				return nil
+			}
+
 			rawURL, err := cmd.Flags().GetString(varAgentURL)
 			if err != nil {
 				return xerrors.Errorf("CODER_AGENT_URL must be set: %w", err)
diff --git a/go.mod b/go.mod
@@ -134,7 +134,10 @@ require (
 	storj.io/drpc v0.0.30
 )
 
-require github.com/xi2/xz v0.0.0-20171230120015-48954b6210f8 // indirect
+require (
+	github.com/hashicorp/go-reap v0.0.0-20170704170343-bf58d8a43e7b // indirect
+	github.com/xi2/xz v0.0.0-20171230120015-48954b6210f8 // indirect
+)
 
 require (
 	github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1 // indirect
diff --git a/go.sum b/go.sum
@@ -899,6 +899,8 @@ github.com/hashicorp/go-multierror v1.0.0/go.mod h1:dHtQlpGsu+cZNNAkkCN/P3hoUDHh
 github.com/hashicorp/go-multierror v1.1.0/go.mod h1:spPvp8C1qA32ftKqdAHm4hHTbPw+vmowP0z+KUhOZdA=
 github.com/hashicorp/go-multierror v1.1.1 h1:H5DkEtf6CXdFp0N0Em5UCwQpXMWke8IA0+lD48awMYo=
 github.com/hashicorp/go-multierror v1.1.1/go.mod h1:iw975J/qwKPdAO1clOe2L8331t/9/fmwbPZ6JB6eMoM=
+github.com/hashicorp/go-reap v0.0.0-20170704170343-bf58d8a43e7b h1:3GrpnZQBxcMj1gCXQLelfjCT1D5MPGTuGMKHVzSIH6A=
+github.com/hashicorp/go-reap v0.0.0-20170704170343-bf58d8a43e7b/go.mod h1:qIFzeFcJU3OIFk/7JreWXcUjFmcCaeHTH9KoNyHYVCs=
 github.com/hashicorp/go-rootcerts v1.0.0/go.mod h1:K6zTfqpRlCUIjkwsN4Z+hiSfzSTQa6eBIzfwKfwNnHU=
 github.com/hashicorp/go-sockaddr v1.0.0/go.mod h1:7Xibr9yA9JjQq1JpNB2Vw7kxv8xerXegt+ozgdvDeDU=
 github.com/hashicorp/go-syslog v1.0.0/go.mod h1:qPfqrKkXGihmCqbJM2mZgkZGvKG1dFdvsLplgctolz4=
