From ebab733c930832c301dd0b076e53db353ee679e2 Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Mon, 19 Feb 2024 13:35:02 +0200 Subject: [PATCH 01/11] 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 --- agent/agent.go | 17 +++++++++++++++ agent/agent_test.go | 9 ++++++++ agent/agentscripts/agentscripts.go | 34 +++++++++++++++++++++++++++++- cli/agent.go | 9 ++++++++ 4 files changed, 68 insertions(+), 1 deletion(-) diff --git a/agent/agent.go b/agent/agent.go index a369432c0390f..fc1dad101bef2 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -66,6 +66,7 @@ type Options struct { Filesystem afero.Fs LogDir string TempDir string + ScriptDataDir string ExchangeToken func(ctx context.Context) (string, error) Client Client ReconnectingPTYTimeout time.Duration @@ -115,6 +116,12 @@ func New(options Options) Agent { } options.LogDir = options.TempDir } + if options.ScriptDataDir == "" { + if options.TempDir != os.TempDir() { + options.Logger.Debug(context.Background(), "script data dir not set, using temp dir", slog.F("temp_dir", options.TempDir)) + } + options.ScriptDataDir = options.TempDir + } if options.ExchangeToken == nil { options.ExchangeToken = func(ctx context.Context) (string, error) { return "", nil @@ -152,6 +159,7 @@ func New(options Options) Agent { filesystem: options.Filesystem, logDir: options.LogDir, tempDir: options.TempDir, + scriptDataDir: options.ScriptDataDir, lifecycleUpdate: make(chan struct{}, 1), lifecycleReported: make(chan codersdk.WorkspaceAgentLifecycle, 1), lifecycleStates: []agentsdk.PostLifecycleRequest{{State: codersdk.WorkspaceAgentLifecycleCreated}}, @@ -183,6 +191,7 @@ type agent struct { filesystem afero.Fs logDir string tempDir string + scriptDataDir string // ignorePorts tells the api handler which ports to ignore when // listing all listening ports. This is helpful to hide ports that // are used by the agent, that the user does not care about. @@ -250,6 +259,7 @@ func (a *agent) init(ctx context.Context) { a.sshServer = sshSrv a.scriptRunner = agentscripts.New(agentscripts.Options{ LogDir: a.logDir, + DataDir: a.scriptDataDir, Logger: a.logger, SSHServer: sshSrv, Filesystem: a.filesystem, @@ -954,6 +964,13 @@ func (a *agent) updateCommandEnv(current []string) (updated []string, err error) envs[k] = v } + // Prepend the agent script bin directory to the PATH + // (this is where Coder modules place their binaries). + if _, ok := envs["PATH"]; !ok { + envs["PATH"] = os.Getenv("PATH") + } + envs["PATH"] = a.scriptRunner.ScriptBinDir() + ":" + envs["PATH"] + for k, v := range envs { updated = append(updated, fmt.Sprintf("%s=%s", k, v)) } diff --git a/agent/agent_test.go b/agent/agent_test.go index b894beeca905f..dd830ee16093e 100644 --- a/agent/agent_test.go +++ b/agent/agent_test.go @@ -286,6 +286,11 @@ func TestAgent_SessionExec(t *testing.T) { func TestAgent_Session_EnvironmentVariables(t *testing.T) { t.Parallel() + tmpdir := t.TempDir() + + // Defined by the coder script runner. + scriptBinDir := filepath.Join(tmpdir, "coder-script-data", "bin") + manifest := agentsdk.Manifest{ EnvironmentVariables: map[string]string{ "MY_MANIFEST": "true", @@ -294,7 +299,10 @@ func TestAgent_Session_EnvironmentVariables(t *testing.T) { }, } banner := codersdk.ServiceBannerConfig{} + fs := afero.NewMemMapFs() session := setupSSHSession(t, manifest, banner, nil, func(_ *agenttest.Client, opts *agent.Options) { + opts.Filesystem = fs + opts.ScriptDataDir = tmpdir opts.EnvironmentVariables["MY_OVERRIDE"] = "true" }) @@ -341,6 +349,7 @@ func TestAgent_Session_EnvironmentVariables(t *testing.T) { "MY_OVERRIDE": "true", // From the agent environment variables option, overrides manifest. "MY_SESSION_MANIFEST": "false", // From the manifest, overrides session env. "MY_SESSION": "true", // From the session. + "PATH": scriptBinDir + ":", } { t.Run(k, func(t *testing.T) { echoEnv(t, stdin, k) diff --git a/agent/agentscripts/agentscripts.go b/agent/agentscripts/agentscripts.go index f6052d605432e..d3990d228f45a 100644 --- a/agent/agentscripts/agentscripts.go +++ b/agent/agentscripts/agentscripts.go @@ -43,6 +43,7 @@ var ( // Options are a set of options for the runner. type Options struct { + DataDir string LogDir string Logger slog.Logger SSHServer *agentssh.Server @@ -59,6 +60,7 @@ func New(opts Options) *Runner { cronCtxCancel: cronCtxCancel, cron: cron.New(cron.WithParser(parser)), closed: make(chan struct{}), + dataDir: filepath.Join(opts.DataDir, "coder-script-data"), scriptsExecuted: prometheus.NewCounterVec(prometheus.CounterOpts{ Namespace: "agent", Subsystem: "scripts", @@ -78,6 +80,7 @@ type Runner struct { cron *cron.Cron initialized atomic.Bool scripts []codersdk.WorkspaceAgentScript + dataDir string // scriptsExecuted includes all scripts executed by the workspace agent. Agents // execute startup scripts, and scripts on a cron schedule. Both will increment @@ -85,6 +88,12 @@ type Runner struct { scriptsExecuted *prometheus.CounterVec } +// ScriptBinDir returns the directory where scripts can store executable +// binaries. +func (r *Runner) ScriptBinDir() string { + return filepath.Join(r.dataDir, "bin") +} + func (r *Runner) RegisterMetrics(reg prometheus.Registerer) { if reg == nil { // If no registry, do nothing. @@ -104,6 +113,11 @@ func (r *Runner) Init(scripts []codersdk.WorkspaceAgentScript) error { r.scripts = scripts r.Logger.Info(r.cronCtx, "initializing agent scripts", slog.F("script_count", len(scripts)), slog.F("log_dir", r.LogDir)) + err := r.Filesystem.MkdirAll(r.ScriptBinDir(), 0o700) + if err != nil { + return xerrors.Errorf("create script bin dir: %w", err) + } + for _, script := range scripts { if script.Cron == "" { continue @@ -208,7 +222,18 @@ func (r *Runner) run(ctx context.Context, script codersdk.WorkspaceAgentScript) if !filepath.IsAbs(logPath) { logPath = filepath.Join(r.LogDir, logPath) } - logger := r.Logger.With(slog.F("log_path", logPath)) + + scriptDataDir := filepath.Join(r.dataDir, script.LogSourceID.String()) + err := r.Filesystem.MkdirAll(scriptDataDir, 0o700) + if err != nil { + return xerrors.Errorf("%s script: create script temp dir: %w", scriptDataDir, err) + } + + logger := r.Logger.With( + slog.F("log_source_id", script.LogSourceID), + slog.F("log_path", logPath), + slog.F("script_data_dir", scriptDataDir), + ) logger.Info(ctx, "running agent script", slog.F("script", script.Script)) 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) cmd.WaitDelay = 10 * time.Second cmd.Cancel = cmdCancel(cmd) + // Expose env vars that can be used in the script for storing data + // and binaries. In the future, we may want to expose more env vars + // for the script to use, like CODER_SCRIPT_DATA_DIR for persistent + // storage. + cmd.Env = append(cmd.Env, "CODER_SCRIPT_DATA_DIR="+scriptDataDir) + cmd.Env = append(cmd.Env, "CODER_SCRIPT_BIN_DIR="+r.ScriptBinDir()) + send, flushAndClose := agentsdk.LogsSender(script.LogSourceID, r.PatchLogs, logger) // If ctx is canceled here (or in a writer below), we may be // discarding logs, but that's okay because we're shutting down diff --git a/cli/agent.go b/cli/agent.go index c951ec7509901..b2af9f6b5172b 100644 --- a/cli/agent.go +++ b/cli/agent.go @@ -40,6 +40,7 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd { var ( auth string logDir string + scriptDataDir string pprofAddress string noReap bool sshMaxTimeout time.Duration @@ -289,6 +290,7 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd { Client: client, Logger: logger, LogDir: logDir, + ScriptDataDir: scriptDataDir, TailnetListenPort: uint16(tailnetListenPort), ExchangeToken: func(ctx context.Context) (string, error) { if exchangeToken == nil { @@ -339,6 +341,13 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd { Env: "CODER_AGENT_LOG_DIR", Value: clibase.StringOf(&logDir), }, + { + Flag: "script-data-dir", + Default: os.TempDir(), + Description: "Specify the location for storing script data.", + Env: "CODER_AGENT_SCRIPT_DIR", + Value: clibase.StringOf(&scriptDataDir), + }, { Flag: "pprof-address", Default: "127.0.0.1:6060", From 61007e8bf9cfc0eabd43cbdb41b5671e29891389 Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Mon, 19 Feb 2024 13:50:01 +0200 Subject: [PATCH 02/11] add agentscript test --- agent/agentscripts/agentscripts_test.go | 33 ++++++++++++++++++++++--- 1 file changed, 30 insertions(+), 3 deletions(-) diff --git a/agent/agentscripts/agentscripts_test.go b/agent/agentscripts/agentscripts_test.go index bb3f842a45962..de70c043a0701 100644 --- a/agent/agentscripts/agentscripts_test.go +++ b/agent/agentscripts/agentscripts_test.go @@ -2,9 +2,11 @@ package agentscripts_test import ( "context" + "path/filepath" "testing" "time" + "github.com/google/uuid" "github.com/prometheus/client_golang/prometheus" "github.com/spf13/afero" "github.com/stretchr/testify/require" @@ -30,7 +32,8 @@ func TestExecuteBasic(t *testing.T) { }) defer runner.Close() err := runner.Init([]codersdk.WorkspaceAgentScript{{ - Script: "echo hello", + LogSourceID: uuid.New(), + Script: "echo hello", }}) require.NoError(t, err) require.NoError(t, runner.Execute(context.Background(), func(script codersdk.WorkspaceAgentScript) bool { @@ -40,13 +43,36 @@ func TestExecuteBasic(t *testing.T) { require.Equal(t, "hello", log.Logs[0].Output) } +func TestEnv(t *testing.T) { + t.Parallel() + logs := make(chan agentsdk.PatchLogs, 1) + runner := setup(t, func(ctx context.Context, req agentsdk.PatchLogs) error { + logs <- req + return nil + }) + defer runner.Close() + id := uuid.New() + err := runner.Init([]codersdk.WorkspaceAgentScript{{ + LogSourceID: id, + Script: "echo $CODER_SCRIPT_DATA_DIR\necho $CODER_SCRIPT_BIN_DIR\n", + }}) + require.NoError(t, err) + require.NoError(t, runner.Execute(context.Background(), func(script codersdk.WorkspaceAgentScript) bool { + return true + })) + log := <-logs + require.Contains(t, log.Logs[0].Output, filepath.Join(runner.DataDir, "coder-script-data", id.String())) + require.Contains(t, log.Logs[1].Output, runner.ScriptBinDir()) +} + func TestTimeout(t *testing.T) { t.Parallel() runner := setup(t, nil) defer runner.Close() err := runner.Init([]codersdk.WorkspaceAgentScript{{ - Script: "sleep infinity", - Timeout: time.Millisecond, + LogSourceID: uuid.New(), + Script: "sleep infinity", + Timeout: time.Millisecond, }}) require.NoError(t, err) require.ErrorIs(t, runner.Execute(context.Background(), nil), agentscripts.ErrTimeout) @@ -78,6 +104,7 @@ func setup(t *testing.T, patchLogs func(ctx context.Context, req agentsdk.PatchL }) return agentscripts.New(agentscripts.Options{ LogDir: t.TempDir(), + DataDir: t.TempDir(), Logger: logger, SSHServer: s, Filesystem: fs, From 7f438cc9c54f50efdaed4987296e8f8f27ff6463 Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Mon, 19 Feb 2024 14:09:00 +0200 Subject: [PATCH 03/11] handle window paths --- agent/agent.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/agent/agent.go b/agent/agent.go index fc1dad101bef2..ec0e23cdf49e3 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -969,7 +969,7 @@ func (a *agent) updateCommandEnv(current []string) (updated []string, err error) if _, ok := envs["PATH"]; !ok { envs["PATH"] = os.Getenv("PATH") } - envs["PATH"] = a.scriptRunner.ScriptBinDir() + ":" + envs["PATH"] + envs["PATH"] = fmt.Sprintf("%s%c%s", a.scriptRunner.ScriptBinDir(), filepath.ListSeparator, envs["PATH"]) for k, v := range envs { updated = append(updated, fmt.Sprintf("%s=%s", k, v)) From c30a1c31d157c303f1e46b6238cbb501e73462d6 Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Mon, 19 Feb 2024 14:30:20 +0200 Subject: [PATCH 04/11] fix path test --- agent/agent_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/agent/agent_test.go b/agent/agent_test.go index dd830ee16093e..031a59fc4919b 100644 --- a/agent/agent_test.go +++ b/agent/agent_test.go @@ -349,7 +349,7 @@ func TestAgent_Session_EnvironmentVariables(t *testing.T) { "MY_OVERRIDE": "true", // From the agent environment variables option, overrides manifest. "MY_SESSION_MANIFEST": "false", // From the manifest, overrides session env. "MY_SESSION": "true", // From the session. - "PATH": scriptBinDir + ":", + "PATH": scriptBinDir + string(filepath.ListSeparator), } { t.Run(k, func(t *testing.T) { echoEnv(t, stdin, k) From e4053c82887bb3e9b10291ceecdb5ad16f952faa Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Mon, 19 Feb 2024 14:35:50 +0200 Subject: [PATCH 05/11] cleanup datadir --- agent/agent.go | 12 ++++++------ agent/agent_test.go | 3 ++- agent/agentscripts/agentscripts.go | 21 +++++++++++++-------- agent/agentscripts/agentscripts_test.go | 24 +++++++++++++++--------- 4 files changed, 36 insertions(+), 24 deletions(-) diff --git a/agent/agent.go b/agent/agent.go index ec0e23cdf49e3..a833d677e2999 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -258,12 +258,12 @@ func (a *agent) init(ctx context.Context) { } a.sshServer = sshSrv a.scriptRunner = agentscripts.New(agentscripts.Options{ - LogDir: a.logDir, - DataDir: a.scriptDataDir, - Logger: a.logger, - SSHServer: sshSrv, - Filesystem: a.filesystem, - PatchLogs: a.client.PatchLogs, + LogDir: a.logDir, + DataDirBase: a.scriptDataDir, + Logger: a.logger, + SSHServer: sshSrv, + Filesystem: a.filesystem, + PatchLogs: a.client.PatchLogs, }) // Register runner metrics. If the prom registry is nil, the metrics // will not report anywhere. diff --git a/agent/agent_test.go b/agent/agent_test.go index 031a59fc4919b..e72543e5e733e 100644 --- a/agent/agent_test.go +++ b/agent/agent_test.go @@ -288,7 +288,8 @@ func TestAgent_Session_EnvironmentVariables(t *testing.T) { tmpdir := t.TempDir() - // Defined by the coder script runner. + // Defined by the coder script runner, hardcoded here since we don't + // have a reference to it. scriptBinDir := filepath.Join(tmpdir, "coder-script-data", "bin") manifest := agentsdk.Manifest{ diff --git a/agent/agentscripts/agentscripts.go b/agent/agentscripts/agentscripts.go index d3990d228f45a..e7169f9fdb699 100644 --- a/agent/agentscripts/agentscripts.go +++ b/agent/agentscripts/agentscripts.go @@ -43,12 +43,12 @@ var ( // Options are a set of options for the runner. type Options struct { - DataDir string - LogDir string - Logger slog.Logger - SSHServer *agentssh.Server - Filesystem afero.Fs - PatchLogs func(ctx context.Context, req agentsdk.PatchLogs) error + DataDirBase string + LogDir string + Logger slog.Logger + SSHServer *agentssh.Server + Filesystem afero.Fs + PatchLogs func(ctx context.Context, req agentsdk.PatchLogs) error } // New creates a runner for the provided scripts. @@ -60,7 +60,7 @@ func New(opts Options) *Runner { cronCtxCancel: cronCtxCancel, cron: cron.New(cron.WithParser(parser)), closed: make(chan struct{}), - dataDir: filepath.Join(opts.DataDir, "coder-script-data"), + dataDir: filepath.Join(opts.DataDirBase, "coder-script-data"), scriptsExecuted: prometheus.NewCounterVec(prometheus.CounterOpts{ Namespace: "agent", Subsystem: "scripts", @@ -88,6 +88,11 @@ type Runner struct { scriptsExecuted *prometheus.CounterVec } +// DataDir returns the directory where scripts data is stored. +func (r *Runner) DataDir() string { + return r.dataDir +} + // ScriptBinDir returns the directory where scripts can store executable // binaries. func (r *Runner) ScriptBinDir() string { @@ -223,7 +228,7 @@ func (r *Runner) run(ctx context.Context, script codersdk.WorkspaceAgentScript) logPath = filepath.Join(r.LogDir, logPath) } - scriptDataDir := filepath.Join(r.dataDir, script.LogSourceID.String()) + scriptDataDir := filepath.Join(r.DataDir(), script.LogSourceID.String()) err := r.Filesystem.MkdirAll(scriptDataDir, 0o700) if err != nil { return xerrors.Errorf("%s script: create script temp dir: %w", scriptDataDir, err) diff --git a/agent/agentscripts/agentscripts_test.go b/agent/agentscripts/agentscripts_test.go index de70c043a0701..e51a6c2d2dd4a 100644 --- a/agent/agentscripts/agentscripts_test.go +++ b/agent/agentscripts/agentscripts_test.go @@ -27,7 +27,10 @@ func TestExecuteBasic(t *testing.T) { t.Parallel() logs := make(chan agentsdk.PatchLogs, 1) runner := setup(t, func(ctx context.Context, req agentsdk.PatchLogs) error { - logs <- req + select { + case <-ctx.Done(): + case logs <- req: + } return nil }) defer runner.Close() @@ -47,7 +50,10 @@ func TestEnv(t *testing.T) { t.Parallel() logs := make(chan agentsdk.PatchLogs, 1) runner := setup(t, func(ctx context.Context, req agentsdk.PatchLogs) error { - logs <- req + select { + case <-ctx.Done(): + case logs <- req: + } return nil }) defer runner.Close() @@ -61,7 +67,7 @@ func TestEnv(t *testing.T) { return true })) log := <-logs - require.Contains(t, log.Logs[0].Output, filepath.Join(runner.DataDir, "coder-script-data", id.String())) + require.Contains(t, log.Logs[0].Output, filepath.Join(runner.DataDir(), id.String())) require.Contains(t, log.Logs[1].Output, runner.ScriptBinDir()) } @@ -103,11 +109,11 @@ func setup(t *testing.T, patchLogs func(ctx context.Context, req agentsdk.PatchL _ = s.Close() }) return agentscripts.New(agentscripts.Options{ - LogDir: t.TempDir(), - DataDir: t.TempDir(), - Logger: logger, - SSHServer: s, - Filesystem: fs, - PatchLogs: patchLogs, + LogDir: t.TempDir(), + DataDirBase: t.TempDir(), + Logger: logger, + SSHServer: s, + Filesystem: fs, + PatchLogs: patchLogs, }) } From 35968dde8bbcd598b6ecd4081c7f5dd4479f278a Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Mon, 19 Feb 2024 14:38:10 +0200 Subject: [PATCH 06/11] always log --- agent/agent.go | 2 ++ 1 file changed, 2 insertions(+) diff --git a/agent/agent.go b/agent/agent.go index a833d677e2999..ffab3b17008fc 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -119,6 +119,8 @@ func New(options Options) Agent { if options.ScriptDataDir == "" { if options.TempDir != os.TempDir() { options.Logger.Debug(context.Background(), "script data dir not set, using temp dir", slog.F("temp_dir", options.TempDir)) + } else { + options.Logger.Debug(context.Background(), "using script data dir", slog.F("script_data_dir", options.ScriptDataDir)) } options.ScriptDataDir = options.TempDir } From 0f6f2547562af96f1a2d2f4227440680d2e1548b Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Mon, 19 Feb 2024 14:38:48 +0200 Subject: [PATCH 07/11] add log dir log for good measure --- agent/agent.go | 2 ++ 1 file changed, 2 insertions(+) diff --git a/agent/agent.go b/agent/agent.go index ffab3b17008fc..c0a61fa97fe98 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -113,6 +113,8 @@ func New(options Options) Agent { if options.LogDir == "" { if options.TempDir != os.TempDir() { options.Logger.Debug(context.Background(), "log dir not set, using temp dir", slog.F("temp_dir", options.TempDir)) + } else { + options.Logger.Debug(context.Background(), "using log dir", slog.F("log_dir", options.LogDir)) } options.LogDir = options.TempDir } From fddefbdca07ee04551ff234a3c413d0d6f0283a7 Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Mon, 19 Feb 2024 16:02:29 +0000 Subject: [PATCH 08/11] update golden files --- cli/testdata/coder_agent_--help.golden | 3 +++ 1 file changed, 3 insertions(+) diff --git a/cli/testdata/coder_agent_--help.golden b/cli/testdata/coder_agent_--help.golden index 08dab47a21e14..47f14d4b4d93f 100644 --- a/cli/testdata/coder_agent_--help.golden +++ b/cli/testdata/coder_agent_--help.golden @@ -33,6 +33,9 @@ OPTIONS: --prometheus-address string, $CODER_AGENT_PROMETHEUS_ADDRESS (default: 127.0.0.1:2112) The bind address to serve Prometheus metrics. + --script-data-dir string, $CODER_AGENT_SCRIPT_DIR (default: /tmp) + Specify the location for storing script data. + --ssh-max-timeout duration, $CODER_AGENT_SSH_MAX_TIMEOUT (default: 72h) Specify the max timeout for a SSH connection, it is advisable to set it to a minimum of 60s, but no more than 72h. From 58ee327f7284b48f64fd0ee5dd713f8e85699813 Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Mon, 19 Feb 2024 18:51:36 +0200 Subject: [PATCH 09/11] try to fix test on Win --- agent/agentscripts/agentscripts_test.go | 38 +++++++++++++++++++++---- 1 file changed, 32 insertions(+), 6 deletions(-) diff --git a/agent/agentscripts/agentscripts_test.go b/agent/agentscripts/agentscripts_test.go index e51a6c2d2dd4a..277d1216a42a1 100644 --- a/agent/agentscripts/agentscripts_test.go +++ b/agent/agentscripts/agentscripts_test.go @@ -3,6 +3,7 @@ package agentscripts_test import ( "context" "path/filepath" + "runtime" "testing" "time" @@ -17,6 +18,7 @@ import ( "github.com/coder/coder/v2/agent/agentssh" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/agentsdk" + "github.com/coder/coder/v2/testutil" ) func TestMain(m *testing.M) { @@ -48,7 +50,7 @@ func TestExecuteBasic(t *testing.T) { func TestEnv(t *testing.T) { t.Parallel() - logs := make(chan agentsdk.PatchLogs, 1) + logs := make(chan agentsdk.PatchLogs, 2) runner := setup(t, func(ctx context.Context, req agentsdk.PatchLogs) error { select { case <-ctx.Done(): @@ -58,17 +60,41 @@ func TestEnv(t *testing.T) { }) defer runner.Close() id := uuid.New() + script := "echo $CODER_SCRIPT_DATA_DIR\necho $CODER_SCRIPT_BIN_DIR\n" + if runtime.GOOS == "windows" { + script = ` + cmd.exe /c echo %CODER_SCRIPT_DATA_DIR% + cmd.exe /c echo %CODER_SCRIPT_BIN_DIR% + ` + } err := runner.Init([]codersdk.WorkspaceAgentScript{{ LogSourceID: id, - Script: "echo $CODER_SCRIPT_DATA_DIR\necho $CODER_SCRIPT_BIN_DIR\n", + Script: script, }}) require.NoError(t, err) - require.NoError(t, runner.Execute(context.Background(), func(script codersdk.WorkspaceAgentScript) bool { + + ctx := testutil.Context(t, testutil.WaitLong) + + require.NoError(t, runner.Execute(ctx, func(script codersdk.WorkspaceAgentScript) bool { return true })) - log := <-logs - require.Contains(t, log.Logs[0].Output, filepath.Join(runner.DataDir(), id.String())) - require.Contains(t, log.Logs[1].Output, runner.ScriptBinDir()) + var log []agentsdk.Log + for { + select { + case <-ctx.Done(): + require.Fail(t, "timed out waiting for logs") + case l := <-logs: + for _, l := range l.Logs { + t.Logf("log: %s", l.Output) + } + log = append(log, l.Logs...) + } + if len(log) >= 2 { + break + } + } + require.Contains(t, log[0].Output, filepath.Join(runner.DataDir(), id.String())) + require.Contains(t, log[1].Output, runner.ScriptBinDir()) } func TestTimeout(t *testing.T) { From 7dbfc2b997102a6813f4e3f834a43af44646f7df Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Mon, 19 Feb 2024 20:38:52 +0200 Subject: [PATCH 10/11] Apply suggestions from code review --- agent/agent_test.go | 2 -- cli/agent.go | 2 +- cli/testdata/coder_agent_--help.golden | 2 +- 3 files changed, 2 insertions(+), 4 deletions(-) diff --git a/agent/agent_test.go b/agent/agent_test.go index e72543e5e733e..573440769806c 100644 --- a/agent/agent_test.go +++ b/agent/agent_test.go @@ -300,9 +300,7 @@ func TestAgent_Session_EnvironmentVariables(t *testing.T) { }, } banner := codersdk.ServiceBannerConfig{} - fs := afero.NewMemMapFs() session := setupSSHSession(t, manifest, banner, nil, func(_ *agenttest.Client, opts *agent.Options) { - opts.Filesystem = fs opts.ScriptDataDir = tmpdir opts.EnvironmentVariables["MY_OVERRIDE"] = "true" }) diff --git a/cli/agent.go b/cli/agent.go index b2af9f6b5172b..23473022abea7 100644 --- a/cli/agent.go +++ b/cli/agent.go @@ -345,7 +345,7 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd { Flag: "script-data-dir", Default: os.TempDir(), Description: "Specify the location for storing script data.", - Env: "CODER_AGENT_SCRIPT_DIR", + Env: "CODER_AGENT_SCRIPT_DATA_DIR", Value: clibase.StringOf(&scriptDataDir), }, { diff --git a/cli/testdata/coder_agent_--help.golden b/cli/testdata/coder_agent_--help.golden index 47f14d4b4d93f..372395c4ba5fe 100644 --- a/cli/testdata/coder_agent_--help.golden +++ b/cli/testdata/coder_agent_--help.golden @@ -33,7 +33,7 @@ OPTIONS: --prometheus-address string, $CODER_AGENT_PROMETHEUS_ADDRESS (default: 127.0.0.1:2112) The bind address to serve Prometheus metrics. - --script-data-dir string, $CODER_AGENT_SCRIPT_DIR (default: /tmp) + --script-data-dir string, $CODER_AGENT_SCRIPT_DATA_DIR (default: /tmp) Specify the location for storing script data. --ssh-max-timeout duration, $CODER_AGENT_SSH_MAX_TIMEOUT (default: 72h) From 32341430c991e3b3e8a2b5d76dbd1f28a46a06a6 Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Tue, 20 Feb 2024 13:15:14 +0200 Subject: [PATCH 11/11] execute async to ensure log processing --- agent/agentscripts/agentscripts_test.go | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/agent/agentscripts/agentscripts_test.go b/agent/agentscripts/agentscripts_test.go index 277d1216a42a1..d7fce25fda1fa 100644 --- a/agent/agentscripts/agentscripts_test.go +++ b/agent/agentscripts/agentscripts_test.go @@ -10,6 +10,7 @@ import ( "github.com/google/uuid" "github.com/prometheus/client_golang/prometheus" "github.com/spf13/afero" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/goleak" @@ -75,9 +76,13 @@ func TestEnv(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) - require.NoError(t, runner.Execute(ctx, func(script codersdk.WorkspaceAgentScript) bool { - return true - })) + testutil.Go(t, func() { + err := runner.Execute(ctx, func(script codersdk.WorkspaceAgentScript) bool { + return true + }) + assert.NoError(t, err) + }) + var log []agentsdk.Log for { select {