tr: Fetch Wait channel before killTask in restart

Currently, if killTask results in the termination of a process before
calling WaitTask, Restart() will incorrectly return a TaskNotFound
error when using the raw_exec driver on Windows.

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():