Skip to content

Commit 986e18e

Browse files
sreyasreya
committed
use flag parsing instead
1 parent ae30643 commit 986e18e

File tree

3 files changed

+96
-26
lines changed

3 files changed

+96
-26
lines changed

‎agent/agentexec/cli_unix.go

Lines changed: 16 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
package agentexec
55

66
import (
7+
"flag"
78
"fmt"
89
"os"
910
"os/exec"
@@ -19,8 +20,15 @@ import (
1920

2021
// CLI runs the agent-exec command. It should only be called by the cli package.
2122
func CLI(args []string, environ []string) error {
23+
// We lock the OS thread here to avoid a race conditino where the nice priority
24+
// we get is on a different thread from the one we set it on.
2225
runtime.LockOSThread()
2326

27+
nice := flag.Int(niceArg, 0, "")
28+
oom := flag.Int(oomArg, 0, "")
29+
30+
flag.Parse()
31+
2432
if runtime.GOOS != "linux" {
2533
return xerrors.Errorf("agent-exec is only supported on Linux")
2634
}
@@ -35,30 +43,30 @@ func CLI(args []string, environ []string) error {
3543
pid := os.Getpid()
3644

3745
var err error
38-
nice, ok := envValInt(environ, EnvProcNiceScore)
39-
if !ok {
46+
if nice == nil {
4047
// If an explicit nice score isn't set, we use the default.
41-
nice, err = defaultNiceScore()
48+
n, err := defaultNiceScore()
4249
if err != nil {
4350
return xerrors.Errorf("get default nice score: %w", err)
4451
}
52+
nice = &n
4553
}
4654

47-
oomscore, ok := envValInt(environ, EnvProcOOMScore)
48-
if !ok {
55+
if oom == nil {
4956
// If an explicit oom score isn't set, we use the default.
50-
oomscore, err = defaultOOMScore()
57+
o, err := defaultOOMScore()
5158
if err != nil {
5259
return xerrors.Errorf("get default oom score: %w", err)
5360
}
61+
oom = &o
5462
}
5563

56-
err = unix.Setpriority(unix.PRIO_PROCESS, 0, nice)
64+
err = unix.Setpriority(unix.PRIO_PROCESS, 0, *nice)
5765
if err != nil {
5866
return xerrors.Errorf("set nice score: %w", err)
5967
}
6068

61-
err = writeOOMScoreAdj(pid, oomscore)
69+
err = writeOOMScoreAdj(pid, *oom)
6270
if err != nil {
6371
return xerrors.Errorf("set oom score: %w", err)
6472
}

‎agent/agentexec/exec.go

Lines changed: 27 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,12 @@ package agentexec
22

33
import (
44
"context"
5+
"fmt"
56
"os"
67
"os/exec"
78
"path/filepath"
89
"runtime"
910
"strconv"
10-
"strings"
1111

1212
"golang.org/x/xerrors"
1313
)
@@ -25,8 +25,7 @@ const (
2525
// is returned. All instances of exec.Cmd should flow through this function to ensure
2626
// proper resource constraints are applied to the child process.
2727
func CommandContext(ctx context.Context, cmd string, args ...string) (*exec.Cmd, error) {
28-
environ := os.Environ()
29-
_, enabled := envVal(environ, EnvProcPrioMgmt)
28+
_, enabled := os.LookupEnv(EnvProcPrioMgmt)
3029
if runtime.GOOS != "linux" || !enabled {
3130
return exec.CommandContext(ctx, cmd, args...), nil
3231
}
@@ -41,14 +40,24 @@ func CommandContext(ctx context.Context, cmd string, args ...string) (*exec.Cmd,
4140
return nil, xerrors.Errorf("eval symlinks: %w", err)
4241
}
4342

44-
args = append([]string{"agent-exec", cmd}, args...)
45-
return exec.CommandContext(ctx, bin, args...), nil
43+
execArgs := []string{"agent-exec"}
44+
if score, ok := envValInt(EnvProcOOMScore); ok {
45+
execArgs = append(execArgs, oomScoreArg(score))
46+
}
47+
48+
if score, ok := envValInt(EnvProcNiceScore); ok {
49+
execArgs = append(execArgs, niceScoreArg(score))
50+
}
51+
execArgs = append(execArgs, "--", cmd)
52+
execArgs = append(execArgs, args...)
53+
54+
return exec.CommandContext(ctx, bin, execArgs...), nil
4655
}
4756

4857
// envValInt searches for a key in a list of environment variables and parses it to an int.
4958
// If the key is not found or cannot be parsed, returns 0 and false.
50-
func envValInt(env []string, key string) (int, bool) {
51-
val, ok := envVal(env, key)
59+
func envValInt(key string) (int, bool) {
60+
val, ok := os.LookupEnv(key)
5261
if !ok {
5362
return 0, false
5463
}
@@ -60,14 +69,15 @@ func envValInt(env []string, key string) (int, bool) {
6069
return i, true
6170
}
6271

63-
// envVal searches for a key in a list of environment variables and returns its value.
64-
// If the key is not found, returns empty string and false.
65-
func envVal(env []string, key string) (string, bool) {
66-
prefix := key + "="
67-
for _, e := range env {
68-
if strings.HasPrefix(e, prefix) {
69-
return strings.TrimPrefix(e, prefix), true
70-
}
71-
}
72-
return "", false
72+
const (
73+
niceArg = "coder-nice"
74+
oomArg = "coder-oom"
75+
)
76+
77+
func niceScoreArg(score int) string {
78+
return fmt.Sprintf("--%s=%d", niceArg, score)
79+
}
80+
81+
func oomScoreArg(score int) string {
82+
return fmt.Sprintf("--%s=%d", oomArg, score)
7383
}

‎agent/agentexec/exec_test.go

Lines changed: 53 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,59 @@ func TestExec(t *testing.T) {
6161
cmd, err := agentexec.CommandContext(context.Background(), "sh", "-c", "sleep")
6262
require.NoError(t, err)
6363
require.Equal(t, executable, cmd.Path)
64-
require.Equal(t, []string{executable, "agent-exec", "sh", "-c", "sleep"}, cmd.Args)
64+
require.Equal(t, []string{executable, "agent-exec", "--", "sh", "-c", "sleep"}, cmd.Args)
65+
})
66+
67+
t.Run("Nice", func(t *testing.T) {
68+
t.Setenv(agentexec.EnvProcPrioMgmt, "hello")
69+
t.Setenv(agentexec.EnvProcNiceScore, "10")
70+
71+
if runtime.GOOS != "linux" {
72+
t.Skip("skipping on linux")
73+
}
74+
75+
executable, err := os.Executable()
76+
require.NoError(t, err)
77+
78+
cmd, err := agentexec.CommandContext(context.Background(), "sh", "-c", "sleep")
79+
require.NoError(t, err)
80+
require.Equal(t, executable, cmd.Path)
81+
require.Equal(t, []string{executable, "agent-exec", "--coder-nice=10", "--", "sh", "-c", "sleep"}, cmd.Args)
82+
})
83+
84+
t.Run("OOM", func(t *testing.T) {
85+
t.Setenv(agentexec.EnvProcPrioMgmt, "hello")
86+
t.Setenv(agentexec.EnvProcOOMScore, "123")
87+
88+
if runtime.GOOS != "linux" {
89+
t.Skip("skipping on linux")
90+
}
91+
92+
executable, err := os.Executable()
93+
require.NoError(t, err)
94+
95+
cmd, err := agentexec.CommandContext(context.Background(), "sh", "-c", "sleep")
96+
require.NoError(t, err)
97+
require.Equal(t, executable, cmd.Path)
98+
require.Equal(t, []string{executable, "agent-exec", "--coder-oom=123", "--", "sh", "-c", "sleep"}, cmd.Args)
99+
})
100+
101+
t.Run("Both", func(t *testing.T) {
102+
t.Setenv(agentexec.EnvProcPrioMgmt, "hello")
103+
t.Setenv(agentexec.EnvProcOOMScore, "432")
104+
t.Setenv(agentexec.EnvProcNiceScore, "14")
105+
106+
if runtime.GOOS != "linux" {
107+
t.Skip("skipping on linux")
108+
}
109+
110+
executable, err := os.Executable()
111+
require.NoError(t, err)
112+
113+
cmd, err := agentexec.CommandContext(context.Background(), "sh", "-c", "sleep")
114+
require.NoError(t, err)
115+
require.Equal(t, executable, cmd.Path)
116+
require.Equal(t, []string{executable, "agent-exec", "--coder-oom=432", "--coder-nice=14", "--", "sh", "-c", "sleep"}, cmd.Args)
65117
})
66118
})
67119
}

0 commit comments

Comments
 (0)