more tests

coder · sreya · Nov 25, 2024 · Oct 23, 2024 · Nov 18, 2024 · Nov 18, 2024
commit 712b328fde102aa0cf098f8ee8e9f77dba5a7615
diff --git a/agent/agentexec/cli.go b/agent/agentexec/cli.go
@@ -21,6 +21,8 @@
 
 // CLI runs the agent-exec command. It should only be called by the cli package.
 func CLI(args []string, environ []string) error {
+	runtime.LockOSThread()
+
 	if runtime.GOOS != "linux" {
 		return xerrors.Errorf("agent-exec is only supported on Linux")
 	}
@@ -42,7 +44,6 @@
 		if err != nil {
 			return xerrors.Errorf("get default nice score: %w", err)
 		}
-		fmt.Println("nice score", nice, "pid", pid)
 	}
 
 	oomscore, ok := envValInt(environ, EnvProcOOMScore)
@@ -54,7 +55,7 @@
 		}
 	}
 
-	err = unix.Setpriority(unix.PRIO_PROCESS, pid, nice)
+	err = syscall.Setpriority(syscall.PRIO_PROCESS, 0, nice)
 	if err != nil {
 		return xerrors.Errorf("set nice score: %w", err)
 	}
@@ -78,12 +79,12 @@
 }

 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
+	// See https://linux.die.net/man/2/setpriority#Notes
+	score = 20 - score
 
 	score += 5
 	if score > 19 {

diff --git a/agent/agentexec/cli_test.go b/agent/agentexec/cli_test.go
@@ -26,15 +26,15 @@ func TestCLI(t *testing.T) {
 
 		ctx := testutil.Context(t, testutil.WaitMedium)
 		cmd, path := cmd(ctx, t)
-		cmd.Env = append(cmd.Env, "CODER_PROC_NICE_SCORE=10")
+		cmd.Env = append(cmd.Env, "CODER_PROC_NICE_SCORE=12")
 		cmd.Env = append(cmd.Env, "CODER_PROC_OOM_SCORE=123")
 		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, 10)
+		requireNiceScore(t, cmd.Process.Pid, 12)
 	})
 
 	t.Run("Defaults", func(t *testing.T) {
@@ -50,7 +50,6 @@ func TestCLI(t *testing.T) {
 
 		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)
 	})
@@ -61,7 +60,8 @@ func requireNiceScore(t *testing.T, pid int, score int) {
 
 	nice, err := unix.Getpriority(unix.PRIO_PROCESS, pid)
 	require.NoError(t, err)
-	require.Equal(t, score, nice)
+	// See https://linux.die.net/man/2/setpriority#Notes
+	require.Equal(t, score, 20-nice)
 }
 
 func requireOOMScore(t *testing.T, pid int, expected int) {
@@ -108,6 +108,10 @@ func cmd(ctx context.Context, t *testing.T, args ...string) (*exec.Cmd, string)
 		file = filepath.Join(dir, "sentinel")
 		//nolint:gosec
 		cmd = exec.CommandContext(ctx, TestBin, "agent-exec", "sh", "-c", fmt.Sprintf("touch %s && sleep 10m", file))
+		// 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
@@ -118,10 +122,10 @@ func cmd(ctx context.Context, t *testing.T, args ...string) (*exec.Cmd, string)
 		if t.Failed() {
 			t.Logf("cmd %q output: %s", cmd.Args, buf.String())
 		}
-
-		// if cmd.Process != nil {
-		// 	_ = cmd.Process.Kill()
-		// }
+		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
 }
@@ -151,7 +155,7 @@ func expectedNiceScore(t *testing.T) int {
 	require.NoError(t, err)
 
 	// Priority is niceness + 20.
-	score -= 20
+	score = 20 - score
 	score += 5
 	if score > 19 {
 		return 19

diff --git a/agent/agentexec/exec.go b/agent/agentexec/exec.go
@@ -20,8 +20,9 @@ const (
 // the provided command if CODER_PROC_PRIO_MGMT is set, otherwise a normal exec.Cmd
 // 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) {
-	_, enabled := envVal(os.Environ(), EnvProcPrioMgmt)
+func CommandContext(ctx context.Context, cmd string, args ...string) (*exec.Cmd, error) {
+	environ := os.Environ()
+	_, enabled := envVal(environ, EnvProcPrioMgmt)
 	if runtime.GOOS != "linux" || !enabled {
 		return exec.CommandContext(ctx, cmd, args...), nil
 	}

diff --git a/agent/agentexec/exec_test.go b/agent/agentexec/exec_test.go
@@ -0,0 +1,61 @@
+package agentexec_test
+
+import (
+	"context"
+	"os"
+	"os/exec"
+	"runtime"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/v2/agent/agentexec"
+)
+
+func TestExec(t *testing.T) {
+	t.Run("NonLinux", func(t *testing.T) {
+
+		t.Setenv(agentexec.EnvProcPrioMgmt, "true")
+
+		if runtime.GOOS == "linux" {
+			t.Skip("skipping on linux")
+		}
+
+		cmd, err := agentexec.CommandContext(context.Background(), "sh", "-c", "sleep")
+		require.NoError(t, err)
+		require.Equal(t, "sh", cmd.Path)
+		require.Equal(t, []string{"-c", "sleep"}, cmd.Args[1:])
+	})
+
+	t.Run("Linux", func(t *testing.T) {
+
+		t.Run("Disabled", func(t *testing.T) {
+			if runtime.GOOS != "linux" {
+				t.Skip("skipping on linux")
+			}
+
+			cmd, err := agentexec.CommandContext(context.Background(), "sh", "-c", "sleep")
+			require.NoError(t, err)
+			path, err := exec.LookPath("sh")
+			require.NoError(t, err)
+			require.Equal(t, path, cmd.Path)
+			require.Equal(t, []string{"sh", "-c", "sleep"}, cmd.Args)
+		})
+
+		t.Run("Enabled", func(t *testing.T) {
+			t.Setenv(agentexec.EnvProcPrioMgmt, "hello")
+
+			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", "sh", "-c", "sleep"}, cmd.Args)
+		})
+	})
+}
diff --git a/cli/root.go b/cli/root.go
@@ -121,8 +121,6 @@ func (r *RootCmd) CoreSubcommands() []*serpent.Command {
 		r.update(),
 		r.whoami(),
 
-		// Hidden
-		r.agentExec(),
 		r.expCmd(),
 		r.gitssh(),
 		r.support(),