Skip to content

Commit b1c0b39

Browse files
mafredrimafredri
authored
feat(agent): add script data dir for binaries and files (#12205)
The agent is extended with a `--script-data-dir` flag, defaulting to the OS temp dir. This dir is used for storing `coder-script-data/bin` and `coder-script/[script uuid]`. The former is a place for all scripts to place executable binaries that will be available by other scripts, SSH sessions, etc. The latter is a place for the script to store files. Since we default to OS temp dir, files are ephemeral by default. In the future, we may consider adding new env vars or changing the default storage location. Workspace startup speed could potentially benefit from scripts being able to skip steps that require downloading software. We may also extend this with more env variables (e.g. persistent storage in HOME). Fixes #11131
1 parent ab4cb66 commit b1c0b39

File tree

6 files changed

+162
-20
lines changed

6 files changed

+162
-20
lines changed

‎agent/agent.go

Lines changed: 26 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ type Options struct {
6666
Filesystem afero.Fs
6767
LogDir string
6868
TempDir string
69+
ScriptDataDir string
6970
ExchangeToken func(ctx context.Context) (string, error)
7071
Client Client
7172
ReconnectingPTYTimeout time.Duration
@@ -112,9 +113,19 @@ func New(options Options) Agent {
112113
if options.LogDir == "" {
113114
if options.TempDir != os.TempDir() {
114115
options.Logger.Debug(context.Background(), "log dir not set, using temp dir", slog.F("temp_dir", options.TempDir))
116+
} else {
117+
options.Logger.Debug(context.Background(), "using log dir", slog.F("log_dir", options.LogDir))
115118
}
116119
options.LogDir = options.TempDir
117120
}
121+
if options.ScriptDataDir == "" {
122+
if options.TempDir != os.TempDir() {
123+
options.Logger.Debug(context.Background(), "script data dir not set, using temp dir", slog.F("temp_dir", options.TempDir))
124+
} else {
125+
options.Logger.Debug(context.Background(), "using script data dir", slog.F("script_data_dir", options.ScriptDataDir))
126+
}
127+
options.ScriptDataDir = options.TempDir
128+
}
118129
if options.ExchangeToken == nil {
119130
options.ExchangeToken = func(ctx context.Context) (string, error) {
120131
return "", nil
@@ -152,6 +163,7 @@ func New(options Options) Agent {
152163
filesystem: options.Filesystem,
153164
logDir: options.LogDir,
154165
tempDir: options.TempDir,
166+
scriptDataDir: options.ScriptDataDir,
155167
lifecycleUpdate: make(chan struct{}, 1),
156168
lifecycleReported: make(chan codersdk.WorkspaceAgentLifecycle, 1),
157169
lifecycleStates: []agentsdk.PostLifecycleRequest{{State: codersdk.WorkspaceAgentLifecycleCreated}},
@@ -183,6 +195,7 @@ type agent struct {
183195
filesystem afero.Fs
184196
logDir string
185197
tempDir string
198+
scriptDataDir string
186199
// ignorePorts tells the api handler which ports to ignore when
187200
// listing all listening ports. This is helpful to hide ports that
188201
// are used by the agent, that the user does not care about.
@@ -249,11 +262,12 @@ func (a *agent) init(ctx context.Context) {
249262
}
250263
a.sshServer = sshSrv
251264
a.scriptRunner = agentscripts.New(agentscripts.Options{
252-
LogDir: a.logDir,
253-
Logger: a.logger,
254-
SSHServer: sshSrv,
255-
Filesystem: a.filesystem,
256-
PatchLogs: a.client.PatchLogs,
265+
LogDir: a.logDir,
266+
DataDirBase: a.scriptDataDir,
267+
Logger: a.logger,
268+
SSHServer: sshSrv,
269+
Filesystem: a.filesystem,
270+
PatchLogs: a.client.PatchLogs,
257271
})
258272
// Register runner metrics. If the prom registry is nil, the metrics
259273
// will not report anywhere.
@@ -954,6 +968,13 @@ func (a *agent) updateCommandEnv(current []string) (updated []string, err error)
954968
envs[k] = v
955969
}
956970

971+
// Prepend the agent script bin directory to the PATH
972+
// (this is where Coder modules place their binaries).
973+
if _, ok := envs["PATH"]; !ok {
974+
envs["PATH"] = os.Getenv("PATH")
975+
}
976+
envs["PATH"] = fmt.Sprintf("%s%c%s", a.scriptRunner.ScriptBinDir(), filepath.ListSeparator, envs["PATH"])
977+
957978
for k, v := range envs {
958979
updated = append(updated, fmt.Sprintf("%s=%s", k, v))
959980
}

‎agent/agent_test.go

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -286,6 +286,12 @@ func TestAgent_SessionExec(t *testing.T) {
286286
func TestAgent_Session_EnvironmentVariables(t *testing.T) {
287287
t.Parallel()
288288

289+
tmpdir := t.TempDir()
290+
291+
// Defined by the coder script runner, hardcoded here since we don't
292+
// have a reference to it.
293+
scriptBinDir := filepath.Join(tmpdir, "coder-script-data", "bin")
294+
289295
manifest := agentsdk.Manifest{
290296
EnvironmentVariables: map[string]string{
291297
"MY_MANIFEST": "true",
@@ -295,6 +301,7 @@ func TestAgent_Session_EnvironmentVariables(t *testing.T) {
295301
}
296302
banner := codersdk.ServiceBannerConfig{}
297303
session := setupSSHSession(t, manifest, banner, nil, func(_ *agenttest.Client, opts *agent.Options) {
304+
opts.ScriptDataDir = tmpdir
298305
opts.EnvironmentVariables["MY_OVERRIDE"] = "true"
299306
})
300307

@@ -341,6 +348,7 @@ func TestAgent_Session_EnvironmentVariables(t *testing.T) {
341348
"MY_OVERRIDE": "true", // From the agent environment variables option, overrides manifest.
342349
"MY_SESSION_MANIFEST": "false", // From the manifest, overrides session env.
343350
"MY_SESSION": "true", // From the session.
351+
"PATH": scriptBinDir + string(filepath.ListSeparator),
344352
} {
345353
t.Run(k, func(t *testing.T) {
346354
echoEnv(t, stdin, k)

‎agent/agentscripts/agentscripts.go

Lines changed: 43 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -43,11 +43,12 @@ var (
4343

4444
// Options are a set of options for the runner.
4545
type Options struct {
46-
LogDir string
47-
Logger slog.Logger
48-
SSHServer *agentssh.Server
49-
Filesystem afero.Fs
50-
PatchLogs func(ctx context.Context, req agentsdk.PatchLogs) error
46+
DataDirBase string
47+
LogDir string
48+
Logger slog.Logger
49+
SSHServer *agentssh.Server
50+
Filesystem afero.Fs
51+
PatchLogs func(ctx context.Context, req agentsdk.PatchLogs) error
5152
}
5253

5354
// New creates a runner for the provided scripts.
@@ -59,6 +60,7 @@ func New(opts Options) *Runner {
5960
cronCtxCancel: cronCtxCancel,
6061
cron: cron.New(cron.WithParser(parser)),
6162
closed: make(chan struct{}),
63+
dataDir: filepath.Join(opts.DataDirBase, "coder-script-data"),
6264
scriptsExecuted: prometheus.NewCounterVec(prometheus.CounterOpts{
6365
Namespace: "agent",
6466
Subsystem: "scripts",
@@ -78,13 +80,25 @@ type Runner struct {
7880
cron *cron.Cron
7981
initialized atomic.Bool
8082
scripts []codersdk.WorkspaceAgentScript
83+
dataDir string
8184

8285
// scriptsExecuted includes all scripts executed by the workspace agent. Agents
8386
// execute startup scripts, and scripts on a cron schedule. Both will increment
8487
// this counter.
8588
scriptsExecuted *prometheus.CounterVec
8689
}
8790

91+
// DataDir returns the directory where scripts data is stored.
92+
func (r *Runner) DataDir() string {
93+
return r.dataDir
94+
}
95+
96+
// ScriptBinDir returns the directory where scripts can store executable
97+
// binaries.
98+
func (r *Runner) ScriptBinDir() string {
99+
return filepath.Join(r.dataDir, "bin")
100+
}
101+
88102
func (r *Runner) RegisterMetrics(reg prometheus.Registerer) {
89103
if reg == nil {
90104
// If no registry, do nothing.
@@ -104,6 +118,11 @@ func (r *Runner) Init(scripts []codersdk.WorkspaceAgentScript) error {
104118
r.scripts = scripts
105119
r.Logger.Info(r.cronCtx, "initializing agent scripts", slog.F("script_count", len(scripts)), slog.F("log_dir", r.LogDir))
106120

121+
err := r.Filesystem.MkdirAll(r.ScriptBinDir(), 0o700)
122+
if err != nil {
123+
return xerrors.Errorf("create script bin dir: %w", err)
124+
}
125+
107126
for _, script := range scripts {
108127
if script.Cron == "" {
109128
continue
@@ -208,7 +227,18 @@ func (r *Runner) run(ctx context.Context, script codersdk.WorkspaceAgentScript)
208227
if !filepath.IsAbs(logPath) {
209228
logPath = filepath.Join(r.LogDir, logPath)
210229
}
211-
logger := r.Logger.With(slog.F("log_path", logPath))
230+
231+
scriptDataDir := filepath.Join(r.DataDir(), script.LogSourceID.String())
232+
err := r.Filesystem.MkdirAll(scriptDataDir, 0o700)
233+
if err != nil {
234+
return xerrors.Errorf("%s script: create script temp dir: %w", scriptDataDir, err)
235+
}
236+
237+
logger := r.Logger.With(
238+
slog.F("log_source_id", script.LogSourceID),
239+
slog.F("log_path", logPath),
240+
slog.F("script_data_dir", scriptDataDir),
241+
)
212242
logger.Info(ctx, "running agent script", slog.F("script", script.Script))
213243

214244
fileWriter, err := r.Filesystem.OpenFile(logPath, os.O_CREATE|os.O_RDWR, 0o600)
@@ -238,6 +268,13 @@ func (r *Runner) run(ctx context.Context, script codersdk.WorkspaceAgentScript)
238268
cmd.WaitDelay = 10 * time.Second
239269
cmd.Cancel = cmdCancel(cmd)
240270

271+
// Expose env vars that can be used in the script for storing data
272+
// and binaries. In the future, we may want to expose more env vars
273+
// for the script to use, like CODER_SCRIPT_DATA_DIR for persistent
274+
// storage.
275+
cmd.Env = append(cmd.Env, "CODER_SCRIPT_DATA_DIR="+scriptDataDir)
276+
cmd.Env = append(cmd.Env, "CODER_SCRIPT_BIN_DIR="+r.ScriptBinDir())
277+
241278
send, flushAndClose := agentsdk.LogsSender(script.LogSourceID, r.PatchLogs, logger)
242279
// If ctx is canceled here (or in a writer below), we may be
243280
// discarding logs, but that's okay because we're shutting down

‎agent/agentscripts/agentscripts_test.go

Lines changed: 73 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,15 @@ package agentscripts_test
22

33
import (
44
"context"
5+
"path/filepath"
6+
"runtime"
57
"testing"
68
"time"
79

10+
"github.com/google/uuid"
811
"github.com/prometheus/client_golang/prometheus"
912
"github.com/spf13/afero"
13+
"github.com/stretchr/testify/assert"
1014
"github.com/stretchr/testify/require"
1115
"go.uber.org/goleak"
1216

@@ -15,6 +19,7 @@ import (
1519
"github.com/coder/coder/v2/agent/agentssh"
1620
"github.com/coder/coder/v2/codersdk"
1721
"github.com/coder/coder/v2/codersdk/agentsdk"
22+
"github.com/coder/coder/v2/testutil"
1823
)
1924

2025
func TestMain(m *testing.M) {
@@ -25,12 +30,16 @@ func TestExecuteBasic(t *testing.T) {
2530
t.Parallel()
2631
logs := make(chan agentsdk.PatchLogs, 1)
2732
runner := setup(t, func(ctx context.Context, req agentsdk.PatchLogs) error {
28-
logs <- req
33+
select {
34+
case <-ctx.Done():
35+
case logs <- req:
36+
}
2937
return nil
3038
})
3139
defer runner.Close()
3240
err := runner.Init([]codersdk.WorkspaceAgentScript{{
33-
Script: "echo hello",
41+
LogSourceID: uuid.New(),
42+
Script: "echo hello",
3443
}})
3544
require.NoError(t, err)
3645
require.NoError(t, runner.Execute(context.Background(), func(script codersdk.WorkspaceAgentScript) bool {
@@ -40,13 +49,67 @@ func TestExecuteBasic(t *testing.T) {
4049
require.Equal(t, "hello", log.Logs[0].Output)
4150
}
4251

52+
func TestEnv(t *testing.T) {
53+
t.Parallel()
54+
logs := make(chan agentsdk.PatchLogs, 2)
55+
runner := setup(t, func(ctx context.Context, req agentsdk.PatchLogs) error {
56+
select {
57+
case <-ctx.Done():
58+
case logs <- req:
59+
}
60+
return nil
61+
})
62+
defer runner.Close()
63+
id := uuid.New()
64+
script := "echo $CODER_SCRIPT_DATA_DIR\necho $CODER_SCRIPT_BIN_DIR\n"
65+
if runtime.GOOS == "windows" {
66+
script = `
67+
cmd.exe /c echo %CODER_SCRIPT_DATA_DIR%
68+
cmd.exe /c echo %CODER_SCRIPT_BIN_DIR%
69+
`
70+
}
71+
err := runner.Init([]codersdk.WorkspaceAgentScript{{
72+
LogSourceID: id,
73+
Script: script,
74+
}})
75+
require.NoError(t, err)
76+
77+
ctx := testutil.Context(t, testutil.WaitLong)
78+
79+
testutil.Go(t, func() {
80+
err := runner.Execute(ctx, func(script codersdk.WorkspaceAgentScript) bool {
81+
return true
82+
})
83+
assert.NoError(t, err)
84+
})
85+
86+
var log []agentsdk.Log
87+
for {
88+
select {
89+
case <-ctx.Done():
90+
require.Fail(t, "timed out waiting for logs")
91+
case l := <-logs:
92+
for _, l := range l.Logs {
93+
t.Logf("log: %s", l.Output)
94+
}
95+
log = append(log, l.Logs...)
96+
}
97+
if len(log) >= 2 {
98+
break
99+
}
100+
}
101+
require.Contains(t, log[0].Output, filepath.Join(runner.DataDir(), id.String()))
102+
require.Contains(t, log[1].Output, runner.ScriptBinDir())
103+
}
104+
43105
func TestTimeout(t *testing.T) {
44106
t.Parallel()
45107
runner := setup(t, nil)
46108
defer runner.Close()
47109
err := runner.Init([]codersdk.WorkspaceAgentScript{{
48-
Script: "sleep infinity",
49-
Timeout: time.Millisecond,
110+
LogSourceID: uuid.New(),
111+
Script: "sleep infinity",
112+
Timeout: time.Millisecond,
50113
}})
51114
require.NoError(t, err)
52115
require.ErrorIs(t, runner.Execute(context.Background(), nil), agentscripts.ErrTimeout)
@@ -77,10 +140,11 @@ func setup(t *testing.T, patchLogs func(ctx context.Context, req agentsdk.PatchL
77140
_ = s.Close()
78141
})
79142
return agentscripts.New(agentscripts.Options{
80-
LogDir: t.TempDir(),
81-
Logger: logger,
82-
SSHServer: s,
83-
Filesystem: fs,
84-
PatchLogs: patchLogs,
143+
LogDir: t.TempDir(),
144+
DataDirBase: t.TempDir(),
145+
Logger: logger,
146+
SSHServer: s,
147+
Filesystem: fs,
148+
PatchLogs: patchLogs,
85149
})
86150
}

‎cli/agent.go

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
4040
var (
4141
auth string
4242
logDir string
43+
scriptDataDir string
4344
pprofAddress string
4445
noReap bool
4546
sshMaxTimeout time.Duration
@@ -289,6 +290,7 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
289290
Client: client,
290291
Logger: logger,
291292
LogDir: logDir,
293+
ScriptDataDir: scriptDataDir,
292294
TailnetListenPort: uint16(tailnetListenPort),
293295
ExchangeToken: func(ctx context.Context) (string, error) {
294296
if exchangeToken == nil {
@@ -339,6 +341,13 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
339341
Env: "CODER_AGENT_LOG_DIR",
340342
Value: clibase.StringOf(&logDir),
341343
},
344+
{
345+
Flag: "script-data-dir",
346+
Default: os.TempDir(),
347+
Description: "Specify the location for storing script data.",
348+
Env: "CODER_AGENT_SCRIPT_DATA_DIR",
349+
Value: clibase.StringOf(&scriptDataDir),
350+
},
342351
{
343352
Flag: "pprof-address",
344353
Default: "127.0.0.1:6060",

‎cli/testdata/coder_agent_--help.golden

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,9 @@ OPTIONS:
3333
--prometheus-address string, $CODER_AGENT_PROMETHEUS_ADDRESS (default: 127.0.0.1:2112)
3434
The bind address to serve Prometheus metrics.
3535

36+
--script-data-dir string, $CODER_AGENT_SCRIPT_DATA_DIR (default: /tmp)
37+
Specify the location for storing script data.
38+
3639
--ssh-max-timeout duration, $CODER_AGENT_SSH_MAX_TIMEOUT (default: 72h)
3740
Specify the max timeout for a SSH connection, it is advisable to set
3841
it to a minimum of 60s, but no more than 72h.

0 commit comments

Comments
 (0)