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

Commitbbc549d

Browse files
authored
feat: add agent exec pkg (#15577)
1 parent7876dc5 commitbbc549d

File tree

7 files changed

+603
-0
lines changed

7 files changed

+603
-0
lines changed

‎agent/agentexec/cli_linux.go‎

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
//go:build linux
2+
// +build linux
3+
4+
package agentexec
5+
6+
import (
7+
"flag"
8+
"fmt"
9+
"os"
10+
"os/exec"
11+
"runtime"
12+
"strconv"
13+
"strings"
14+
"syscall"
15+
16+
"golang.org/x/sys/unix"
17+
"golang.org/x/xerrors"
18+
)
19+
20+
// unset is set to an invalid value for nice and oom scores.
21+
constunset=-2000
22+
23+
// CLI runs the agent-exec command. It should only be called by the cli package.
24+
funcCLI()error {
25+
// We lock the OS thread here to avoid a race condition where the nice priority
26+
// we get is on a different thread from the one we set it on.
27+
runtime.LockOSThread()
28+
// Nop on success but we do it anyway in case of an error.
29+
deferruntime.UnlockOSThread()
30+
31+
var (
32+
fs=flag.NewFlagSet("agent-exec",flag.ExitOnError)
33+
nice=fs.Int("coder-nice",unset,"")
34+
oom=fs.Int("coder-oom",unset,"")
35+
)
36+
37+
iflen(os.Args)<3 {
38+
returnxerrors.Errorf("malformed command %+v",os.Args)
39+
}
40+
41+
// Parse everything after "coder agent-exec".
42+
err:=fs.Parse(os.Args[2:])
43+
iferr!=nil {
44+
returnxerrors.Errorf("parse flags: %w",err)
45+
}
46+
47+
// Get everything after "coder agent-exec --"
48+
args:=execArgs(os.Args)
49+
iflen(args)==0 {
50+
returnxerrors.Errorf("no exec command provided %+v",os.Args)
51+
}
52+
53+
if*nice==unset {
54+
// If an explicit nice score isn't set, we use the default.
55+
*nice,err=defaultNiceScore()
56+
iferr!=nil {
57+
returnxerrors.Errorf("get default nice score: %w",err)
58+
}
59+
}
60+
61+
if*oom==unset {
62+
// If an explicit oom score isn't set, we use the default.
63+
*oom,err=defaultOOMScore()
64+
iferr!=nil {
65+
returnxerrors.Errorf("get default oom score: %w",err)
66+
}
67+
}
68+
69+
err=unix.Setpriority(unix.PRIO_PROCESS,0,*nice)
70+
iferr!=nil {
71+
returnxerrors.Errorf("set nice score: %w",err)
72+
}
73+
74+
err=writeOOMScoreAdj(*oom)
75+
iferr!=nil {
76+
returnxerrors.Errorf("set oom score: %w",err)
77+
}
78+
79+
path,err:=exec.LookPath(args[0])
80+
iferr!=nil {
81+
returnxerrors.Errorf("look path: %w",err)
82+
}
83+
84+
returnsyscall.Exec(path,args,os.Environ())
85+
}
86+
87+
funcdefaultNiceScore() (int,error) {
88+
score,err:=unix.Getpriority(unix.PRIO_PROCESS,0)
89+
iferr!=nil {
90+
return0,xerrors.Errorf("get nice score: %w",err)
91+
}
92+
// See https://linux.die.net/man/2/setpriority#Notes
93+
score=20-score
94+
95+
score+=5
96+
ifscore>19 {
97+
return19,nil
98+
}
99+
returnscore,nil
100+
}
101+
102+
funcdefaultOOMScore() (int,error) {
103+
score,err:=oomScoreAdj()
104+
iferr!=nil {
105+
return0,xerrors.Errorf("get oom score: %w",err)
106+
}
107+
108+
// If the agent has a negative oom_score_adj, we set the child to 0
109+
// so it's treated like every other process.
110+
ifscore<0 {
111+
return0,nil
112+
}
113+
114+
// If the agent is already almost at the maximum then set it to the max.
115+
ifscore>=998 {
116+
return1000,nil
117+
}
118+
119+
// If the agent oom_score_adj is >=0, we set the child to slightly
120+
// less than the maximum. If users want a different score they set it
121+
// directly.
122+
return998,nil
123+
}
124+
125+
funcoomScoreAdj() (int,error) {
126+
scoreStr,err:=os.ReadFile("/proc/self/oom_score_adj")
127+
iferr!=nil {
128+
return0,xerrors.Errorf("read oom_score_adj: %w",err)
129+
}
130+
returnstrconv.Atoi(strings.TrimSpace(string(scoreStr)))
131+
}
132+
133+
funcwriteOOMScoreAdj(scoreint)error {
134+
returnos.WriteFile("/proc/self/oom_score_adj", []byte(fmt.Sprintf("%d",score)),0o600)
135+
}
136+
137+
// execArgs returns the arguments to pass to syscall.Exec after the "--" delimiter.
138+
funcexecArgs(args []string) []string {
139+
fori,arg:=rangeargs {
140+
ifarg=="--" {
141+
returnargs[i+1:]
142+
}
143+
}
144+
returnnil
145+
}

‎agent/agentexec/cli_linux_test.go‎

Lines changed: 178 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,178 @@
1+
//go:build linux
2+
// +build linux
3+
4+
package agentexec_test
5+
6+
import (
7+
"bytes"
8+
"context"
9+
"fmt"
10+
"os"
11+
"os/exec"
12+
"path/filepath"
13+
"strconv"
14+
"strings"
15+
"syscall"
16+
"testing"
17+
"time"
18+
19+
"github.com/stretchr/testify/require"
20+
"golang.org/x/sys/unix"
21+
22+
"github.com/coder/coder/v2/testutil"
23+
)
24+
25+
funcTestCLI(t*testing.T) {
26+
t.Parallel()
27+
28+
t.Run("OK",func(t*testing.T) {
29+
t.Parallel()
30+
31+
ctx:=testutil.Context(t,testutil.WaitMedium)
32+
cmd,path:=cmd(ctx,t,123,12)
33+
err:=cmd.Start()
34+
require.NoError(t,err)
35+
gocmd.Wait()
36+
37+
waitForSentinel(ctx,t,cmd,path)
38+
requireOOMScore(t,cmd.Process.Pid,123)
39+
requireNiceScore(t,cmd.Process.Pid,12)
40+
})
41+
42+
t.Run("Defaults",func(t*testing.T) {
43+
t.Parallel()
44+
45+
ctx:=testutil.Context(t,testutil.WaitMedium)
46+
cmd,path:=cmd(ctx,t,0,0)
47+
err:=cmd.Start()
48+
require.NoError(t,err)
49+
gocmd.Wait()
50+
51+
waitForSentinel(ctx,t,cmd,path)
52+
53+
expectedNice:=expectedNiceScore(t)
54+
expectedOOM:=expectedOOMScore(t)
55+
requireOOMScore(t,cmd.Process.Pid,expectedOOM)
56+
requireNiceScore(t,cmd.Process.Pid,expectedNice)
57+
})
58+
}
59+
60+
funcrequireNiceScore(t*testing.T,pidint,scoreint) {
61+
t.Helper()
62+
63+
nice,err:=unix.Getpriority(unix.PRIO_PROCESS,pid)
64+
require.NoError(t,err)
65+
// See https://linux.die.net/man/2/setpriority#Notes
66+
require.Equal(t,score,20-nice)
67+
}
68+
69+
funcrequireOOMScore(t*testing.T,pidint,expectedint) {
70+
t.Helper()
71+
72+
actual,err:=os.ReadFile(fmt.Sprintf("/proc/%d/oom_score_adj",pid))
73+
require.NoError(t,err)
74+
score:=strings.TrimSpace(string(actual))
75+
require.Equal(t,strconv.Itoa(expected),score)
76+
}
77+
78+
funcwaitForSentinel(ctx context.Context,t*testing.T,cmd*exec.Cmd,pathstring) {
79+
t.Helper()
80+
81+
ticker:=time.NewTicker(testutil.IntervalFast)
82+
deferticker.Stop()
83+
84+
// RequireEventually doesn't work well with require.NoError or similar require functions.
85+
for {
86+
err:=cmd.Process.Signal(syscall.Signal(0))
87+
require.NoError(t,err)
88+
89+
_,err=os.Stat(path)
90+
iferr==nil {
91+
return
92+
}
93+
94+
select {
95+
case<-ticker.C:
96+
case<-ctx.Done():
97+
require.NoError(t,ctx.Err())
98+
}
99+
}
100+
}
101+
102+
funccmd(ctx context.Context,t*testing.T,oom,niceint) (*exec.Cmd,string) {
103+
var (
104+
args=execArgs(oom,nice)
105+
dir=t.TempDir()
106+
file=filepath.Join(dir,"sentinel")
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,
116+
}
117+
118+
cmd.Env=os.Environ()
119+
varbuf bytes.Buffer
120+
cmd.Stdout=&buf
121+
cmd.Stderr=&buf
122+
t.Cleanup(func() {
123+
// Print output of a command if the test fails.
124+
ift.Failed() {
125+
t.Logf("cmd %q output: %s",cmd.Args,buf.String())
126+
}
127+
ifcmd.Process!=nil {
128+
// We use -cmd.Process.Pid to kill the whole process group.
129+
_=syscall.Kill(-cmd.Process.Pid,syscall.SIGINT)
130+
}
131+
})
132+
returncmd,file
133+
}
134+
135+
funcexpectedOOMScore(t*testing.T)int {
136+
t.Helper()
137+
138+
score,err:=os.ReadFile(fmt.Sprintf("/proc/%d/oom_score_adj",os.Getpid()))
139+
require.NoError(t,err)
140+
141+
scoreInt,err:=strconv.Atoi(strings.TrimSpace(string(score)))
142+
require.NoError(t,err)
143+
144+
ifscoreInt<0 {
145+
return0
146+
}
147+
ifscoreInt>=998 {
148+
return1000
149+
}
150+
return998
151+
}
152+
153+
funcexpectedNiceScore(t*testing.T)int {
154+
t.Helper()
155+
156+
score,err:=unix.Getpriority(unix.PRIO_PROCESS,os.Getpid())
157+
require.NoError(t,err)
158+
159+
// Priority is niceness + 20.
160+
score=20-score
161+
score+=5
162+
ifscore>19 {
163+
return19
164+
}
165+
returnscore
166+
}
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/cli_other.go‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
//go:build !linux
2+
// +build !linux
3+
4+
package agentexec
5+
6+
import"golang.org/x/xerrors"
7+
8+
funcCLI()error {
9+
returnxerrors.New("agent-exec is only supported on Linux")
10+
}
Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
//go:build linux
2+
// +build linux
3+
4+
package main
5+
6+
import (
7+
"fmt"
8+
"os"
9+
10+
"github.com/coder/coder/v2/agent/agentexec"
11+
)
12+
13+
funcmain() {
14+
err:=agentexec.CLI()
15+
iferr!=nil {
16+
_,_=fmt.Fprintln(os.Stderr,err)
17+
os.Exit(1)
18+
}
19+
}

0 commit comments

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