Initialize task state in allocation sent by scheduler

nomad/mock/mock.go

@@ -1,8 +1,9 @@
 package mock
 
 import (
-	"github.com/hashicorp/nomad/nomad/structs"
 	"time"
+
+	"github.com/hashicorp/nomad/nomad/structs"
 )
 
 func Node() *structs.Node {
@@ -221,6 +222,11 @@ func Alloc() *structs.Allocation {
 				},
 			},
 		},
+		TaskStates: map[string]*structs.TaskState{
+			"web": &structs.TaskState{
+				State: structs.TaskStatePending,
+			},
+		},
 		Job:           Job(),
 		DesiredStatus: structs.AllocDesiredStatusRun,
 		ClientStatus:  structs.AllocClientStatusPending,

scheduler/generic_sched.go

@@ -285,11 +285,13 @@ func (s *GenericScheduler) computePlacements(place []allocTuple) error {
 			alloc.TaskResources = option.TaskResources
 			alloc.DesiredStatus = structs.AllocDesiredStatusRun
 			alloc.ClientStatus = structs.AllocClientStatusPending
+			alloc.TaskStates = initTaskState(missing.TaskGroup, structs.TaskStatePending)
 			s.plan.AppendAlloc(alloc)
 		} else {
 			alloc.DesiredStatus = structs.AllocDesiredStatusFailed
 			alloc.DesiredDescription = "failed to find a node for placement"
 			alloc.ClientStatus = structs.AllocClientStatusFailed
+			alloc.TaskStates = initTaskState(missing.TaskGroup, structs.TaskStateDead)
 			s.plan.AppendFailed(alloc)
 			failedTG[missing.TaskGroup] = alloc
 		}

scheduler/system_sched.go

@@ -252,11 +252,13 @@ func (s *SystemScheduler) computePlacements(place []allocTuple) error {
 			alloc.TaskResources = option.TaskResources
 			alloc.DesiredStatus = structs.AllocDesiredStatusRun
 			alloc.ClientStatus = structs.AllocClientStatusPending
+			alloc.TaskStates = initTaskState(missing.TaskGroup, structs.TaskStatePending)
 			s.plan.AppendAlloc(alloc)
 		} else {
 			alloc.DesiredStatus = structs.AllocDesiredStatusFailed
 			alloc.DesiredDescription = "failed to find a node for placement"
 			alloc.ClientStatus = structs.AllocClientStatusFailed
+			alloc.TaskStates = initTaskState(missing.TaskGroup, structs.TaskStateDead)
 			s.plan.AppendFailed(alloc)
 			failedTG[missing.TaskGroup] = alloc
 		}

scheduler/util.go

@@ -445,3 +445,11 @@ func taskGroupConstraints(tg *structs.TaskGroup) tgConstrainTuple {
 
 	return c
 }
+
+func initTaskState(tg *structs.TaskGroup, state string) map[string]*structs.TaskState {
+	states := make(map[string]*structs.TaskState, len(tg.Tasks))
+	for _, task := range tg.Tasks {
+		states[task.Name] = &structs.TaskState{State: state}
+	}
+	return states
+}

scheduler/util_test.go

@@ -648,3 +648,26 @@ func TestTaskGroupConstraints(t *testing.T) {
 	}
 
 }
+
+func TestInitTaskState(t *testing.T) {
+	tg := &structs.TaskGroup{
+		Tasks: []*structs.Task{
+			&structs.Task{Name: "foo"},
+			&structs.Task{Name: "bar"},
+		},
+	}
+	expPending := map[string]*structs.TaskState{
+		"foo": &structs.TaskState{State: structs.TaskStatePending},
+		"bar": &structs.TaskState{State: structs.TaskStatePending},
+	}
+	expDead := map[string]*structs.TaskState{
+		"foo": &structs.TaskState{State: structs.TaskStateDead},
+		"bar": &structs.TaskState{State: structs.TaskStateDead},
+	}
+	actPending := initTaskState(tg, structs.TaskStatePending)
+	actDead := initTaskState(tg, structs.TaskStateDead)
+
+	if !(reflect.DeepEqual(expPending, actPending) && reflect.DeepEqual(expDead, actDead)) {
+		t.Fatal("Expected and actual not equal")
+	}
+}