cleanup datadir

coder · mafredri · Feb 20, 2024 · Feb 19, 2024 · Feb 19, 2024 · Feb 19, 2024
commit e4053c82887bb3e9b10291ceecdb5ad16f952faa
diff --git 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
@@ -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
@@ -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
@@ -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,
 	})
 }