Skip to content

Commit 7f8b6f9

Browse files
DanielleMaywoodDanielleMaywood
committed
feat: capture stage and if script timed out
1 parent 61c26ea commit 7f8b6f9

15 files changed

+452
-299
lines changed

‎agent/agent.go

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -949,9 +949,7 @@ func (a *agent) handleManifest(manifestOK *checkpoint) func(ctx context.Context,
949949
start := time.Now()
950950
// here we use the graceful context because the script runner is not directly tied
951951
// to the agent API.
952-
err := a.scriptRunner.Execute(a.gracefulCtx, func(script codersdk.WorkspaceAgentScript) bool {
953-
return script.RunOnStart
954-
})
952+
err := a.scriptRunner.Execute(a.gracefulCtx, agentscripts.ExecuteStartScripts)
955953
// Measure the time immediately after the script has finished
956954
dur := time.Since(start).Seconds()
957955
if err != nil {
@@ -1844,9 +1842,7 @@ func (a *agent) Close() error {
18441842
a.gracefulCancel()
18451843

18461844
lifecycleState := codersdk.WorkspaceAgentLifecycleOff
1847-
err = a.scriptRunner.Execute(a.hardCtx, func(script codersdk.WorkspaceAgentScript) bool {
1848-
return script.RunOnStop
1849-
})
1845+
err = a.scriptRunner.Execute(a.hardCtx, agentscripts.ExecuteStopScripts)
18501846
if err != nil {
18511847
a.logger.Warn(a.hardCtx, "shutdown script(s) failed", slog.Error(err))
18521848
if errors.Is(err, agentscripts.ErrTimeout) {

‎agent/agentscripts/agentscripts.go

Lines changed: 39 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -141,7 +141,7 @@ func (r *Runner) Init(scripts []codersdk.WorkspaceAgentScript, scriptCompleted S
141141
}
142142
script := script
143143
_, err := r.cron.AddFunc(script.Cron, func() {
144-
err := r.trackRun(r.cronCtx, script)
144+
err := r.trackRun(r.cronCtx, script, ExecuteCronScripts)
145145
if err != nil {
146146
r.Logger.Warn(context.Background(), "run agent script on schedule", slog.Error(err))
147147
}
@@ -178,22 +178,31 @@ func (r *Runner) StartCron() {
178178
}
179179
}
180180

181+
type ExecuteOption int
182+
183+
const (
184+
ExecuteAllScripts ExecuteOption = iota
185+
ExecuteStartScripts
186+
ExecuteStopScripts
187+
ExecuteCronScripts
188+
)
189+
181190
// Execute runs a set of scripts according to a filter.
182-
func (r *Runner) Execute(ctx context.Context, filter func(script codersdk.WorkspaceAgentScript) bool) error {
183-
if filter == nil {
184-
// Execute em' all!
185-
filter = func(script codersdk.WorkspaceAgentScript) bool {
186-
return true
187-
}
188-
}
191+
func (r *Runner) Execute(ctx context.Context, option ExecuteOption) error {
189192
var eg errgroup.Group
190193
for _, script := range r.scripts {
191-
if !filter(script) {
194+
runScript := (option == ExecuteStartScripts && script.RunOnStart) ||
195+
(option == ExecuteStopScripts && script.RunOnStop) ||
196+
(option == ExecuteCronScripts && script.Cron != "") ||
197+
option == ExecuteAllScripts
198+
199+
if !runScript {
192200
continue
193201
}
202+
194203
script := script
195204
eg.Go(func() error {
196-
err := r.trackRun(ctx, script)
205+
err := r.trackRun(ctx, script, option)
197206
if err != nil {
198207
return xerrors.Errorf("run agent script %q: %w", script.LogSourceID, err)
199208
}
@@ -204,8 +213,8 @@ func (r *Runner) Execute(ctx context.Context, filter func(script codersdk.Worksp
204213
}
205214

206215
// trackRun wraps "run" with metrics.
207-
func (r *Runner) trackRun(ctx context.Context, script codersdk.WorkspaceAgentScript) error {
208-
err := r.run(ctx, script)
216+
func (r *Runner) trackRun(ctx context.Context, script codersdk.WorkspaceAgentScript, option ExecuteOption) error {
217+
err := r.run(ctx, script, option)
209218
if err != nil {
210219
r.scriptsExecuted.WithLabelValues("false").Add(1)
211220
} else {
@@ -218,7 +227,7 @@ func (r *Runner) trackRun(ctx context.Context, script codersdk.WorkspaceAgentScr
218227
// If the timeout is exceeded, the process is sent an interrupt signal.
219228
// If the process does not exit after a few seconds, it is forcefully killed.
220229
// This function immediately returns after a timeout, and does not wait for the process to exit.
221-
func (r *Runner) run(ctx context.Context, script codersdk.WorkspaceAgentScript) error {
230+
func (r *Runner) run(ctx context.Context, script codersdk.WorkspaceAgentScript, option ExecuteOption) error {
222231
logPath := script.LogPath
223232
if logPath == "" {
224233
logPath = fmt.Sprintf("coder-script-%s.log", script.LogSourceID)
@@ -321,15 +330,25 @@ func (r *Runner) run(ctx context.Context, script codersdk.WorkspaceAgentScript)
321330
logger.Info(ctx, fmt.Sprintf("%s script completed", logPath), slog.F("execution_time", execTime), slog.F("exit_code", exitCode))
322331
}
323332

333+
var stage proto.Timing_Stage
334+
switch option {
335+
case ExecuteStartScripts:
336+
stage = proto.Timing_START
337+
case ExecuteStopScripts:
338+
stage = proto.Timing_STOP
339+
case ExecuteCronScripts:
340+
stage = proto.Timing_CRON
341+
}
342+
324343
_, err = r.scriptCompleted(ctx, &proto.WorkspaceAgentScriptCompletedRequest{
325344
Timing: &proto.Timing{
326-
ScriptId: script.ID[:],
327-
DisplayName: script.DisplayName,
328-
Start: timestamppb.New(start),
329-
End: timestamppb.New(end),
330-
ExitCode: int32(exitCode),
331-
RanOnStart: script.RunOnStart,
332-
BlockedLogin: script.StartBlocksLogin,
345+
ScriptId: script.ID[:],
346+
DisplayName: script.DisplayName,
347+
Start: timestamppb.New(start),
348+
End: timestamppb.New(end),
349+
ExitCode: int32(exitCode),
350+
Stage: stage,
351+
TimedOut: errors.Is(err, ErrTimeout),
333352
},
334353
})
335354

‎agent/agentscripts/agentscripts_test.go

Lines changed: 4 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -41,9 +41,7 @@ func TestExecuteBasic(t *testing.T) {
4141
Script: "echo hello",
4242
}}, aAPI.ScriptCompleted)
4343
require.NoError(t, err)
44-
require.NoError(t, runner.Execute(context.Background(), func(script codersdk.WorkspaceAgentScript) bool {
45-
return true
46-
}))
44+
require.NoError(t, runner.Execute(context.Background(), agentscripts.ExecuteAllScripts))
4745
log := testutil.RequireRecvCtx(ctx, t, fLogger.logs)
4846
require.Equal(t, "hello", log.Output)
4947
}
@@ -73,9 +71,7 @@ func TestEnv(t *testing.T) {
7371
ctx := testutil.Context(t, testutil.WaitLong)
7472

7573
done := testutil.Go(t, func() {
76-
err := runner.Execute(ctx, func(script codersdk.WorkspaceAgentScript) bool {
77-
return true
78-
})
74+
err := runner.Execute(ctx, agentscripts.ExecuteAllScripts)
7975
assert.NoError(t, err)
8076
})
8177
defer func() {
@@ -113,7 +109,7 @@ func TestTimeout(t *testing.T) {
113109
Timeout: time.Millisecond,
114110
}}, aAPI.ScriptCompleted)
115111
require.NoError(t, err)
116-
require.ErrorIs(t, runner.Execute(context.Background(), nil), agentscripts.ErrTimeout)
112+
require.ErrorIs(t, runner.Execute(context.Background(), agentscripts.ExecuteAllScripts), agentscripts.ErrTimeout)
117113
}
118114

119115
func TestScriptReportsTiming(t *testing.T) {
@@ -133,9 +129,7 @@ func TestScriptReportsTiming(t *testing.T) {
133129
Script: "echo hello",
134130
}}, aAPI.ScriptCompleted)
135131
require.NoError(t, err)
136-
require.NoError(t, runner.Execute(context.Background(), func(script codersdk.WorkspaceAgentScript) bool {
137-
return true
138-
}))
132+
require.NoError(t, runner.Execute(context.Background(), agentscripts.ExecuteAllScripts))
139133

140134
log := testutil.RequireRecvCtx(ctx, t, fLogger.logs)
141135
require.Equal(t, "hello", log.Output)

0 commit comments

Comments
 (0)