coder · sreya · Nov 25, 2024 · Oct 23, 2024 · Nov 18, 2024 · Nov 18, 2024
diff --git a/agent/agentexec/cli_other.go b/agent/agentexec/cli_other.go
@@ -0,0 +1,10 @@
+//go:build !linux
+// +build !linux
+
+package agentexec
+
+import "golang.org/x/xerrors"
+
+func CLI() error {
+	return xerrors.Errorf("agent-exec is only supported on Linux")
+}
diff --git a/agent/agentexec/cli_unix.go b/agent/agentexec/cli_unix.go
@@ -0,0 +1,149 @@
+//go:build linux
+// +build linux
+
+package agentexec
+
+import (
+	"flag"
+	"fmt"
+	"os"
+	"os/exec"
+	"runtime"
+	"strconv"
+	"strings"
+	"syscall"
+
+	"golang.org/x/sys/unix"
+	"golang.org/x/xerrors"
+)
+
+// unset is set to an invalid value for nice and oom scores.
+const unset = -2000
+
+// CLI runs the agent-exec command. It should only be called by the cli package.
+func CLI() error {
+	// We lock the OS thread here to avoid a race conditino where the nice priority
+	// we get is on a different thread from the one we set it on.
+	runtime.LockOSThread()
+
+	var (
+		fs   = flag.NewFlagSet("agent-exec", flag.ExitOnError)
+		nice = fs.Int("coder-nice", unset, "")
+		oom  = fs.Int("coder-oom", unset, "")
+	)
+
+	if len(os.Args) < 3 {
+		return xerrors.Errorf("malformed command %+v", os.Args)
+	}
+
+	// Parse everything after "coder agent-exec".
+	err := fs.Parse(os.Args[2:])
+	if err != nil {
+		return xerrors.Errorf("parse flags: %w", err)
+	}
+
+	if runtime.GOOS != "linux" {
+		return xerrors.Errorf("agent-exec is only supported on Linux")
+	}
+
+	// Get everything after "coder agent-exec --"
+	args := execArgs(os.Args)
+	if len(args) == 0 {
+		return xerrors.Errorf("no exec command provided %+v", os.Args)
+	}
+
+	pid := os.Getpid()
+
+	if *nice == unset {
+		// If an explicit nice score isn't set, we use the default.
+		*nice, err = defaultNiceScore()
+		if err != nil {
+			return xerrors.Errorf("get default nice score: %w", err)
+		}
+	}
+
+	if *oom == unset {
+		// If an explicit oom score isn't set, we use the default.
+		*oom, err = defaultOOMScore()
+		if err != nil {
+			return xerrors.Errorf("get default oom score: %w", err)
+		}
+	}
+
+	err = unix.Setpriority(unix.PRIO_PROCESS, 0, *nice)
+	if err != nil {
+		return xerrors.Errorf("set nice score: %w", err)
+	}
+
+	err = writeOOMScoreAdj(pid, *oom)
+	if err != nil {
+		return xerrors.Errorf("set oom score: %w", err)
+	}
+
+	path, err := exec.LookPath(args[0])
+	if err != nil {
+		return xerrors.Errorf("look path: %w", err)
+	}
+
+	return syscall.Exec(path, args, os.Environ())
+}
+
+func defaultNiceScore() (int, error) {
+	score, err := unix.Getpriority(unix.PRIO_PROCESS, os.Getpid())
+	if err != nil {
+		return 0, xerrors.Errorf("get nice score: %w", err)
+	}
+	// See https://linux.die.net/man/2/setpriority#Notes
+	score = 20 - score
+
+	score += 5
+	if score > 19 {
+		return 19, nil
+	}
+	return score, nil
+}
+
+func defaultOOMScore() (int, error) {
+	score, err := oomScoreAdj(os.Getpid())
+	if err != nil {
+		return 0, xerrors.Errorf("get oom score: %w", err)
+	}
+
+	// If the agent has a negative oom_score_adj, we set the child to 0
+	// so it's treated like every other process.
+	if score < 0 {
+		return 0, nil
+	}
+
+	// If the agent is already almost at the maximum then set it to the max.
+	if score >= 998 {
+		return 1000, nil
+	}
+
+	// If the agent oom_score_adj is >=0, we set the child to slightly
+	// less than the maximum. If users want a different score they set it
+	// directly.
+	return 998, nil
+}
+
+func oomScoreAdj(pid int) (int, error) {
+	scoreStr, err := os.ReadFile(fmt.Sprintf("/proc/%d/oom_score_adj", pid))
+	if err != nil {
+		return 0, xerrors.Errorf("read oom_score_adj: %w", err)
+	}
+	return strconv.Atoi(strings.TrimSpace(string(scoreStr)))
+}
+
+func writeOOMScoreAdj(pid int, score int) error {
+	return os.WriteFile(fmt.Sprintf("/proc/%d/oom_score_adj", pid), []byte(fmt.Sprintf("%d", score)), 0o600)
+}
+
+// execArgs returns the arguments to pass to syscall.Exec after the "--" delimiter.
+func execArgs(args []string) []string {
+	for i, arg := range args {
+		if arg == "--" {
+			return args[i+1:]
+		}
+	}
+	return nil
+}
diff --git a/agent/agentexec/cli_unix_test.go b/agent/agentexec/cli_unix_test.go
@@ -0,0 +1,178 @@
+//go:build linux
+// +build linux
+
+package agentexec_test
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strconv"
+	"strings"
+	"syscall"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"golang.org/x/sys/unix"
+
+	"github.com/coder/coder/v2/testutil"
+)
+
+func TestCLI(t *testing.T) {
+	t.Parallel()
+
+	t.Run("OK", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := testutil.Context(t, testutil.WaitMedium)
+		cmd, path := cmd(ctx, t, 123, 12)
+		err := cmd.Start()
+		require.NoError(t, err)
+		go cmd.Wait()
+
+		waitForSentinel(ctx, t, cmd, path)
+		requireOOMScore(t, cmd.Process.Pid, 123)
+		requireNiceScore(t, cmd.Process.Pid, 12)
+	})
+
+	t.Run("Defaults", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := testutil.Context(t, testutil.WaitMedium)
+		cmd, path := cmd(ctx, t, 0, 0)
+		err := cmd.Start()
+		require.NoError(t, err)
+		go cmd.Wait()
+
+		waitForSentinel(ctx, t, cmd, path)
+
+		expectedNice := expectedNiceScore(t)
+		expectedOOM := expectedOOMScore(t)
+		requireOOMScore(t, cmd.Process.Pid, expectedOOM)
+		requireNiceScore(t, cmd.Process.Pid, expectedNice)
+	})
+}
+
+func requireNiceScore(t *testing.T, pid int, score int) {
+	t.Helper()
+
+	nice, err := unix.Getpriority(unix.PRIO_PROCESS, pid)
+	require.NoError(t, err)
+	// See https://linux.die.net/man/2/setpriority#Notes
+	require.Equal(t, score, 20-nice)
+}
+
+func requireOOMScore(t *testing.T, pid int, expected int) {
+	t.Helper()
+
+	actual, err := os.ReadFile(fmt.Sprintf("/proc/%d/oom_score_adj", pid))
+	require.NoError(t, err)
+	score := strings.TrimSpace(string(actual))
+	require.Equal(t, strconv.Itoa(expected), score)
+}
+
+func waitForSentinel(ctx context.Context, t *testing.T, cmd *exec.Cmd, path string) {
+	t.Helper()
+
+	ticker := time.NewTicker(testutil.IntervalFast)
+	defer ticker.Stop()
+
+	// RequireEventually doesn't work well with require.NoError or similar require functions.
+	for {
+		err := cmd.Process.Signal(syscall.Signal(0))
+		require.NoError(t, err)
+
+		_, err = os.Stat(path)
+		if err == nil {
+			return
+		}
+
+		select {
+		case <-ticker.C:
+		case <-ctx.Done():
+			require.NoError(t, ctx.Err())
+		}
+	}
+}
+
+func cmd(ctx context.Context, t *testing.T, oom, nice int) (*exec.Cmd, string) {
+	var (
+		args = execArgs(oom, nice)
+		dir  = t.TempDir()
+		file = filepath.Join(dir, "sentinel")
+	)
+
+	args = append(args, "sh", "-c", fmt.Sprintf("touch %s && sleep 10m", file))
+	//nolint:gosec
+	cmd := exec.CommandContext(ctx, TestBin, args...)
+
+	// We set this so we can also easily kill the sleep process the shell spawns.
+	cmd.SysProcAttr = &syscall.SysProcAttr{
+		Setpgid: true,
+	}
+
+	cmd.Env = os.Environ()
+	var buf bytes.Buffer
+	cmd.Stdout = &buf
+	cmd.Stderr = &buf
+	t.Cleanup(func() {
+		// Print output of a command if the test fails.
+		if t.Failed() {
+			t.Logf("cmd %q output: %s", cmd.Args, buf.String())
+		}
+		if cmd.Process != nil {
+			// We use -cmd.Process.Pid to kill the whole process group.
+			_ = syscall.Kill(-cmd.Process.Pid, syscall.SIGINT)
+		}
+	})
+	return cmd, file
+}
+
+func expectedOOMScore(t *testing.T) int {
+	t.Helper()
+
+	score, err := os.ReadFile(fmt.Sprintf("/proc/%d/oom_score_adj", os.Getpid()))
+	require.NoError(t, err)
+
+	scoreInt, err := strconv.Atoi(strings.TrimSpace(string(score)))
+	require.NoError(t, err)
+
+	if scoreInt < 0 {
+		return 0
+	}
+	if scoreInt >= 998 {
+		return 1000
+	}
+	return 998
+}
+
+func expectedNiceScore(t *testing.T) int {
+	t.Helper()
+
+	score, err := unix.Getpriority(unix.PRIO_PROCESS, os.Getpid())
+	require.NoError(t, err)
+
+	// Priority is niceness + 20.
+	score = 20 - score
+	score += 5
+	if score > 19 {
+		return 19
+	}
+	return score
+}
+
+func execArgs(oom int, nice int) []string {
+	execArgs := []string{"agent-exec"}
+	if oom != 0 {
+		execArgs = append(execArgs, fmt.Sprintf("--coder-oom=%d", oom))
+	}
+	if nice != 0 {
+		execArgs = append(execArgs, fmt.Sprintf("--coder-nice=%d", nice))
+	}
+	execArgs = append(execArgs, "--")
+	return execArgs
+}
diff --git a/agent/agentexec/cmdtest/main.go b/agent/agentexec/cmdtest/main.go
@@ -0,0 +1,16 @@
+package main
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/coder/coder/v2/agent/agentexec"
+)
+
+func main() {
+	err := agentexec.CLI()
+	if err != nil {
+		_, _ = fmt.Fprintln(os.Stderr, err)
+		os.Exit(1)
+	}
+}