From 72a24aecb028da1219c93967d18fbdf623aebbef Mon Sep 17 00:00:00 2001
From: Michael Schurter <schmichael@hashicorp.com>
Date: Tue, 23 May 2017 15:57:35 -0700
Subject: [PATCH] Add env.Builder.UpdateTask for alloc updates

---
 client/alloc_runner_test.go | 10 +++++++++-
 client/driver/env/env.go    | 23 +++++++++++++++++------
 client/task_runner.go       |  8 ++++++--
 3 files changed, 32 insertions(+), 9 deletions(-)

diff --git a/client/alloc_runner_test.go b/client/alloc_runner_test.go
index 69abbaf7f28..3b68870ccfa 100644
--- a/client/alloc_runner_test.go
+++ b/client/alloc_runner_test.go
@@ -237,6 +237,14 @@ func TestAllocRunner_Destroy(t *testing.T) {
 	ctestutil.ExecCompatible(t)
 	upd, ar := testAllocRunner(false)
 
+	// Shrink chroot
+	ar.config.ChrootEnv = map[string]string{
+		"/bin":   "/bin",
+		"/lib":   "/lib",
+		"/lib32": "/lib32",
+		"/lib64": "/lib64",
+	}
+
 	// Ensure task takes some time
 	task := ar.alloc.Job.TaskGroups[0].Tasks[0]
 	task.Config["command"] = "/bin/sleep"
@@ -269,7 +277,7 @@ func TestAllocRunner_Destroy(t *testing.T) {
 
 			return nil
 		}); err != nil {
-			return false, fmt.Errorf("state not destroyed")
+			return false, fmt.Errorf("state not destroyed: %v", err)
 		}
 
 		// Check the alloc directory was cleaned
diff --git a/client/driver/env/env.go b/client/driver/env/env.go
index f8ffdc10c1e..7a4f3cdacce 100644
--- a/client/driver/env/env.go
+++ b/client/driver/env/env.go
@@ -221,11 +221,8 @@ type Builder struct {
 // NewBuilder creates a new task environment builder.
 func NewBuilder(node *structs.Node, alloc *structs.Allocation, task *structs.Task, region string) *Builder {
 	b := &Builder{
-		region:     region,
-		envvars:    make(map[string]string),
-		nodeAttrs:  make(map[string]string),
-		otherPorts: make(map[string]string),
-		mu:         &sync.RWMutex{},
+		region: region,
+		mu:     &sync.RWMutex{},
 	}
 	return b.setTask(task).setAlloc(alloc).setNode(node)
 }
@@ -344,14 +341,26 @@ func (b *Builder) Build() *TaskEnv {
 	return NewTaskEnv(cleanedEnv, nodeAttrs)
 }
 
+// Update task updates the environment based on a new alloc and task.
+func (b *Builder) UpdateTask(alloc *structs.Allocation, task *structs.Task) *Builder {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	return b.setTask(task).setAlloc(alloc)
+}
+
 // setTask is called from NewBuilder to populate task related environment
 // variables.
 func (b *Builder) setTask(task *structs.Task) *Builder {
 	b.taskName = task.Name
+	b.envvars = make(map[string]string, len(task.Env))
 	for k, v := range task.Env {
 		b.envvars[k] = v
 	}
-	if task.Resources != nil {
+	if task.Resources == nil {
+		b.memLimit = 0
+		b.cpuLimit = 0
+		b.networks = []*structs.NetworkResource{}
+	} else {
 		b.memLimit = task.Resources.MemoryMB
 		b.cpuLimit = task.Resources.CPU
 		// Copy networks to prevent sharing
@@ -380,6 +389,7 @@ func (b *Builder) setAlloc(alloc *structs.Allocation) *Builder {
 	}
 
 	// Add ports from other tasks
+	b.otherPorts = make(map[string]string, len(alloc.TaskResources)*2)
 	for taskName, resources := range alloc.TaskResources {
 		if taskName == b.taskName {
 			continue
@@ -398,6 +408,7 @@ func (b *Builder) setAlloc(alloc *structs.Allocation) *Builder {
 
 // setNode is called from NewBuilder to populate node attributes.
 func (b *Builder) setNode(n *structs.Node) *Builder {
+	b.nodeAttrs = make(map[string]string, 4+len(n.Attributes)+len(n.Meta))
 	b.nodeAttrs[nodeIdKey] = n.ID
 	b.nodeAttrs[nodeNameKey] = n.Name
 	b.nodeAttrs[nodeClassKey] = n.NodeClass
diff --git a/client/task_runner.go b/client/task_runner.go
index 4495a7bba74..dad24a207cb 100644
--- a/client/task_runner.go
+++ b/client/task_runner.go
@@ -1503,6 +1503,9 @@ func (r *TaskRunner) handleUpdate(update *structs.Allocation) error {
 	// Merge in the task resources
 	updatedTask.Resources = update.TaskResources[updatedTask.Name]
 
+	// Update the task's environment
+	r.envBuilder.UpdateTask(update, updatedTask)
+
 	var mErr multierror.Error
 	r.handleLock.Lock()
 	if r.handle != nil {
@@ -1518,7 +1521,8 @@ func (r *TaskRunner) handleUpdate(update *structs.Allocation) error {
 			mErr.Errors = append(mErr.Errors, fmt.Errorf("updating task resources failed: %v", err))
 		}
 
-		if err := r.updateServices(drv, r.handle, r.task, updatedTask, update); err != nil {
+		// Update services in Consul
+		if err := r.updateServices(drv, r.handle, r.task, updatedTask); err != nil {
 			mErr.Errors = append(mErr.Errors, fmt.Errorf("error updating services and checks in Consul: %v", err))
 		}
 	}
@@ -1536,7 +1540,7 @@ func (r *TaskRunner) handleUpdate(update *structs.Allocation) error {
 }
 
 // updateServices and checks with Consul.
-func (r *TaskRunner) updateServices(d driver.Driver, h driver.ScriptExecutor, old, new *structs.Task, newAlloc *structs.Allocation) error {
+func (r *TaskRunner) updateServices(d driver.Driver, h driver.ScriptExecutor, old, new *structs.Task) error {
 	var exec driver.ScriptExecutor
 	if d.Abilities().Exec {
 		// Allow set the script executor if the driver supports it