diff --git a/client/allocrunner/interfaces/task_lifecycle.go b/client/allocrunner/interfaces/task_lifecycle.go
index 680c0d9c2..ee99a507b 100644
--- a/client/allocrunner/interfaces/task_lifecycle.go
+++ b/client/allocrunner/interfaces/task_lifecycle.go
@@ -89,7 +89,7 @@ type TaskPrestartHook interface {
 	// Prestart is called before the task is started including after every
 	// restart. Prestart is not called if the allocation is terminal.
 	//
-	// The context is cancelled if the task is killed.
+	// The context is cancelled if the task is killed or shutdown.
 	Prestart(context.Context, *TaskPrestartRequest, *TaskPrestartResponse) error
 }
 
diff --git a/client/allocrunner/taskrunner/task_runner_test.go b/client/allocrunner/taskrunner/task_runner_test.go
index e1f49157d..25124c3fd 100644
--- a/client/allocrunner/taskrunner/task_runner_test.go
+++ b/client/allocrunner/taskrunner/task_runner_test.go
@@ -1766,6 +1766,7 @@ func TestTaskRunner_Template_BlockingPreStart(t *testing.T) {
 	tr, err := NewTaskRunner(conf)
 	require.NoError(t, err)
 	go tr.Run()
+	defer tr.Shutdown()
 
 	testutil.WaitForResult(func() (bool, error) {
 		ts := tr.TaskState()
@@ -1774,18 +1775,14 @@ func TestTaskRunner_Template_BlockingPreStart(t *testing.T) {
 			return false, fmt.Errorf("no events yet")
 		}
 
-		foundMissingVaultSecretEvent := false
 		for _, e := range ts.Events {
 			if e.Type == "Template" && strings.Contains(e.DisplayMessage, "vault.read(foo/secret)") {
-				foundMissingVaultSecretEvent = true
+				return true, nil
 			}
 		}
 
-		if !foundMissingVaultSecretEvent {
-			return false, fmt.Errorf("no missing vault secret template event yet: %#v", ts.Events)
-		}
+		return false, fmt.Errorf("no missing vault secret template event yet: %#v", ts.Events)
 
-		return true, nil
 	}, func(err error) {
 		require.NoError(t, err)
 	})