idk

coder · sreya · Nov 25, 2024 · Oct 23, 2024 · Nov 18, 2024 · Nov 18, 2024
commit 97e68f410dc8f09dbba8917f0d79ca380c2b361f
diff --git a/agent/agentexec/cli.go b/agent/agentexec/cli.go
@@ -10,6 +10,7 @@
 	"strings"
 	"syscall"
 
+	"golang.org/x/sys/unix"
 	"golang.org/x/xerrors"
 )
 
@@ -28,32 +29,37 @@
 		return xerrors.Errorf("malformed command %q", args)
 	}
 
+	// Slice off 'coder agent-exec'
 	args = args[2:]
 
 	pid := os.Getpid()
 
-	oomScore, ok := envVal(environ, EnvProcOOMScore)
+	var err error
+	nice, ok := envValInt(environ, EnvProcNiceScore)
 	if !ok {
-		return xerrors.Errorf("missing %q", EnvProcOOMScore)
+		// 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)
+		}
+		fmt.Println("nice score", nice, "pid", pid)
 	}
 
-	niceScore, ok := envVal(environ, EnvProcNiceScore)
+	oomscore, ok := envValInt(environ, EnvProcOOMScore)
 	if !ok {
-		return xerrors.Errorf("missing %q", EnvProcNiceScore)
-	}
-
-	score, err := strconv.Atoi(niceScore)
-	if err != nil {
-		return xerrors.Errorf("invalid nice score: %w", err)
+		// If an explicit oom score isn't set, we use the default.
+		oomscore, err = defaultOOMScore()
+		if err != nil {
+			return xerrors.Errorf("get default oom score: %w", err)
+		}
 	}
 
-	err = syscall.Setpriority(syscall.PRIO_PROCESS, pid, score)
+	err = unix.Setpriority(unix.PRIO_PROCESS, pid, nice)
 	if err != nil {
 		return xerrors.Errorf("set nice score: %w", err)
 	}
 
-	oomPath := fmt.Sprintf("/proc/%d/oom_score_adj", pid)
-	err = os.WriteFile(oomPath, []byte(oomScore), 0o600)
+	err = writeOOMScoreAdj(pid, oomscore)
 	if err != nil {
 		return xerrors.Errorf("set oom score: %w", err)
 	}
@@ -63,17 +69,86 @@
 		return xerrors.Errorf("look path: %w", err)
 	}
 
+	// Remove environments variables specifically set for the agent-exec command.
 	env := slices.DeleteFunc(environ, func(env string) bool {
 		return strings.HasPrefix(env, EnvProcOOMScore) || strings.HasPrefix(env, EnvProcNiceScore)
 	})
 
 	return syscall.Exec(path, args, env)
 }
 
-func envVal(environ []string, key string) (string, bool) {
-	for _, env := range environ {
-		if strings.HasPrefix(env, key+"=") {
-			return strings.TrimPrefix(env, key+"="), true
+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)
+	}
+	// Priority is niceness + 20.
+	score -= 20
+
+	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)
+}
+
+// envValInt searches for a key in a list of environment variables and parses it to an int.
+// If the key is not found or cannot be parsed, returns 0 and false.
+func envValInt(env []string, key string) (int, bool) {
+	val, ok := envVal(env, key)
+	if !ok {
+		return 0, false
+	}
+
+	i, err := strconv.Atoi(val)
+	if err != nil {
+		return 0, false
+	}
+	return i, true
+}
+
+// envVal searches for a key in a list of environment variables and returns its value.
+// If the key is not found, returns empty string and false.
+func envVal(env []string, key string) (string, bool) {
+	prefix := key + "="
+	for _, e := range env {
+		if strings.HasPrefix(e, prefix) {
+			return strings.TrimPrefix(e, prefix), true
 		}
 	}
 	return "", false

diff --git a/agent/agentexec/cli_test.go b/agent/agentexec/cli_test.go
@@ -32,16 +32,34 @@
 		require.NoError(t, err)
 		go cmd.Wait()
 
-		waitForSentinel(t, ctx, cmd, path)
+		waitForSentinel(ctx, t, cmd, path)
 		requireOOMScore(t, cmd.Process.Pid, 123)
 		requireNiceScore(t, cmd.Process.Pid, 10)
 	})
+
+	t.Run("Defaults", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := testutil.Context(t, testutil.WaitMedium)
+		cmd, path := cmd(ctx, t)
+		err := cmd.Start()
+		require.NoError(t, err)
+		go cmd.Wait()
+
+		waitForSentinel(ctx, t, cmd, path)
+
+		expectedNice := expectedNiceScore(t)
+		expectedOOM := expectedOOMScore(t)
+		fmt.Println("expected nice", expectedNice, "expected oom", expectedOOM)
+		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(0, pid)
+	nice, err := unix.Getpriority(unix.PRIO_PROCESS, pid)
 	require.NoError(t, err)
 	require.Equal(t, score, nice)
 }
@@ -55,7 +73,7 @@
 	require.Equal(t, strconv.Itoa(expected), score)
 }
 
-func waitForSentinel(t *testing.T, ctx context.Context, cmd *exec.Cmd, path string) {
+func waitForSentinel(ctx context.Context, t *testing.T, cmd *exec.Cmd, path string) {
 	t.Helper()
 
 	ticker := time.NewTicker(testutil.IntervalFast)
@@ -74,17 +92,20 @@
 		select {
 		case <-ticker.C:
 		case <-ctx.Done():
-			return
+			require.NoError(t, ctx.Err())
 		}
 	}
 }
 
 func cmd(ctx context.Context, t *testing.T, args ...string) (*exec.Cmd, string) {
 	file := ""
-	cmd := exec.Command(TestBin, args...)
+	//nolint:gosec
+	cmd := exec.Command(TestBin, append([]string{"agent-exec"}, args...)...)
 	if len(args) == 0 {
+		// Generate a unique path that we can touch to indicate that we've progressed past the
+		// syscall.Exec.
 		dir := t.TempDir()
-		file = filepath.Join(dir, uniqueFile(t))
+		file = filepath.Join(dir, "sentinel")
 		//nolint:gosec
 		cmd = exec.CommandContext(ctx, TestBin, "agent-exec", "sh", "-c", fmt.Sprintf("touch %s && sleep 10m", file))
 	}
@@ -98,13 +119,42 @@
 			t.Logf("cmd %q output: %s", cmd.Args, buf.String())
 		}
 
-		if cmd.Process != nil {
-			_ = cmd.Process.Kill()
-		}
+		// if cmd.Process != nil {
+		// 	_ = cmd.Process.Kill()
+		// }
 	})
 	return cmd, file
 }
 
-func uniqueFile(t *testing.T) string {
-	return fmt.Sprintf("%s-%d", strings.ReplaceAll(t.Name(), "/", "_"), time.Now().UnixNano())
+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 += 5
+	if score > 19 {
+		return 19
+	}
+	return score
 }