Skip to content

Commit e804a03

Browse files
mafredrimafredri
committed
feat(agent): add script data dir for binaries and files
The agent is extended with a `--script-data-dir` flag, defaulting to the OS temp dir. This dir is used for storing `coder-script/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 2fe0a69 commit e804a03

File tree

4 files changed

+68
-1
lines changed

4 files changed

+68
-1
lines changed

‎agent/agent.go

Lines changed: 17 additions & 0 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
@@ -115,6 +116,12 @@ func New(options Options) Agent {
115116
}
116117
options.LogDir = options.TempDir
117118
}
119+
if options.ScriptDataDir == "" {
120+
if options.TempDir != os.TempDir() {
121+
options.Logger.Debug(context.Background(), "script data dir not set, using temp dir", slog.F("temp_dir", options.TempDir))
122+
}
123+
options.ScriptDataDir = options.TempDir
124+
}
118125
if options.ExchangeToken == nil {
119126
options.ExchangeToken = func(ctx context.Context) (string, error) {
120127
return "", nil
@@ -152,6 +159,7 @@ func New(options Options) Agent {
152159
filesystem: options.Filesystem,
153160
logDir: options.LogDir,
154161
tempDir: options.TempDir,
162+
scriptDataDir: options.ScriptDataDir,
155163
lifecycleUpdate: make(chan struct{}, 1),
156164
lifecycleReported: make(chan codersdk.WorkspaceAgentLifecycle, 1),
157165
lifecycleStates: []agentsdk.PostLifecycleRequest{{State: codersdk.WorkspaceAgentLifecycleCreated}},
@@ -182,6 +190,7 @@ type agent struct {
182190
filesystem afero.Fs
183191
logDir string
184192
tempDir string
193+
scriptDataDir string
185194
// ignorePorts tells the api handler which ports to ignore when
186195
// listing all listening ports. This is helpful to hide ports that
187196
// are used by the agent, that the user does not care about.
@@ -249,6 +258,7 @@ func (a *agent) init(ctx context.Context) {
249258
a.sshServer = sshSrv
250259
a.scriptRunner = agentscripts.New(agentscripts.Options{
251260
LogDir: a.logDir,
261+
DataDir: a.scriptDataDir,
252262
Logger: a.logger,
253263
SSHServer: sshSrv,
254264
Filesystem: a.filesystem,
@@ -953,6 +963,13 @@ func (a *agent) updateCommandEnv(current []string) (updated []string, err error)
953963
envs[k] = v
954964
}
955965

966+
// Prepend the agent script bin directory to the PATH
967+
// (this is where Coder modules place their binaries).
968+
if _, ok := envs["PATH"]; !ok {
969+
envs["PATH"] = os.Getenv("PATH")
970+
}
971+
envs["PATH"] = a.scriptRunner.ScriptBinDir() + ":" + envs["PATH"]
972+
956973
for k, v := range envs {
957974
updated = append(updated, fmt.Sprintf("%s=%s", k, v))
958975
}

‎agent/agent_test.go

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -285,6 +285,11 @@ func TestAgent_SessionExec(t *testing.T) {
285285
func TestAgent_Session_EnvironmentVariables(t *testing.T) {
286286
t.Parallel()
287287

288+
tmpdir := t.TempDir()
289+
290+
// Defined by the coder script runner.
291+
scriptBinDir := filepath.Join(tmpdir, "coder-script", "bin")
292+
288293
manifest := agentsdk.Manifest{
289294
EnvironmentVariables: map[string]string{
290295
"MY_MANIFEST": "true",
@@ -293,7 +298,10 @@ func TestAgent_Session_EnvironmentVariables(t *testing.T) {
293298
},
294299
}
295300
banner := codersdk.ServiceBannerConfig{}
301+
fs := afero.NewMemMapFs()
296302
session := setupSSHSession(t, manifest, banner, nil, func(_ *agenttest.Client, opts *agent.Options) {
303+
opts.Filesystem = fs
304+
opts.ScriptDataDir = tmpdir
297305
opts.EnvironmentVariables["MY_OVERRIDE"] = "true"
298306
})
299307

@@ -335,6 +343,7 @@ func TestAgent_Session_EnvironmentVariables(t *testing.T) {
335343
"MY_OVERRIDE": "true", // From the agent environment variables option, overrides manifest.
336344
"MY_SESSION_MANIFEST": "false", // From the manifest, overrides session env.
337345
"MY_SESSION": "true", // From the session.
346+
"PATH": scriptBinDir + ":",
338347
} {
339348
t.Run(k, func(t *testing.T) {
340349
out := echoEnv(t, stdin, stdout, k)

‎agent/agentscripts/agentscripts.go

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ var (
4343

4444
// Options are a set of options for the runner.
4545
type Options struct {
46+
DataDir string
4647
LogDir string
4748
Logger slog.Logger
4849
SSHServer *agentssh.Server
@@ -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.DataDir, "coder-script"),
6264
scriptsExecuted: prometheus.NewCounterVec(prometheus.CounterOpts{
6365
Namespace: "agent",
6466
Subsystem: "scripts",
@@ -78,13 +80,20 @@ 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+
// ScriptBinDir returns the directory where scripts can store executable
92+
// binaries.
93+
func (r *Runner) ScriptBinDir() string {
94+
return filepath.Join(r.dataDir, "bin")
95+
}
96+
8897
func (r *Runner) RegisterMetrics(reg prometheus.Registerer) {
8998
if reg == nil {
9099
// If no registry, do nothing.
@@ -104,6 +113,11 @@ func (r *Runner) Init(scripts []codersdk.WorkspaceAgentScript) error {
104113
r.scripts = scripts
105114
r.Logger.Info(r.cronCtx, "initializing agent scripts", slog.F("script_count", len(scripts)), slog.F("log_dir", r.LogDir))
106115

116+
err := r.Filesystem.MkdirAll(r.ScriptBinDir(), 0o700)
117+
if err != nil {
118+
return xerrors.Errorf("create script bin dir: %w", err)
119+
}
120+
107121
for _, script := range scripts {
108122
if script.Cron == "" {
109123
continue
@@ -208,7 +222,18 @@ func (r *Runner) run(ctx context.Context, script codersdk.WorkspaceAgentScript)
208222
if !filepath.IsAbs(logPath) {
209223
logPath = filepath.Join(r.LogDir, logPath)
210224
}
211-
logger := r.Logger.With(slog.F("log_path", logPath))
225+
226+
scriptDataDir := filepath.Join(r.dataDir, script.LogSourceID.String())
227+
err := r.Filesystem.MkdirAll(scriptDataDir, 0o700)
228+
if err != nil {
229+
return xerrors.Errorf("%s script: create script temp dir: %w", scriptDataDir, err)
230+
}
231+
232+
logger := r.Logger.With(
233+
slog.F("log_source_id", script.LogSourceID),
234+
slog.F("log_path", logPath),
235+
slog.F("script_data_dir", scriptDataDir),
236+
)
212237
logger.Info(ctx, "running agent script", slog.F("script", script.Script))
213238

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

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

‎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_DIR",
349+
Value: clibase.StringOf(&scriptDataDir),
350+
},
342351
{
343352
Flag: "pprof-address",
344353
Default: "127.0.0.1:6060",

0 commit comments

Comments
 (0)