client: properly support hook env vars

The old approach was incomplete. Hook env vars are now:

 * persisted and restored between agent restarts
 * deterministic (LWW if 2 hooks set the same key)

client/allocrunner/taskrunner/state/state.go

@@ -63,7 +63,13 @@ type HookState struct {
 	// Prestart is true if the hook has run Prestart successfully and does
 	// not need to run again
 	PrestartDone bool
-	Data         map[string]string
+
+	// Data allows hooks to persist arbitrary state.
+	Data map[string]string
+
+	// Environment variables set by the hook that will continue to be set
+	// even if PrestartDone=true.
+	Env map[string]string
 }
 
 func (h *HookState) Copy() *HookState {

client/allocrunner/taskrunner/task_runner_hooks.go

@@ -130,6 +130,7 @@ func (tr *TaskRunner) prestart() error {
 		}
 
 		name := pre.Name()
+
 		// Build the request
 		req := interfaces.TaskPrestartRequest{
 			Task:          tr.Task(),
@@ -148,6 +149,8 @@ func (tr *TaskRunner) prestart() error {
 		if origHookState != nil {
 			if origHookState.PrestartDone {
 				tr.logger.Trace("skipping done prestart hook", "name", pre.Name())
+				// Always set env vars from hooks
+				tr.envBuilder.SetHookEnv(name, origHookState.Env)
 				continue
 			}
 
@@ -175,6 +178,7 @@ func (tr *TaskRunner) prestart() error {
 			hookState := &state.HookState{
 				Data:         resp.HookData,
 				PrestartDone: resp.Done,
+				Env:          resp.Env,
 			}
 
 			// Store and persist local state if the hook state has changed
@@ -190,9 +194,7 @@ func (tr *TaskRunner) prestart() error {
 		}
 
 		// Store the environment variables returned by the hook
-		if len(resp.Env) != 0 {
-			tr.envBuilder.SetGenericEnv(resp.Env)
-		}
+		tr.envBuilder.SetHookEnv(name, resp.Env)
 
 		// Store the resources
 		if len(resp.Devices) != 0 {

client/allocrunner/taskrunner/task_runner_test.go

@@ -8,6 +8,7 @@ import (
 	"time"
 
 	"github.com/hashicorp/nomad/client/allocdir"
+	"github.com/hashicorp/nomad/client/allocrunner/interfaces"
 	"github.com/hashicorp/nomad/client/config"
 	consulapi "github.com/hashicorp/nomad/client/consul"
 	"github.com/hashicorp/nomad/client/devicemanager"
@@ -238,7 +239,7 @@ func TestTaskRunner_DevicePropogation(t *testing.T) {
 	dm.ReserveF = func(d *structs.AllocatedDeviceResource) (*device.ContainerReservation, error) {
 		res := &device.ContainerReservation{
 			Envs: map[string]string{
-				"123": "456",
+				"ABC": "123",
 			},
 			Mounts: []*device.Mount{
 				{
@@ -287,5 +288,64 @@ func TestTaskRunner_DevicePropogation(t *testing.T) {
 	require.Equal(driverCfg.Devices[0].Permissions, "123")
 	require.Len(driverCfg.Mounts, 1)
 	require.Equal(driverCfg.Mounts[0].TaskPath, "foo")
-	require.Contains(driverCfg.Env, "123")
+	require.Contains(driverCfg.Env, "ABC")
+}
+
+// mockEnvHook is a test hook that sets an env var and done=true. It fails if
+// it's called more than once.
+type mockEnvHook struct {
+	called int
+}
+
+func (*mockEnvHook) Name() string {
+	return "mock_env_hook"
+}
+
+func (h *mockEnvHook) Prestart(ctx context.Context, req *interfaces.TaskPrestartRequest, resp *interfaces.TaskPrestartResponse) error {
+	h.called++
+
+	resp.Done = true
+	resp.Env = map[string]string{
+		"mock_hook": "1",
+	}
+
+	return nil
+}
+
+// TestTaskRunner_Restore_HookEnv asserts that re-running prestart hooks with
+// hook environments set restores the environment without re-running done
+// hooks.
+func TestTaskRunner_Restore_HookEnv(t *testing.T) {
+	t.Parallel()
+	require := require.New(t)
+
+	alloc := mock.BatchAlloc()
+	task := alloc.Job.TaskGroups[0].Tasks[0]
+	conf, cleanup := testTaskRunnerConfig(t, alloc, task.Name)
+	conf.StateDB = cstate.NewMemDB() // "persist" state between prestart calls
+	defer cleanup()
+
+	tr, err := NewTaskRunner(conf)
+	require.NoError(err)
+
+	// Override the default hooks to only run the mock hook
+	mockHook := &mockEnvHook{}
+	tr.runnerHooks = []interfaces.TaskHook{mockHook}
+
+	// Manually run prestart hooks
+	require.NoError(tr.prestart())
+
+	// Assert env was called
+	require.Equal(1, mockHook.called)
+
+	// Re-running prestart hooks should *not* call done mock hook
+	require.NoError(tr.prestart())
+
+	// Assert env was called
+	require.Equal(1, mockHook.called)
+
+	// Assert the env is still set
+	env := tr.envBuilder.Build().All()
+	require.Contains(env, "mock_hook")
+	require.Equal("1", env["mock_hook"])
 }

client/allocrunner/taskrunner/validate_hook_test.go

@@ -52,12 +52,12 @@ func TestTaskRunner_Validate_ServiceName(t *testing.T) {
 	require.NoError(t, validateTask(task, builder.Build(), conf))
 
 	// Add an env var that should validate
-	builder.SetGenericEnv(map[string]string{"FOO": "bar"})
+	builder.SetHookEnv("test", map[string]string{"FOO": "bar"})
 	task.Services[0].Name = "${FOO}"
 	require.NoError(t, validateTask(task, builder.Build(), conf))
 
 	// Add an env var that should *not* validate
-	builder.SetGenericEnv(map[string]string{"BAD": "invalid/in/consul"})
+	builder.SetHookEnv("test", map[string]string{"BAD": "invalid/in/consul"})
 	task.Services[0].Name = "${BAD}"
 	require.Error(t, validateTask(task, builder.Build(), conf))
 }