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

Commit986e18e

Browse files
committed
use flag parsing instead
1 parentae30643 commit986e18e

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
funcCLI(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
ifruntime.GOOS!="linux" {
2533
returnxerrors.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
varerrerror
38-
nice,ok:=envValInt(environ,EnvProcNiceScore)
39-
if!ok {
46+
ifnice==nil {
4047
// If an explicit nice score isn't set, we use the default.
41-
nice,err=defaultNiceScore()
48+
n,err:=defaultNiceScore()
4249
iferr!=nil {
4350
returnxerrors.Errorf("get default nice score: %w",err)
4451
}
52+
nice=&n
4553
}
4654

47-
oomscore,ok:=envValInt(environ,EnvProcOOMScore)
48-
if!ok {
55+
ifoom==nil {
4956
// If an explicit oom score isn't set, we use the default.
50-
oomscore,err=defaultOOMScore()
57+
o,err:=defaultOOMScore()
5158
iferr!=nil {
5259
returnxerrors.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
iferr!=nil {
5866
returnxerrors.Errorf("set nice score: %w",err)
5967
}
6068

61-
err=writeOOMScoreAdj(pid,oomscore)
69+
err=writeOOMScoreAdj(pid,*oom)
6270
iferr!=nil {
6371
returnxerrors.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
funcCommandContext(ctx context.Context,cmdstring,args...string) (*exec.Cmd,error) {
28-
environ:=os.Environ()
29-
_,enabled:=envVal(environ,EnvProcPrioMgmt)
28+
_,enabled:=os.LookupEnv(EnvProcPrioMgmt)
3029
ifruntime.GOOS!="linux"||!enabled {
3130
returnexec.CommandContext(ctx,cmd,args...),nil
3231
}
@@ -41,14 +40,24 @@ func CommandContext(ctx context.Context, cmd string, args ...string) (*exec.Cmd,
4140
returnnil,xerrors.Errorf("eval symlinks: %w",err)
4241
}
4342

44-
args=append([]string{"agent-exec",cmd},args...)
45-
returnexec.CommandContext(ctx,bin,args...),nil
43+
execArgs:= []string{"agent-exec"}
44+
ifscore,ok:=envValInt(EnvProcOOMScore);ok {
45+
execArgs=append(execArgs,oomScoreArg(score))
46+
}
47+
48+
ifscore,ok:=envValInt(EnvProcNiceScore);ok {
49+
execArgs=append(execArgs,niceScoreArg(score))
50+
}
51+
execArgs=append(execArgs,"--",cmd)
52+
execArgs=append(execArgs,args...)
53+
54+
returnexec.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-
funcenvValInt(env []string,keystring) (int,bool) {
51-
val,ok:=envVal(env,key)
59+
funcenvValInt(keystring) (int,bool) {
60+
val,ok:=os.LookupEnv(key)
5261
if!ok {
5362
return0,false
5463
}
@@ -60,14 +69,15 @@ func envValInt(env []string, key string) (int, bool) {
6069
returni,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-
funcenvVal(env []string,keystring) (string,bool) {
66-
prefix:=key+"="
67-
for_,e:=rangeenv {
68-
ifstrings.HasPrefix(e,prefix) {
69-
returnstrings.TrimPrefix(e,prefix),true
70-
}
71-
}
72-
return"",false
72+
const (
73+
niceArg="coder-nice"
74+
oomArg="coder-oom"
75+
)
76+
77+
funcniceScoreArg(scoreint)string {
78+
returnfmt.Sprintf("--%s=%d",niceArg,score)
79+
}
80+
81+
funcoomScoreArg(scoreint)string {
82+
returnfmt.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+
ifruntime.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+
ifruntime.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+
ifruntime.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)
[8]ページ先頭
©2009-2025 Movatter.jp