Always populate task dir environment variables

Fixes an issue where if a task was restarted after restating the client,
the task dir environment variables would not be populated. This PR fixes
this for both upgrades from 0.8.X and for normal 0.9 restarts.

client/allocrunner/taskrunner/task_dir_hook.go

@@ -13,6 +13,15 @@ import (
 	"github.com/hashicorp/nomad/plugins/drivers"
 )
 
+const (
+	// TaskDirHookIsDoneDataKey is used to mark whether the hook is done. We
+	// do not use the Done response value because we still need to set the
+	// environment variables every time a task starts.
+	// TODO(0.9.1): Use the resp.Env map and switch to resp.Done. We need to
+	// remove usage of the envBuilder
+	TaskDirHookIsDoneDataKey = "is_done"
+)
+
 type taskDirHook struct {
 	runner *TaskRunner
 	logger log.Logger
@@ -33,6 +42,15 @@ func (h *taskDirHook) Name() string {
 }
 
 func (h *taskDirHook) Prestart(ctx context.Context, req *interfaces.TaskPrestartRequest, resp *interfaces.TaskPrestartResponse) error {
+	if v, ok := req.HookData[TaskDirHookIsDoneDataKey]; ok && v == "true" {
+		fsi := h.runner.driverCapabilities.FSIsolation
+		setEnvvars(h.runner.envBuilder, fsi, h.runner.taskDir, h.runner.clientConfig)
+		resp.HookData = map[string]string{
+			TaskDirHookIsDoneDataKey: "true",
+		}
+		return nil
+	}
+
 	cc := h.runner.clientConfig
 	chroot := cconfig.DefaultChrootEnv
 	if len(cc.ChrootEnv) > 0 {
@@ -51,7 +69,9 @@ func (h *taskDirHook) Prestart(ctx context.Context, req *interfaces.TaskPrestart
 
 	// Update the environment variables based on the built task directory
 	setEnvvars(h.runner.envBuilder, fsi, h.runner.taskDir, h.runner.clientConfig)
-	resp.Done = true
+	resp.HookData = map[string]string{
+		TaskDirHookIsDoneDataKey: "true",
+	}
 	return nil
 }
 

client/state/08types.go

@@ -83,7 +83,11 @@ func (t *taskRunnerState08) Upgrade(allocID, taskName string) (*state.LocalState
 
 	// Upgrade task dir state
 	ls.Hooks["task_dir"] = &state.HookState{
-		PrestartDone: t.TaskDirBuilt,
+		Data: map[string]string{
+			// "is_done" is equivalent to task_dir_hook.TaskDirHookIsDoneKey
+			// Does not import to avoid import cycle
+			"is_done": fmt.Sprintf("%v", t.TaskDirBuilt),
+		},
 	}
 
 	// Upgrade dispatch payload state