From 4f0a9f195a68b2c855eff1449cb2cc6dcd8fa382 Mon Sep 17 00:00:00 2001 From: Spike Curtis Date: Thu, 1 May 2025 13:46:38 +0400 Subject: [PATCH] fix: fix data race in agentscripts.Runner --- agent/agentscripts/agentscripts.go | 23 +++++++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/agent/agentscripts/agentscripts.go b/agent/agentscripts/agentscripts.go index 4e4921b87ee5b..79606a80233b9 100644 --- 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) ||