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

Skip to content

Commit ce573b9

Browse files
sreyasreya
authored
fix: add agent exec abstraction (coder#15717)
1 parent 6c9ccca commit ce573b9

16 files changed

+210
-192
lines changed

‎agent/agent.go

+9-1
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ import (
3333
"tailscale.com/util/clientmetric"
3434

3535
"cdr.dev/slog"
36+
"github.com/coder/coder/v2/agent/agentexec"
3637
"github.com/coder/coder/v2/agent/agentscripts"
3738
"github.com/coder/coder/v2/agent/agentssh"
3839
"github.com/coder/coder/v2/agent/proto"
@@ -80,6 +81,7 @@ type Options struct {
8081
ReportMetadataInterval time.Duration
8182
ServiceBannerRefreshInterval time.Duration
8283
BlockFileTransfer bool
84+
Execer agentexec.Execer
8385
}
8486

8587
type Client interface {
@@ -139,6 +141,10 @@ func New(options Options) Agent {
139141
prometheusRegistry = prometheus.NewRegistry()
140142
}
141143

144+
if options.Execer == nil {
145+
options.Execer = agentexec.DefaultExecer
146+
}
147+
142148
hardCtx, hardCancel := context.WithCancel(context.Background())
143149
gracefulCtx, gracefulCancel := context.WithCancel(hardCtx)
144150
a := &agent{
@@ -171,6 +177,7 @@ func New(options Options) Agent {
171177

172178
prometheusRegistry: prometheusRegistry,
173179
metrics: newAgentMetrics(prometheusRegistry),
180+
execer: options.Execer,
174181
}
175182
// Initially, we have a closed channel, reflecting the fact that we are not initially connected.
176183
// Each time we connect we replace the channel (while holding the closeMutex) with a new one
@@ -239,6 +246,7 @@ type agent struct {
239246
// metrics are prometheus registered metrics that will be collected and
240247
// labeled in Coder with the agent + workspace.
241248
metrics *agentMetrics
249+
execer agentexec.Execer
242250
}
243251

244252
func (a *agent) TailnetConn() *tailnet.Conn {
@@ -247,7 +255,7 @@ func (a *agent) TailnetConn() *tailnet.Conn {
247255

248256
func (a *agent) init() {
249257
// pass the "hard" context because we explicitly close the SSH server as part of graceful shutdown.
250-
sshSrv, err := agentssh.NewServer(a.hardCtx, a.logger.Named("ssh-server"), a.prometheusRegistry, a.filesystem, &agentssh.Config{
258+
sshSrv, err := agentssh.NewServer(a.hardCtx, a.logger.Named("ssh-server"), a.prometheusRegistry, a.filesystem, a.execer, &agentssh.Config{
251259
MaxTimeout: a.sshMaxTimeout,
252260
MOTDFile: func() string { return a.manifest.Load().MOTDFile },
253261
AnnouncementBanners: func() *[]codersdk.BannerConfig { return a.announcementBanners.Load() },

‎agent/agentexec/cli_linux.go

-3
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,6 @@ import (
1717
"golang.org/x/xerrors"
1818
)
1919

20-
// unset is set to an invalid value for nice and oom scores.
21-
const unset = -2000
22-
2320
// CLI runs the agent-exec command. It should only be called by the cli package.
2421
func CLI() error {
2522
// We lock the OS thread here to avoid a race condition where the nice priority

‎agent/agentexec/exec.go

+72-31
Original file line numberDiff line numberDiff line change
@@ -20,60 +20,101 @@ const (
2020
EnvProcPrioMgmt = "CODER_PROC_PRIO_MGMT"
2121
EnvProcOOMScore = "CODER_PROC_OOM_SCORE"
2222
EnvProcNiceScore = "CODER_PROC_NICE_SCORE"
23-
)
2423

25-
// CommandContext returns an exec.Cmd that calls "coder agent-exec" prior to exec'ing
26-
// the provided command if CODER_PROC_PRIO_MGMT is set, otherwise a normal exec.Cmd
27-
// is returned. All instances of exec.Cmd should flow through this function to ensure
28-
// proper resource constraints are applied to the child process.
29-
func CommandContext(ctx context.Context, cmd string, args ...string) (*exec.Cmd, error) {
30-
cmd, args, err := agentExecCmd(cmd, args...)
31-
if err != nil {
32-
return nil, xerrors.Errorf("agent exec cmd: %w", err)
33-
}
34-
return exec.CommandContext(ctx, cmd, args...), nil
35-
}
24+
// unset is set to an invalid value for nice and oom scores.
25+
unset = -2000
26+
)
3627

37-
// PTYCommandContext returns an pty.Cmd that calls "coder agent-exec" prior to exec'ing
38-
// the provided command if CODER_PROC_PRIO_MGMT is set, otherwise a normal pty.Cmd
39-
// is returned. All instances of pty.Cmd should flow through this function to ensure
40-
// proper resource constraints are applied to the child process.
41-
func PTYCommandContext(ctx context.Context, cmd string, args ...string) (*pty.Cmd, error) {
42-
cmd, args, err := agentExecCmd(cmd, args...)
43-
if err != nil {
44-
return nil, xerrors.Errorf("agent exec cmd: %w", err)
45-
}
46-
return pty.CommandContext(ctx, cmd, args...), nil
28+
var DefaultExecer Execer = execer{}
29+
30+
// Execer defines an abstraction for creating exec.Cmd variants. It's unfortunately
31+
// necessary because we need to be able to wrap child processes with "coder agent-exec"
32+
// for templates that expect the agent to manage process priority.
33+
type Execer interface {
34+
// CommandContext returns an exec.Cmd that calls "coder agent-exec" prior to exec'ing
35+
// the provided command if CODER_PROC_PRIO_MGMT is set, otherwise a normal exec.Cmd
36+
// is returned. All instances of exec.Cmd should flow through this function to ensure
37+
// proper resource constraints are applied to the child process.
38+
CommandContext(ctx context.Context, cmd string, args ...string) *exec.Cmd
39+
// PTYCommandContext returns an pty.Cmd that calls "coder agent-exec" prior to exec'ing
40+
// the provided command if CODER_PROC_PRIO_MGMT is set, otherwise a normal pty.Cmd
41+
// is returned. All instances of pty.Cmd should flow through this function to ensure
42+
// proper resource constraints are applied to the child process.
43+
PTYCommandContext(ctx context.Context, cmd string, args ...string) *pty.Cmd
4744
}
4845

49-
func agentExecCmd(cmd string, args ...string) (string, []string, error) {
46+
func NewExecer() (Execer, error) {
5047
_, enabled := os.LookupEnv(EnvProcPrioMgmt)
5148
if runtime.GOOS != "linux" || !enabled {
52-
return cmd, args, nil
49+
return DefaultExecer, nil
5350
}
5451

5552
executable, err := os.Executable()
5653
if err != nil {
57-
return "", nil, xerrors.Errorf("get executable: %w", err)
54+
return nil, xerrors.Errorf("get executable: %w", err)
5855
}
5956

6057
bin, err := filepath.EvalSymlinks(executable)
6158
if err != nil {
62-
return "", nil, xerrors.Errorf("eval symlinks: %w", err)
59+
return nil, xerrors.Errorf("eval symlinks: %w", err)
60+
}
61+
62+
oomScore, ok := envValInt(EnvProcOOMScore)
63+
if !ok {
64+
oomScore = unset
65+
}
66+
67+
niceScore, ok := envValInt(EnvProcNiceScore)
68+
if !ok {
69+
niceScore = unset
6370
}
6471

72+
return priorityExecer{
73+
binPath: bin,
74+
oomScore: oomScore,
75+
niceScore: niceScore,
76+
}, nil
77+
}
78+
79+
type execer struct{}
80+
81+
func (execer) CommandContext(ctx context.Context, cmd string, args ...string) *exec.Cmd {
82+
return exec.CommandContext(ctx, cmd, args...)
83+
}
84+
85+
func (execer) PTYCommandContext(ctx context.Context, cmd string, args ...string) *pty.Cmd {
86+
return pty.CommandContext(ctx, cmd, args...)
87+
}
88+
89+
type priorityExecer struct {
90+
binPath string
91+
oomScore int
92+
niceScore int
93+
}
94+
95+
func (e priorityExecer) CommandContext(ctx context.Context, cmd string, args ...string) *exec.Cmd {
96+
cmd, args = e.agentExecCmd(cmd, args...)
97+
return exec.CommandContext(ctx, cmd, args...)
98+
}
99+
100+
func (e priorityExecer) PTYCommandContext(ctx context.Context, cmd string, args ...string) *pty.Cmd {
101+
cmd, args = e.agentExecCmd(cmd, args...)
102+
return pty.CommandContext(ctx, cmd, args...)
103+
}
104+
105+
func (e priorityExecer) agentExecCmd(cmd string, args ...string) (string, []string) {
65106
execArgs := []string{"agent-exec"}
66-
if score, ok := envValInt(EnvProcOOMScore); ok {
67-
execArgs = append(execArgs, oomScoreArg(score))
107+
if e.oomScore != unset {
108+
execArgs = append(execArgs, oomScoreArg(e.oomScore))
68109
}
69110

70-
if score, ok := envValInt(EnvProcNiceScore); ok {
71-
execArgs = append(execArgs, niceScoreArg(score))
111+
if e.niceScore != unset {
112+
execArgs = append(execArgs, niceScoreArg(e.niceScore))
72113
}
73114
execArgs = append(execArgs, "--", cmd)
74115
execArgs = append(execArgs, args...)
75116

76-
return bin, execArgs, nil
117+
return e.binPath, execArgs
77118
}
78119

79120
// envValInt searches for a key in a list of environment variables and parses it to an int.

‎agent/agentexec/exec_internal_test.go

+84
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
package agentexec
2+
3+
import (
4+
"context"
5+
"os/exec"
6+
"testing"
7+
8+
"github.com/stretchr/testify/require"
9+
)
10+
11+
func TestExecer(t *testing.T) {
12+
t.Parallel()
13+
14+
t.Run("Default", func(t *testing.T) {
15+
t.Parallel()
16+
17+
cmd := DefaultExecer.CommandContext(context.Background(), "sh", "-c", "sleep")
18+
19+
path, err := exec.LookPath("sh")
20+
require.NoError(t, err)
21+
require.Equal(t, path, cmd.Path)
22+
require.Equal(t, []string{"sh", "-c", "sleep"}, cmd.Args)
23+
})
24+
25+
t.Run("Priority", func(t *testing.T) {
26+
t.Parallel()
27+
28+
t.Run("OK", func(t *testing.T) {
29+
t.Parallel()
30+
31+
e := priorityExecer{
32+
binPath: "/foo/bar/baz",
33+
oomScore: unset,
34+
niceScore: unset,
35+
}
36+
37+
cmd := e.CommandContext(context.Background(), "sh", "-c", "sleep")
38+
require.Equal(t, e.binPath, cmd.Path)
39+
require.Equal(t, []string{e.binPath, "agent-exec", "--", "sh", "-c", "sleep"}, cmd.Args)
40+
})
41+
42+
t.Run("Nice", func(t *testing.T) {
43+
t.Parallel()
44+
45+
e := priorityExecer{
46+
binPath: "/foo/bar/baz",
47+
oomScore: unset,
48+
niceScore: 10,
49+
}
50+
51+
cmd := e.CommandContext(context.Background(), "sh", "-c", "sleep")
52+
require.Equal(t, e.binPath, cmd.Path)
53+
require.Equal(t, []string{e.binPath, "agent-exec", "--coder-nice=10", "--", "sh", "-c", "sleep"}, cmd.Args)
54+
})
55+
56+
t.Run("OOM", func(t *testing.T) {
57+
t.Parallel()
58+
59+
e := priorityExecer{
60+
binPath: "/foo/bar/baz",
61+
oomScore: 123,
62+
niceScore: unset,
63+
}
64+
65+
cmd := e.CommandContext(context.Background(), "sh", "-c", "sleep")
66+
require.Equal(t, e.binPath, cmd.Path)
67+
require.Equal(t, []string{e.binPath, "agent-exec", "--coder-oom=123", "--", "sh", "-c", "sleep"}, cmd.Args)
68+
})
69+
70+
t.Run("Both", func(t *testing.T) {
71+
t.Parallel()
72+
73+
e := priorityExecer{
74+
binPath: "/foo/bar/baz",
75+
oomScore: 432,
76+
niceScore: 14,
77+
}
78+
79+
cmd := e.CommandContext(context.Background(), "sh", "-c", "sleep")
80+
require.Equal(t, e.binPath, cmd.Path)
81+
require.Equal(t, []string{e.binPath, "agent-exec", "--coder-oom=432", "--coder-nice=14", "--", "sh", "-c", "sleep"}, cmd.Args)
82+
})
83+
})
84+
}

‎agent/agentexec/exec_test.go

-119
This file was deleted.

0 commit comments

Comments
 (0)