Skip to content

Commit 97e68f4

Browse files
sreyasreya
committed
idk
1 parent 196c8a9 commit 97e68f4

File tree

2 files changed

+152
-27
lines changed

2 files changed

+152
-27
lines changed

‎agent/agentexec/cli.go

Lines changed: 91 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
"strings"
1111
"syscall"
1212

13+
"golang.org/x/sys/unix"
1314
"golang.org/x/xerrors"
1415
)
1516

@@ -28,32 +29,37 @@ func CLI(args []string, environ []string) error {
2829
return xerrors.Errorf("malformed command %q", args)
2930
}
3031

32+
// Slice off 'coder agent-exec'
3133
args = args[2:]
3234

3335
pid := os.Getpid()
3436

35-
oomScore, ok := envVal(environ, EnvProcOOMScore)
37+
var err error
38+
nice, ok := envValInt(environ, EnvProcNiceScore)
3639
if !ok {
37-
return xerrors.Errorf("missing %q", EnvProcOOMScore)
40+
// If an explicit nice score isn't set, we use the default.
41+
nice, err = defaultNiceScore()
42+
if err != nil {
43+
return xerrors.Errorf("get default nice score: %w", err)
44+
}
45+
fmt.Println("nice score", nice, "pid", pid)
3846
}
3947

40-
niceScore, ok := envVal(environ, EnvProcNiceScore)
48+
oomscore, ok := envValInt(environ, EnvProcOOMScore)
4149
if !ok {
42-
return xerrors.Errorf("missing %q", EnvProcNiceScore)
43-
}
44-
45-
score, err := strconv.Atoi(niceScore)
46-
if err != nil {
47-
return xerrors.Errorf("invalid nice score: %w", err)
50+
// If an explicit oom score isn't set, we use the default.
51+
oomscore, err = defaultOOMScore()
52+
if err != nil {
53+
return xerrors.Errorf("get default oom score: %w", err)
54+
}
4855
}
4956

50-
err = syscall.Setpriority(syscall.PRIO_PROCESS, pid, score)
57+
err = unix.Setpriority(unix.PRIO_PROCESS, pid, nice)
5158
if err != nil {
5259
return xerrors.Errorf("set nice score: %w", err)
5360
}
5461

55-
oomPath := fmt.Sprintf("/proc/%d/oom_score_adj", pid)
56-
err = os.WriteFile(oomPath, []byte(oomScore), 0o600)
62+
err = writeOOMScoreAdj(pid, oomscore)
5763
if err != nil {
5864
return xerrors.Errorf("set oom score: %w", err)
5965
}
@@ -63,17 +69,86 @@ func CLI(args []string, environ []string) error {
6369
return xerrors.Errorf("look path: %w", err)
6470
}
6571

72+
// Remove environments variables specifically set for the agent-exec command.
6673
env := slices.DeleteFunc(environ, func(env string) bool {
6774
return strings.HasPrefix(env, EnvProcOOMScore) || strings.HasPrefix(env, EnvProcNiceScore)
6875
})
6976

7077
return syscall.Exec(path, args, env)
7178
}
7279

73-
func envVal(environ []string, key string) (string, bool) {
74-
for _, env := range environ {
75-
if strings.HasPrefix(env, key+"=") {
76-
return strings.TrimPrefix(env, key+"="), true
80+
func defaultNiceScore() (int, error) {
81+
score, err := unix.Getpriority(unix.PRIO_PROCESS, os.Getpid())
82+
if err != nil {
83+
return 0, xerrors.Errorf("get nice score: %w", err)
84+
}
85+
// Priority is niceness + 20.
86+
score -= 20
87+
88+
score += 5
89+
if score > 19 {
90+
return 19, nil
91+
}
92+
return score, nil
93+
}
94+
95+
func defaultOOMScore() (int, error) {
96+
score, err := oomScoreAdj(os.Getpid())
97+
if err != nil {
98+
return 0, xerrors.Errorf("get oom score: %w", err)
99+
}
100+
101+
// If the agent has a negative oom_score_adj, we set the child to 0
102+
// so it's treated like every other process.
103+
if score < 0 {
104+
return 0, nil
105+
}
106+
107+
// If the agent is already almost at the maximum then set it to the max.
108+
if score >= 998 {
109+
return 1000, nil
110+
}
111+
112+
// If the agent oom_score_adj is >=0, we set the child to slightly
113+
// less than the maximum. If users want a different score they set it
114+
// directly.
115+
return 998, nil
116+
}
117+
118+
func oomScoreAdj(pid int) (int, error) {
119+
scoreStr, err := os.ReadFile(fmt.Sprintf("/proc/%d/oom_score_adj", pid))
120+
if err != nil {
121+
return 0, xerrors.Errorf("read oom_score_adj: %w", err)
122+
}
123+
return strconv.Atoi(strings.TrimSpace(string(scoreStr)))
124+
}
125+
126+
func writeOOMScoreAdj(pid int, score int) error {
127+
return os.WriteFile(fmt.Sprintf("/proc/%d/oom_score_adj", pid), []byte(fmt.Sprintf("%d", score)), 0o600)
128+
}
129+
130+
// envValInt searches for a key in a list of environment variables and parses it to an int.
131+
// If the key is not found or cannot be parsed, returns 0 and false.
132+
func envValInt(env []string, key string) (int, bool) {
133+
val, ok := envVal(env, key)
134+
if !ok {
135+
return 0, false
136+
}
137+
138+
i, err := strconv.Atoi(val)
139+
if err != nil {
140+
return 0, false
141+
}
142+
return i, true
143+
}
144+
145+
// envVal searches for a key in a list of environment variables and returns its value.
146+
// If the key is not found, returns empty string and false.
147+
func envVal(env []string, key string) (string, bool) {
148+
prefix := key + "="
149+
for _, e := range env {
150+
if strings.HasPrefix(e, prefix) {
151+
return strings.TrimPrefix(e, prefix), true
77152
}
78153
}
79154
return "", false

‎agent/agentexec/cli_test.go

Lines changed: 61 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -32,16 +32,34 @@ func TestCLI(t *testing.T) {
3232
require.NoError(t, err)
3333
go cmd.Wait()
3434

35-
waitForSentinel(t, ctx, cmd, path)
35+
waitForSentinel(ctx, t, cmd, path)
3636
requireOOMScore(t, cmd.Process.Pid, 123)
3737
requireNiceScore(t, cmd.Process.Pid, 10)
3838
})
39+
40+
t.Run("Defaults", func(t *testing.T) {
41+
t.Parallel()
42+
43+
ctx := testutil.Context(t, testutil.WaitMedium)
44+
cmd, path := cmd(ctx, t)
45+
err := cmd.Start()
46+
require.NoError(t, err)
47+
go cmd.Wait()
48+
49+
waitForSentinel(ctx, t, cmd, path)
50+
51+
expectedNice := expectedNiceScore(t)
52+
expectedOOM := expectedOOMScore(t)
53+
fmt.Println("expected nice", expectedNice, "expected oom", expectedOOM)
54+
requireOOMScore(t, cmd.Process.Pid, expectedOOM)
55+
requireNiceScore(t, cmd.Process.Pid, expectedNice)
56+
})
3957
}
4058

4159
func requireNiceScore(t *testing.T, pid int, score int) {
4260
t.Helper()
4361

44-
nice, err := unix.Getpriority(0, pid)
62+
nice, err := unix.Getpriority(unix.PRIO_PROCESS, pid)
4563
require.NoError(t, err)
4664
require.Equal(t, score, nice)
4765
}
@@ -55,7 +73,7 @@ func requireOOMScore(t *testing.T, pid int, expected int) {
5573
require.Equal(t, strconv.Itoa(expected), score)
5674
}
5775

58-
func waitForSentinel(t *testing.T, ctx context.Context, cmd *exec.Cmd, path string) {
76+
func waitForSentinel(ctx context.Context, t *testing.T, cmd *exec.Cmd, path string) {
5977
t.Helper()
6078

6179
ticker := time.NewTicker(testutil.IntervalFast)
@@ -74,17 +92,20 @@ func waitForSentinel(t *testing.T, ctx context.Context, cmd *exec.Cmd, path stri
7492
select {
7593
case <-ticker.C:
7694
case <-ctx.Done():
77-
return
95+
require.NoError(t, ctx.Err())
7896
}
7997
}
8098
}
8199

82100
func cmd(ctx context.Context, t *testing.T, args ...string) (*exec.Cmd, string) {
83101
file := ""
84-
cmd := exec.Command(TestBin, args...)
102+
//nolint:gosec
103+
cmd := exec.Command(TestBin, append([]string{"agent-exec"}, args...)...)
85104
if len(args) == 0 {
105+
// Generate a unique path that we can touch to indicate that we've progressed past the
106+
// syscall.Exec.
86107
dir := t.TempDir()
87-
file = filepath.Join(dir, uniqueFile(t))
108+
file = filepath.Join(dir, "sentinel")
88109
//nolint:gosec
89110
cmd = exec.CommandContext(ctx, TestBin, "agent-exec", "sh", "-c", fmt.Sprintf("touch %s && sleep 10m", file))
90111
}
@@ -98,13 +119,42 @@ func cmd(ctx context.Context, t *testing.T, args ...string) (*exec.Cmd, string)
98119
t.Logf("cmd %q output: %s", cmd.Args, buf.String())
99120
}
100121

101-
if cmd.Process != nil {
102-
_ = cmd.Process.Kill()
103-
}
122+
// if cmd.Process != nil {
123+
// _ = cmd.Process.Kill()
124+
// }
104125
})
105126
return cmd, file
106127
}
107128

108-
func uniqueFile(t *testing.T) string {
109-
return fmt.Sprintf("%s-%d", strings.ReplaceAll(t.Name(), "/", "_"), time.Now().UnixNano())
129+
func expectedOOMScore(t *testing.T) int {
130+
t.Helper()
131+
132+
score, err := os.ReadFile(fmt.Sprintf("/proc/%d/oom_score_adj", os.Getpid()))
133+
require.NoError(t, err)
134+
135+
scoreInt, err := strconv.Atoi(strings.TrimSpace(string(score)))
136+
require.NoError(t, err)
137+
138+
if scoreInt < 0 {
139+
return 0
140+
}
141+
if scoreInt >= 998 {
142+
return 1000
143+
}
144+
return 998
145+
}
146+
147+
func expectedNiceScore(t *testing.T) int {
148+
t.Helper()
149+
150+
score, err := unix.Getpriority(unix.PRIO_PROCESS, os.Getpid())
151+
require.NoError(t, err)
152+
153+
// Priority is niceness + 20.
154+
score -= 20
155+
score += 5
156+
if score > 19 {
157+
return 19
158+
}
159+
return score
110160
}

0 commit comments

Comments
 (0)