Use joint context to cancel prestart hooks

fixes https://github.com/hashicorp/nomad/issues/6382

The prestart hook for templates blocks while it resolves vault secrets.
If the secret is not found it continues to retry. If a task is shutdown
during this time, the prestart hook currently does not receive
shutdownCtxCancel, causing it to hang.

This PR joins the two contexts so either killCtx or shutdownCtx cancel
and stop the task.

client/allocrunner/taskrunner/task_runner_hooks.go

@@ -6,6 +6,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/LK4D4/joincontext"
 	multierror "github.com/hashicorp/go-multierror"
 	"github.com/hashicorp/nomad/client/allocrunner/interfaces"
 	"github.com/hashicorp/nomad/client/allocrunner/taskrunner/state"
@@ -192,8 +193,11 @@ 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(tr.killCtx, &req, &resp); err != nil {
+		if err := pre.Prestart(joinedCtx, &req, &resp); err != nil {
 			tr.emitHookError(err, name)
 			return structs.WrapRecoverable(fmt.Sprintf("prestart hook %q failed: %v", name, err), err)
 		}

client/allocrunner/taskrunner/task_runner_test.go

@@ -1742,6 +1742,72 @@ func TestTaskRunner_Template_Artifact(t *testing.T) {
 	require.NoErrorf(t, err, "%v not rendered", f2)
 }
 
+// TestTaskRunner_Template_BlockingPreStart asserts that a template
+// that fails to render in PreStart can gracefully be shutdown by
+// either killCtx or shutdownCtx
+func TestTaskRunner_Template_BlockingPreStart(t *testing.T) {
+	t.Parallel()
+
+	alloc := mock.BatchAlloc()
+	task := alloc.Job.TaskGroups[0].Tasks[0]
+	task.Templates = []*structs.Template{
+		{
+			EmbeddedTmpl: `{{ with secret "foo/secret" }}{{ .Data.certificate }}{{ end }}`,
+			DestPath:     "local/test",
+			ChangeMode:   structs.TemplateChangeModeNoop,
+		},
+	}
+
+	task.Vault = &structs.Vault{Policies: []string{"default"}}
+
+	conf, cleanup := testTaskRunnerConfig(t, alloc, task.Name)
+	defer cleanup()
+
+	tr, err := NewTaskRunner(conf)
+	require.NoError(t, err)
+	go tr.Run()
+
+	testutil.WaitForResult(func() (bool, error) {
+		ts := tr.TaskState()
+
+		if len(ts.Events) == 0 {
+			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
+			}
+		}
+
+		if !foundMissingVaultSecretEvent {
+			return false, fmt.Errorf("no missing vault secret template event yet: %#v", ts.Events)
+		}
+
+		return true, nil
+	}, func(err error) {
+		require.NoError(t, err)
+	})
+
+	shutdown := func() <-chan bool {
+		finished := make(chan bool)
+		go func() {
+			tr.Shutdown()
+			finished <- true
+		}()
+
+		return finished
+	}
+
+	select {
+	case <-shutdown():
+		// it shut down like it should have
+	case <-time.After(10 * time.Second):
+		require.Fail(t, "timeout shutting down task")
+	}
+}
+
 // TestTaskRunner_Template_NewVaultToken asserts that a new vault token is
 // created when rendering template and that it is revoked on alloc completion
 func TestTaskRunner_Template_NewVaultToken(t *testing.T) {