diff --git a/client/allocrunner/taskrunner/lifecycle.go b/client/allocrunner/taskrunner/lifecycle.go index 404c9596f..c9fc10a53 100644 --- a/client/allocrunner/taskrunner/lifecycle.go +++ b/client/allocrunner/taskrunner/lifecycle.go @@ -28,18 +28,19 @@ func (tr *TaskRunner) Restart(ctx context.Context, event *structs.TaskEvent, fai // Tell the restart tracker that a restart triggered the exit tr.restartTracker.SetRestartTriggered(failure) + // Grab a handle to the wait channel that will timeout with context cancelation + // _before_ killing the task. + waitCh, err := handle.WaitCh(ctx) + if err != nil { + return err + } + // Kill the task using an exponential backoff in-case of failures. if err := tr.killTask(handle); err != nil { // We couldn't successfully destroy the resource created. tr.logger.Error("failed to kill task. Resources may have been leaked", "error", err) } - // Drain the wait channel or wait for the request context to be canceled - waitCh, err := handle.WaitCh(ctx) - if err != nil { - return err - } - select { case <-waitCh: case <-ctx.Done():