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

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+
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

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+
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

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

‎agent/agentexec/cmdtest/main_linux.go

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+
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)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy