task runner: fix goroutine leak in prestart hook

.changelog/11741.txt

@@ -0,0 +1,3 @@
+```release-note:bug
+client: Fixed a memory and goroutine leak for batch tasks and any task that exits without being shut down from the server
+```

client/allocrunner/interfaces/task_lifecycle.go

@@ -89,7 +89,9 @@ 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 or shutdown.
+	// The context is cancelled if the task is killed or shutdown but
+	// should not be stored any persistent goroutines this Prestart
+	// creates.
 	Prestart(context.Context, *TaskPrestartRequest, *TaskPrestartResponse) error
 }
 

client/allocrunner/taskrunner/task_runner_hooks.go

@@ -190,6 +190,11 @@ func (tr *TaskRunner) prestart() error {
 		}()
 	}
 
+	// use a join context to allow any blocking pre-start hooks
+	// to be canceled by either killCtx or shutdownCtx
+	joinedCtx, joinedCancel := joincontext.Join(tr.killCtx, tr.shutdownCtx)
+	defer joinedCancel()
+
 	for _, hook := range tr.runnerHooks {
 		pre, ok := hook.(interfaces.TaskPrestartHook)
 		if !ok {
@@ -235,9 +240,6 @@ func (tr *TaskRunner) prestart() error {
 		}
 
 		// Run the prestart hook
-		// use a joint context to allow any blocking pre-start hooks
-		// to be canceled by either killCtx or shutdownCtx
-		joinedCtx, _ := joincontext.Join(tr.killCtx, tr.shutdownCtx)
 		var resp interfaces.TaskPrestartResponse
 		if err := pre.Prestart(joinedCtx, &req, &resp); err != nil {
 			tr.emitHookError(err, name)