From 97e3603043ae7a49302bf962bd61969bf611e708 Mon Sep 17 00:00:00 2001
From: Alex Dadgar <alex.dadgar@gmail.com>
Date: Tue, 29 Jan 2019 13:17:10 -0800
Subject: [PATCH] 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.
---
 .../allocrunner/taskrunner/task_dir_hook.go   | 22 ++++++++++++++++++-
 client/state/08types.go                       |  6 ++++-
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/client/allocrunner/taskrunner/task_dir_hook.go b/client/allocrunner/taskrunner/task_dir_hook.go
index 9e54b40b6..da0d84839 100644
--- a/client/allocrunner/taskrunner/task_dir_hook.go
+++ b/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
 }
 
diff --git a/client/state/08types.go b/client/state/08types.go
index 95acd7328..aaa839803 100644
--- a/client/state/08types.go
+++ b/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