fix: fix data race in agentscripts.Runner

spikecurtis · spikecurtis · commit 4f0a9f195a68 · 2025-05-01T13:46:38.000+04:00
diff --git a/agent/agentscripts/agentscripts.go b/agent/agentscripts/agentscripts.go
@@ -10,7 +10,6 @@ import (
 	"os/user"
 	"path/filepath"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"github.com/google/uuid"
@@ -104,7 +103,6 @@ type Runner struct {
 	closed          chan struct{}
 	closeMutex      sync.Mutex
 	cron            *cron.Cron
-	initialized     atomic.Bool
 	scripts         []runnerScript
 	dataDir         string
 	scriptCompleted ScriptCompletedFunc
@@ -113,6 +111,9 @@ type Runner struct {
 	// execute startup scripts, and scripts on a cron schedule. Both will increment
 	// this counter.
 	scriptsExecuted *prometheus.CounterVec
+
+	initMutex   sync.Mutex
+	initialized bool
 }
 
 // DataDir returns the directory where scripts data is stored.
@@ -154,10 +155,12 @@ func WithPostStartScripts(scripts ...codersdk.WorkspaceAgentScript) InitOption {
 // It also schedules any scripts that have a schedule.
 // This function must be called before Execute.
 func (r *Runner) Init(scripts []codersdk.WorkspaceAgentScript, scriptCompleted ScriptCompletedFunc, opts ...InitOption) error {
-	if r.initialized.Load() {
+	r.initMutex.Lock()
+	defer r.initMutex.Unlock()
+	if r.initialized {
 		return xerrors.New("init: already initialized")
 	}
-	r.initialized.Store(true)
+	r.initialized = true
 	r.scripts = toRunnerScript(scripts...)
 	r.scriptCompleted = scriptCompleted
 	for _, opt := range opts {
@@ -227,6 +230,18 @@ const (
 
 // Execute runs a set of scripts according to a filter.
 func (r *Runner) Execute(ctx context.Context, option ExecuteOption) error {
+	initErr := func() error {
+		r.initMutex.Lock()
+		defer r.initMutex.Unlock()
+		if !r.initialized {
+			return xerrors.New("execute: not initialized")
+		}
+		return nil
+	}()
+	if initErr != nil {
+		return initErr
+	}
+
 	var eg errgroup.Group
 	for _, script := range r.scripts {
 		runScript := (option == ExecuteStartScripts && script.RunOnStart) ||
