use flag parsing instead

coder · sreya · Nov 25, 2024 · Oct 23, 2024 · Nov 18, 2024 · Nov 18, 2024
commit 986e18e73f087e060d67f613147197df6b706502
diff --git a/agent/agentexec/cli_unix.go b/agent/agentexec/cli_unix.go
@@ -4,6 +4,7 @@
 package agentexec
 
 import (
+	"flag"
 	"fmt"
 	"os"
 	"os/exec"
@@ -19,8 +20,15 @@ import (
 
 // CLI runs the agent-exec command. It should only be called by the cli package.
 func CLI(args []string, environ []string) 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()
 
+	nice := flag.Int(niceArg, 0, "")
+	oom := flag.Int(oomArg, 0, "")
+
+	flag.Parse()
+
 	if runtime.GOOS != "linux" {
 		return xerrors.Errorf("agent-exec is only supported on Linux")
 	}
@@ -35,30 +43,30 @@ func CLI(args []string, environ []string) error {
 	pid := os.Getpid()
 
 	var err error
-	nice, ok := envValInt(environ, EnvProcNiceScore)
-	if !ok {
+	if nice == nil {
 		// If an explicit nice score isn't set, we use the default.
-		nice, err = defaultNiceScore()
+		n, err := defaultNiceScore()
 		if err != nil {
 			return xerrors.Errorf("get default nice score: %w", err)
 		}
+		nice = &n
 	}
 
-	oomscore, ok := envValInt(environ, EnvProcOOMScore)
-	if !ok {
+	if oom == nil {
 		// If an explicit oom score isn't set, we use the default.
-		oomscore, err = defaultOOMScore()
+		o, err := defaultOOMScore()
 		if err != nil {
 			return xerrors.Errorf("get default oom score: %w", err)
 		}
+		oom = &o
 	}
 
-	err = unix.Setpriority(unix.PRIO_PROCESS, 0, nice)
+	err = unix.Setpriority(unix.PRIO_PROCESS, 0, *nice)
 	if err != nil {
 		return xerrors.Errorf("set nice score: %w", err)
 	}
 
-	err = writeOOMScoreAdj(pid, oomscore)
+	err = writeOOMScoreAdj(pid, *oom)
 	if err != nil {
 		return xerrors.Errorf("set oom score: %w", err)
 	}

diff --git a/agent/agentexec/exec.go b/agent/agentexec/exec.go
@@ -2,12 +2,12 @@ package agentexec
 
 import (
 	"context"
+	"fmt"
 	"os"
 	"os/exec"
 	"path/filepath"
 	"runtime"
 	"strconv"
-	"strings"
 
 	"golang.org/x/xerrors"
 )
@@ -25,8 +25,7 @@ const (
 // is returned. All instances of exec.Cmd should flow through this function to ensure
 // proper resource constraints are applied to the child process.
 func CommandContext(ctx context.Context, cmd string, args ...string) (*exec.Cmd, error) {
-	environ := os.Environ()
-	_, enabled := envVal(environ, EnvProcPrioMgmt)
+	_, enabled := os.LookupEnv(EnvProcPrioMgmt)
 	if runtime.GOOS != "linux" || !enabled {
 		return exec.CommandContext(ctx, cmd, args...), nil
 	}
@@ -41,14 +40,24 @@ func CommandContext(ctx context.Context, cmd string, args ...string) (*exec.Cmd,
 		return nil, xerrors.Errorf("eval symlinks: %w", err)
 	}
 
-	args = append([]string{"agent-exec", cmd}, args...)
-	return exec.CommandContext(ctx, bin, args...), nil
+	execArgs := []string{"agent-exec"}
+	if score, ok := envValInt(EnvProcOOMScore); ok {
+		execArgs = append(execArgs, oomScoreArg(score))
+	}
+
+	if score, ok := envValInt(EnvProcNiceScore); ok {
+		execArgs = append(execArgs, niceScoreArg(score))
+	}
+	execArgs = append(execArgs, "--", cmd)
+	execArgs = append(execArgs, args...)
+
+	return exec.CommandContext(ctx, bin, execArgs...), nil
 }
 
 // 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)
+func envValInt(key string) (int, bool) {
+	val, ok := os.LookupEnv(key)
 	if !ok {
 		return 0, false
 	}
@@ -60,14 +69,15 @@ func envValInt(env []string, key string) (int, bool) {
 	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
+const (
+	niceArg = "coder-nice"
+	oomArg  = "coder-oom"
+)
+
+func niceScoreArg(score int) string {
+	return fmt.Sprintf("--%s=%d", niceArg, score)
+}
+
+func oomScoreArg(score int) string {
+	return fmt.Sprintf("--%s=%d", oomArg, score)
 }
diff --git a/agent/agentexec/exec_test.go b/agent/agentexec/exec_test.go
@@ -61,7 +61,59 @@ func TestExec(t *testing.T) {
 			cmd, err := agentexec.CommandContext(context.Background(), "sh", "-c", "sleep")
 			require.NoError(t, err)
 			require.Equal(t, executable, cmd.Path)
-			require.Equal(t, []string{executable, "agent-exec", "sh", "-c", "sleep"}, cmd.Args)
+			require.Equal(t, []string{executable, "agent-exec", "--", "sh", "-c", "sleep"}, cmd.Args)
+		})
+
+		t.Run("Nice", func(t *testing.T) {
+			t.Setenv(agentexec.EnvProcPrioMgmt, "hello")
+			t.Setenv(agentexec.EnvProcNiceScore, "10")
+
+			if runtime.GOOS != "linux" {
+				t.Skip("skipping on linux")
+			}
+
+			executable, err := os.Executable()
+			require.NoError(t, err)
+
+			cmd, err := agentexec.CommandContext(context.Background(), "sh", "-c", "sleep")
+			require.NoError(t, err)
+			require.Equal(t, executable, cmd.Path)
+			require.Equal(t, []string{executable, "agent-exec", "--coder-nice=10", "--", "sh", "-c", "sleep"}, cmd.Args)
+		})
+
+		t.Run("OOM", func(t *testing.T) {
+			t.Setenv(agentexec.EnvProcPrioMgmt, "hello")
+			t.Setenv(agentexec.EnvProcOOMScore, "123")
+
+			if runtime.GOOS != "linux" {
+				t.Skip("skipping on linux")
+			}
+
+			executable, err := os.Executable()
+			require.NoError(t, err)
+
+			cmd, err := agentexec.CommandContext(context.Background(), "sh", "-c", "sleep")
+			require.NoError(t, err)
+			require.Equal(t, executable, cmd.Path)
+			require.Equal(t, []string{executable, "agent-exec", "--coder-oom=123", "--", "sh", "-c", "sleep"}, cmd.Args)
+		})
+
+		t.Run("Both", func(t *testing.T) {
+			t.Setenv(agentexec.EnvProcPrioMgmt, "hello")
+			t.Setenv(agentexec.EnvProcOOMScore, "432")
+			t.Setenv(agentexec.EnvProcNiceScore, "14")
+
+			if runtime.GOOS != "linux" {
+				t.Skip("skipping on linux")
+			}
+
+			executable, err := os.Executable()
+			require.NoError(t, err)
+
+			cmd, err := agentexec.CommandContext(context.Background(), "sh", "-c", "sleep")
+			require.NoError(t, err)
+			require.Equal(t, executable, cmd.Path)
+			require.Equal(t, []string{executable, "agent-exec", "--coder-oom=432", "--coder-nice=14", "--", "sh", "-c", "sleep"}, cmd.Args)
 		})
 	})
 }