diff --git a/internal/worker/listener.go b/internal/worker/listener.go index ec1c88b9..c1c50eae 100644 --- a/internal/worker/listener.go +++ b/internal/worker/listener.go @@ -109,12 +109,8 @@ func (l *Listener) OutputProduced(ctx context.Context, taskID string, outputLoca } func (l *Listener) sendTaskUpdate(ctx context.Context, taskID string, update api.TaskUpdateJSONRequestBody) error { - // Check whether the context is closed before doing anything. - select { - default: - case <-ctx.Done(): + if ctx.Err() != nil { return ctx.Err() } - return l.buffer.SendTaskUpdate(ctx, taskID, update) } diff --git a/internal/worker/task_executor.go b/internal/worker/task_executor.go index c56c33f8..91cc111a 100644 --- a/internal/worker/task_executor.go +++ b/internal/worker/task_executor.go @@ -59,13 +59,11 @@ func (te *TaskExecutor) Run(ctx context.Context, task api.AssignedTask) error { } for _, cmd := range task.Commands { - select { - case <-ctx.Done(): + if ctx.Err() != nil { // Shutdown does not mean task failure; cleanly shutting down will hand // back the task for requeueing on the Manager. logger.Warn().Msg("task execution aborted due to context shutdown") - return nil - default: + return ctx.Err() } runErr := te.cmdRunner.Run(ctx, task.Uuid, cmd)