From 3291523d8cd5b35e7d97d0ffda4c4e8dec8261df Mon Sep 17 00:00:00 2001
From: Mahmood Ali <mahmood@hashicorp.com>
Date: Wed, 15 Jan 2020 08:57:05 -0500
Subject: [PATCH] address review comments

---
 client/allocrunner/taskrunner/task_runner.go | 14 ++++-----
 client/state/state_database.go               |  1 +
 client/taskenv/env.go                        |  3 ++
 nomad/structs/structs.go                     | 31 ++++++++++----------
 4 files changed, 26 insertions(+), 23 deletions(-)

diff --git a/client/allocrunner/taskrunner/task_runner.go b/client/allocrunner/taskrunner/task_runner.go
index 506a3aca9..bb6c6a445 100644
--- a/client/allocrunner/taskrunner/task_runner.go
+++ b/client/allocrunner/taskrunner/task_runner.go
@@ -294,16 +294,16 @@ func NewTaskRunner(config *Config) (*TaskRunner, error) {
 
 	// Pull out the task's resources
 	ares := tr.alloc.AllocatedResources
-	if ares != nil {
-		tres, ok := ares.Tasks[tr.taskName]
-		if !ok {
-			return nil, fmt.Errorf("no task resources found on allocation")
-		}
-		tr.taskResources = tres
-	} else {
+	if ares == nil {
 		return nil, fmt.Errorf("no task resources found on allocation")
 	}
 
+	tres, ok := ares.Tasks[tr.taskName]
+	if !ok {
+		return nil, fmt.Errorf("no task resources found on allocation")
+	}
+	tr.taskResources = tres
+
 	// Build the restart tracker.
 	tg := tr.alloc.Job.LookupTaskGroup(tr.alloc.TaskGroup)
 	if tg == nil {
diff --git a/client/state/state_database.go b/client/state/state_database.go
index c513e13da..6d1e65fb2 100644
--- a/client/state/state_database.go
+++ b/client/state/state_database.go
@@ -209,6 +209,7 @@ func (s *BoltStateDB) getAllAllocations(tx *boltdd.Tx) ([]*structs.Allocation, m
 
 		// Handle upgrade path
 		ae.Alloc.Canonicalize()
+		ae.Alloc.Job.Canonicalize()
 
 		allocs = append(allocs, ae.Alloc)
 	}
diff --git a/client/taskenv/env.go b/client/taskenv/env.go
index 0be24d63f..8612651c3 100644
--- a/client/taskenv/env.go
+++ b/client/taskenv/env.go
@@ -604,6 +604,9 @@ func (b *Builder) setAlloc(alloc *structs.Allocation) *Builder {
 	tg := alloc.Job.LookupTaskGroup(alloc.TaskGroup)
 
 	b.otherPorts = make(map[string]string, len(tg.Tasks)*2)
+
+	// Protect against invalid allocs where AllocatedResources isn't set.
+	// TestClient_AddAllocError explicitly tests for this condition
 	if alloc.AllocatedResources != nil {
 		// Populate task resources
 		if tr, ok := alloc.AllocatedResources.Tasks[b.taskName]; ok {
diff --git a/nomad/structs/structs.go b/nomad/structs/structs.go
index 1356e75eb..18c60fbf6 100644
--- a/nomad/structs/structs.go
+++ b/nomad/structs/structs.go
@@ -7634,10 +7634,24 @@ func (a *Allocation) CopySkipJob() *Allocation {
 	return a.copyImpl(false)
 }
 
+// Canonicalize Allocation to ensure fields are initialized to the expectations
+// of this version of Nomad. Should be called when restoring persisted
+// Allocations or receiving Allocations from Nomad agents potentially on an
+// older version of Nomad.
 func (a *Allocation) Canonicalize() {
 	if a.AllocatedResources == nil && a.TaskResources != nil {
 		ar := AllocatedResources{}
-		ar.Tasks = toAllocatedResources(a.TaskResources)
+
+		tasks := make(map[string]*AllocatedTaskResources, len(a.TaskResources))
+		for name, tr := range a.TaskResources {
+			atr := AllocatedTaskResources{}
+			atr.Cpu.CpuShares = int64(tr.CPU)
+			atr.Memory.MemoryMB = int64(tr.MemoryMB)
+			atr.Networks = tr.Networks.Copy()
+
+			tasks[name] = &atr
+		}
+		ar.Tasks = tasks
 
 		if a.SharedResources != nil {
 			ar.Shared.DiskMB = int64(a.SharedResources.DiskMB)
@@ -7652,21 +7666,6 @@ func (a *Allocation) Canonicalize() {
 	// a.Job.Canonicalize()
 }
 
-func toAllocatedResources(taskResources map[string]*Resources) map[string]*AllocatedTaskResources {
-	tasks := make(map[string]*AllocatedTaskResources, len(taskResources))
-
-	for name, tr := range taskResources {
-		atr := AllocatedTaskResources{}
-		atr.Cpu.CpuShares = int64(tr.CPU)
-		atr.Memory.MemoryMB = int64(tr.MemoryMB)
-		atr.Networks = tr.Networks.Copy()
-
-		tasks[name] = &atr
-	}
-
-	return tasks
-}
-
 func (a *Allocation) copyImpl(job bool) *Allocation {
 	if a == nil {
 		return nil