pr changes

sreya · sreya · commit 5076cf08d3c3 · 2024-11-21T21:08:12.000Z
diff --git a/agent/agentexec/cli_linux.go b/agent/agentexec/cli_linux.go
@@ -25,6 +25,8 @@ 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()
+	// Nop on success but we do it anyway in case of an error.
+	defer runtime.UnlockOSThread()
 
 	var (
 		fs   = flag.NewFlagSet("agent-exec", flag.ExitOnError)
@@ -42,18 +44,12 @@ func CLI() error {
 		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()
@@ -75,7 +71,7 @@ func CLI() error {
 		return xerrors.Errorf("set nice score: %w", err)
 	}
 
-	err = writeOOMScoreAdj(pid, *oom)
+	err = writeOOMScoreAdj(*oom)
 	if err != nil {
 		return xerrors.Errorf("set oom score: %w", err)
 	}
@@ -104,7 +100,7 @@ func defaultNiceScore() (int, error) {
 }
 
 func defaultOOMScore() (int, error) {
-	score, err := oomScoreAdj(os.Getpid())
+	score, err := oomScoreAdj()
 	if err != nil {
 		return 0, xerrors.Errorf("get oom score: %w", err)
 	}
@@ -126,16 +122,16 @@ func defaultOOMScore() (int, error) {
 	return 998, nil
 }
 
-func oomScoreAdj(pid int) (int, error) {
-	scoreStr, err := os.ReadFile(fmt.Sprintf("/proc/%d/oom_score_adj", pid))
+func oomScoreAdj() (int, error) {
+	scoreStr, err := os.ReadFile("/proc/self/oom_score_adj")
 	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)
+func writeOOMScoreAdj(score int) error {
+	return os.WriteFile("/proc/self/oom_score_adj", []byte(fmt.Sprintf("%d", score)), 0o600)
 }
 
 // execArgs returns the arguments to pass to syscall.Exec after the "--" delimiter.
diff --git a/agent/agentexec/cli_linux_test.go b/agent/agentexec/cli_linux_test.go
diff --git a/agent/agentexec/cli_other.go b/agent/agentexec/cli_other.go
@@ -6,5 +6,5 @@ package agentexec
 import "golang.org/x/xerrors"
 
 func CLI() error {
-	return xerrors.Errorf("agent-exec is only supported on Linux")
+	return xerrors.New("agent-exec is only supported on Linux")
 }
diff --git a/agent/agentexec/cmdtest/main_linux.go b/agent/agentexec/cmdtest/main_linux.go
@@ -1,3 +1,6 @@
+//go:build linux
+// +build linux
+
 package main
 
 import (
diff --git a/agent/agentexec/main_linux_test.go b/agent/agentexec/main_linux_test.go
@@ -0,0 +1,43 @@
+package agentexec_test
+
+import (
+	"fmt"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"testing"
+)
+
+var TestBin string
+
+func TestMain(m *testing.M) {
+	code := func() int {
+		// We generate a unique directory per test invocation to avoid collisions between two
+		// processes attempting to create the same temp file.
+		dir := genDir()
+		defer os.RemoveAll(dir)
+		TestBin = buildBinary(dir)
+		return m.Run()
+	}()
+
+	os.Exit(code)
+}
+
+func buildBinary(dir string) string {
+	path := filepath.Join(dir, "agent-test")
+	out, err := exec.Command("go", "build", "-o", path, "./cmdtest").CombinedOutput()
+	mustf(err, "build binary: %s", out)
+	return path
+}
+
+func mustf(err error, msg string, args ...any) {
+	if err != nil {
+		panic(fmt.Sprintf(msg, args...))
+	}
+}
+
+func genDir() string {
+	dir, err := os.MkdirTemp(os.TempDir(), "agentexec")
+	mustf(err, "create temp dir: %v", err)
+	return dir
+}
diff --git a/agent/agentexec/main_test.go b/agent/agentexec/main_test.go

Original file line number	Diff line number	Diff line change
`@@ -6,5 +6,5 @@ package agentexec`
`6`	`6`	`import "golang.org/x/xerrors"`
`7`	`7`
`8`	`8`	`func CLI() error {`
`9`		`- return xerrors.Errorf("agent-exec is only supported on Linux")`
	`9`	`+ return xerrors.New("agent-exec is only supported on Linux")`
`10`	`10`	`}`