Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit bbc549d

Browse files
sreyasreya
authored
feat: add agent exec pkg (#15577)
1 parent 7876dc5 commit bbc549d

File tree

7 files changed

+603
-0
lines changed

7 files changed

+603
-0
lines changed

‎agent/agentexec/cli_linux.go

+145
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+
const unset = -2000
22+
23+
// CLI runs the agent-exec command. It should only be called by the cli package.
24+
func CLI() 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+
defer runtime.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+
if len(os.Args) < 3 {
38+
return xerrors.Errorf("malformed command %+v", os.Args)
39+
}
40+
41+
// Parse everything after "coder agent-exec".
42+
err := fs.Parse(os.Args[2:])
43+
if err != nil {
44+
return xerrors.Errorf("parse flags: %w", err)
45+
}
46+
47+
// Get everything after "coder agent-exec --"
48+
args := execArgs(os.Args)
49+
if len(args) == 0 {
50+
return xerrors.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+
if err != nil {
57+
return xerrors.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+
if err != nil {
65+
return xerrors.Errorf("get default oom score: %w", err)
66+
}
67+
}
68+
69+
err = unix.Setpriority(unix.PRIO_PROCESS, 0, *nice)
70+
if err != nil {
71+
return xerrors.Errorf("set nice score: %w", err)
72+
}
73+
74+
err = writeOOMScoreAdj(*oom)
75+
if err != nil {
76+
return xerrors.Errorf("set oom score: %w", err)
77+
}
78+
79+
path, err := exec.LookPath(args[0])
80+
if err != nil {
81+
return xerrors.Errorf("look path: %w", err)
82+
}
83+
84+
return syscall.Exec(path, args, os.Environ())
85+
}
86+
87+
func defaultNiceScore() (int, error) {
88+
score, err := unix.Getpriority(unix.PRIO_PROCESS, 0)
89+
if err != nil {
90+
return 0, 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+
if score > 19 {
97+
return 19, nil
98+
}
99+
return score, nil
100+
}
101+
102+
func defaultOOMScore() (int, error) {
103+
score, err := oomScoreAdj()
104+
if err != nil {
105+
return 0, 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+
if score < 0 {
111+
return 0, nil
112+
}
113+
114+
// If the agent is already almost at the maximum then set it to the max.
115+
if score >= 998 {
116+
return 1000, 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+
return 998, nil
123+
}
124+
125+
func oomScoreAdj() (int, error) {
126+
scoreStr, err := os.ReadFile("/proc/self/oom_score_adj")
127+
if err != nil {
128+
return 0, xerrors.Errorf("read oom_score_adj: %w", err)
129+
}
130+
return strconv.Atoi(strings.TrimSpace(string(scoreStr)))
131+
}
132+
133+
func writeOOMScoreAdj(score int) error {
134+
return os.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+
func execArgs(args []string) []string {
139+
for i, arg := range args {
140+
if arg == "--" {
141+
return args[i+1:]
142+
}
143+
}
144+
return nil
145+
}

‎agent/agentexec/cli_linux_test.go

+178
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+
func TestCLI(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+
go cmd.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+
go cmd.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+
func requireNiceScore(t *testing.T, pid int, score int) {
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+
func requireOOMScore(t *testing.T, pid int, expected int) {
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+
func waitForSentinel(ctx context.Context, t *testing.T, cmd *exec.Cmd, path string) {
79+
t.Helper()
80+
81+
ticker := time.NewTicker(testutil.IntervalFast)
82+
defer ticker.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+
if err == 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+
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()
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+
var buf 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+
if t.Failed() {
125+
t.Logf("cmd %q output: %s", cmd.Args, buf.String())
126+
}
127+
if cmd.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+
return cmd, file
133+
}
134+
135+
func expectedOOMScore(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+
if scoreInt < 0 {
145+
return 0
146+
}
147+
if scoreInt >= 998 {
148+
return 1000
149+
}
150+
return 998
151+
}
152+
153+
func expectedNiceScore(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+
if score > 19 {
163+
return 19
164+
}
165+
return score
166+
}
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/cli_other.go

+10
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+
func CLI() error {
9+
return xerrors.New("agent-exec is only supported on Linux")
10+
}

‎agent/agentexec/cmdtest/main_linux.go

+19
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+
func main() {
14+
err := agentexec.CLI()
15+
if err != nil {
16+
_, _ = fmt.Fprintln(os.Stderr, err)
17+
os.Exit(1)
18+
}
19+
}

0 commit comments

Comments
 (0)