Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commitb158919

Browse files
committed
flag parsing
1 parent986e18e commitb158919

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+
constunset=-2000
22+
2123
// CLI runs the agent-exec command. It should only be called by the cli package.
22-
funcCLI(args []string,environ []string)error {
24+
funcCLI()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+
iflen(os.Args)<3 {
36+
returnxerrors.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+
iferr!=nil {
42+
returnxerrors.Errorf("parse flags: %w",err)
43+
}
3144

3245
ifruntime.GOOS!="linux" {
3346
returnxerrors.Errorf("agent-exec is only supported on Linux")
3447
}
3548

36-
iflen(args)<2 {
37-
returnxerrors.Errorf("malformed command %q",args)
49+
// Get everything after "coder agent-exec --"
50+
args:=execArgs(os.Args)
51+
iflen(args)==0 {
52+
returnxerrors.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-
varerrerror
46-
ifnice==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
iferr!=nil {
5061
returnxerrors.Errorf("get default nice score: %w",err)
5162
}
52-
nice=&n
5363
}
5464

55-
ifoom==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
iferr!=nil {
5969
returnxerrors.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
returnxerrors.Errorf("look path: %w",err)
7786
}
7887

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

8791
funcdefaultNiceScore() (int,error) {
@@ -133,3 +137,13 @@ func oomScoreAdj(pid int) (int, error) {
133137
funcwriteOOMScoreAdj(pidint,scoreint)error {
134138
returnos.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+
funcexecArgs(args []string) []string {
143+
fori,arg:=rangeargs {
144+
ifarg=="--" {
145+
returnargs[i+1:]
146+
}
147+
}
148+
returnnil
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
gocmd.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
gocmd.Wait()
@@ -101,22 +99,22 @@ func waitForSentinel(ctx context.Context, t *testing.T, cmd *exec.Cmd, path stri
10199
}
102100
}
103101

104-
funccmd(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-
iflen(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+
funccmd(ctx context.Context,t*testing.T,oom,niceint) (*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
varbuf bytes.Buffer
122120
cmd.Stdout=&buf
@@ -166,3 +164,15 @@ func expectedNiceScore(t *testing.T) int {
166164
}
167165
returnscore
168166
}
167+
168+
funcexecArgs(oomint,niceint) []string {
169+
execArgs:= []string{"agent-exec"}
170+
ifoom!=0 {
171+
execArgs=append(execArgs,fmt.Sprintf("--coder-oom=%d",oom))
172+
}
173+
ifnice!=0 {
174+
execArgs=append(execArgs,fmt.Sprintf("--coder-nice=%d",nice))
175+
}
176+
execArgs=append(execArgs,"--")
177+
returnexecArgs
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
funcmain() {
11-
err:=agentexec.CLI(os.Args,os.Environ())
11+
err:=agentexec.CLI()
1212
iferr!=nil {
1313
_,_=fmt.Fprintln(os.Stderr,err)
1414
os.Exit(1)

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp