client: block on context as well as waitCh

For lifecycle operations such as Restart and Kill, the client should not
expect driver plugins to be well behaved and close their waitCh on
context cancelation. Always wait on the passed in context as well as the
waitCh.

client/allocrunner/taskrunner/lifecycle.go

@@ -39,7 +39,10 @@ func (tr *TaskRunner) Restart(ctx context.Context, event *structs.TaskEvent, fai
 		return err
 	}
 
-	<-waitCh
+	select {
+	case <-waitCh:
+	case <-ctx.Done():
+	}
 	return nil
 }
 
@@ -104,7 +107,10 @@ func (tr *TaskRunner) Kill(ctx context.Context, event *structs.TaskEvent) error
 		return err
 	}
 
-	<-waitCh
+	select {
+	case <-waitCh:
+	case <-ctx.Done():
+	}
 
 	// Store that the task has been destroyed and any associated error.
 	tr.UpdateState(structs.TaskStateDead, structs.NewTaskEvent(structs.TaskKilled).SetKillError(killErr))