Skip to content

Commit b158919

Browse files
sreyasreya
committed
flag parsing
1 parent 986e18e commit b158919

File tree

3 files changed

+66
-42
lines changed

3 files changed

+66
-42
lines changed

‎agent/agentexec/cli_unix.go

Lines changed: 37 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@ import (
99
"os"
1010
"os/exec"
1111
"runtime"
12-
"slices"
1312
"strconv"
1413
"strings"
1514
"syscall"
@@ -18,47 +17,57 @@ import (
1817
"golang.org/x/xerrors"
1918
)
2019

20+
// unset is set to an invalid value for nice and oom scores.
21+
const unset = -2000
22+
2123
// CLI runs the agent-exec command. It should only be called by the cli package.
22-
func CLI(args []string, environ []string) error {
24+
func CLI() error {
2325
// We lock the OS thread here to avoid a race conditino where the nice priority
2426
// we get is on a different thread from the one we set it on.
2527
runtime.LockOSThread()
2628

27-
nice := flag.Int(niceArg, 0, "")
28-
oom := flag.Int(oomArg, 0, "")
29+
var (
30+
fs = flag.NewFlagSet("agent-exec", flag.ExitOnError)
31+
nice = fs.Int("coder-nice", unset, "")
32+
oom = fs.Int("coder-oom", unset, "")
33+
)
34+
35+
if len(os.Args) < 3 {
36+
return xerrors.Errorf("malformed command %+v", os.Args)
37+
}
2938

30-
flag.Parse()
39+
// Parse everything after "coder agent-exec".
40+
err := fs.Parse(os.Args[2:])
41+
if err != nil {
42+
return xerrors.Errorf("parse flags: %w", err)
43+
}
3144

3245
if runtime.GOOS != "linux" {
3346
return xerrors.Errorf("agent-exec is only supported on Linux")
3447
}
3548

36-
if len(args) < 2 {
37-
return xerrors.Errorf("malformed command %q", args)
49+
// Get everything after "coder agent-exec --"
50+
args := execArgs(os.Args)
51+
if len(args) == 0 {
52+
return xerrors.Errorf("no exec command provided %+v", os.Args)
3853
}
3954

40-
// Slice off 'coder agent-exec'
41-
args = args[2:]
42-
4355
pid := os.Getpid()
4456

45-
var err error
46-
if nice == nil {
57+
if *nice == unset {
4758
// If an explicit nice score isn't set, we use the default.
48-
n, err := defaultNiceScore()
59+
*nice, err = defaultNiceScore()
4960
if err != nil {
5061
return xerrors.Errorf("get default nice score: %w", err)
5162
}
52-
nice = &n
5363
}
5464

55-
if oom == nil {
65+
if *oom == unset {
5666
// If an explicit oom score isn't set, we use the default.
57-
o, err := defaultOOMScore()
67+
*oom, err = defaultOOMScore()
5868
if err != nil {
5969
return xerrors.Errorf("get default oom score: %w", err)
6070
}
61-
oom = &o
6271
}
6372

6473
err = unix.Setpriority(unix.PRIO_PROCESS, 0, *nice)
@@ -76,12 +85,7 @@ func CLI(args []string, environ []string) error {
7685
return xerrors.Errorf("look path: %w", err)
7786
}
7887

79-
// Remove environments variables specifically set for the agent-exec command.
80-
env := slices.DeleteFunc(environ, func(env string) bool {
81-
return strings.HasPrefix(env, EnvProcOOMScore) || strings.HasPrefix(env, EnvProcNiceScore)
82-
})
83-
84-
return syscall.Exec(path, args, env)
88+
return syscall.Exec(path, args, os.Environ())
8589
}
8690

8791
func defaultNiceScore() (int, error) {
@@ -133,3 +137,13 @@ func oomScoreAdj(pid int) (int, error) {
133137
func writeOOMScoreAdj(pid int, score int) error {
134138
return os.WriteFile(fmt.Sprintf("/proc/%d/oom_score_adj", pid), []byte(fmt.Sprintf("%d", score)), 0o600)
135139
}
140+
141+
// execArgs returns the arguments to pass to syscall.Exec after the "--" delimiter.
142+
func execArgs(args []string) []string {
143+
for i, arg := range args {
144+
if arg == "--" {
145+
return args[i+1:]
146+
}
147+
}
148+
return nil
149+
}

‎agent/agentexec/cli_unix_test.go

Lines changed: 28 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -29,9 +29,7 @@ func TestCLI(t *testing.T) {
2929
t.Parallel()
3030

3131
ctx := testutil.Context(t, testutil.WaitMedium)
32-
cmd, path := cmd(ctx, t)
33-
cmd.Env = append(cmd.Env, "CODER_PROC_NICE_SCORE=12")
34-
cmd.Env = append(cmd.Env, "CODER_PROC_OOM_SCORE=123")
32+
cmd, path := cmd(ctx, t, 123, 12)
3533
err := cmd.Start()
3634
require.NoError(t, err)
3735
go cmd.Wait()
@@ -45,7 +43,7 @@ func TestCLI(t *testing.T) {
4543
t.Parallel()
4644

4745
ctx := testutil.Context(t, testutil.WaitMedium)
48-
cmd, path := cmd(ctx, t)
46+
cmd, path := cmd(ctx, t, 0, 0)
4947
err := cmd.Start()
5048
require.NoError(t, err)
5149
go cmd.Wait()
@@ -101,22 +99,22 @@ func waitForSentinel(ctx context.Context, t *testing.T, cmd *exec.Cmd, path stri
10199
}
102100
}
103101

104-
func cmd(ctx context.Context, t *testing.T, args ...string) (*exec.Cmd, string) {
105-
file := ""
106-
//nolint:gosec
107-
cmd := exec.Command(TestBin, append([]string{"agent-exec"}, args...)...)
108-
if len(args) == 0 {
109-
// Generate a unique path that we can touch to indicate that we've progressed past the
110-
// syscall.Exec.
111-
dir := t.TempDir()
102+
func cmd(ctx context.Context, t *testing.T, oom, nice int) (*exec.Cmd, string) {
103+
var (
104+
args = execArgs(oom, nice)
105+
dir = t.TempDir()
112106
file = filepath.Join(dir, "sentinel")
113-
//nolint:gosec
114-
cmd = exec.CommandContext(ctx, TestBin, "agent-exec", "sh", "-c", fmt.Sprintf("touch %s && sleep 10m", file))
115-
// We set this so we can also easily kill the sleep process the shell spawns.
116-
cmd.SysProcAttr = &syscall.SysProcAttr{
117-
Setpgid: true,
118-
}
107+
)
108+
109+
args = append(args, "sh", "-c", fmt.Sprintf("touch %s && sleep 10m", file))
110+
//nolint:gosec
111+
cmd := exec.CommandContext(ctx, TestBin, args...)
112+
113+
// We set this so we can also easily kill the sleep process the shell spawns.
114+
cmd.SysProcAttr = &syscall.SysProcAttr{
115+
Setpgid: true,
119116
}
117+
120118
cmd.Env = os.Environ()
121119
var buf bytes.Buffer
122120
cmd.Stdout = &buf
@@ -166,3 +164,15 @@ func expectedNiceScore(t *testing.T) int {
166164
}
167165
return score
168166
}
167+
168+
func execArgs(oom int, nice int) []string {
169+
execArgs := []string{"agent-exec"}
170+
if oom != 0 {
171+
execArgs = append(execArgs, fmt.Sprintf("--coder-oom=%d", oom))
172+
}
173+
if nice != 0 {
174+
execArgs = append(execArgs, fmt.Sprintf("--coder-nice=%d", nice))
175+
}
176+
execArgs = append(execArgs, "--")
177+
return execArgs
178+
}

‎agent/agentexec/cmdtest/main.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ import (
88
)
99

1010
func main() {
11-
err := agentexec.CLI(os.Args, os.Environ())
11+
err := agentexec.CLI()
1212
if err != nil {
1313
_, _ = fmt.Fprintln(os.Stderr, err)
1414
os.Exit(1)

0 commit comments

Comments
 (0)