From 64a944c7ff5f62f558024e7f8780762b408937e4 Mon Sep 17 00:00:00 2001
From: Diptanu Choudhury <diptanuc@gmail.com>
Date: Mon, 25 Jul 2016 22:11:11 -0700
Subject: [PATCH] Making the queued allocations bind late

---
 nomad/fsm.go               | 5 +++++
 scheduler/generic_sched.go | 6 +++++-
 scheduler/system_sched.go  | 6 +++++-
 3 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/nomad/fsm.go b/nomad/fsm.go
index e1d98ebfb..5a827363c 100644
--- a/nomad/fsm.go
+++ b/nomad/fsm.go
@@ -589,6 +589,8 @@ func (n *nomadFSM) Restore(old io.ReadCloser) error {
 	}
 
 	restore.Commit()
+
+	// Reconciling the queued allocations
 	return n.reconcileSummaries(jobs)
 }
 
@@ -634,6 +636,9 @@ func (n *nomadFSM) reconcileSummaries(jobs []*structs.Job) error {
 		if err != nil {
 			return err
 		}
+		if l := len(planner.Evals); l != 1 {
+			return fmt.Errorf("unexpected number of evals during restore %d. Please file an issue including the logs", l)
+		}
 		for tg, queued := range planner.Evals[0].QueuedAllocations {
 			tgSummary, ok := summary.Summary[tg]
 			if !ok {
diff --git a/scheduler/generic_sched.go b/scheduler/generic_sched.go
index ceed9e88e..acafac9c5 100644
--- a/scheduler/generic_sched.go
+++ b/scheduler/generic_sched.go
@@ -179,7 +179,6 @@ func (s *GenericScheduler) createBlockedEval(planFailure bool) error {
 // process is wrapped in retryMax to iteratively run the handler until we have no
 // further work or we've made the maximum number of attempts.
 func (s *GenericScheduler) process() (bool, error) {
-	s.queuedAllocs = make(map[string]int)
 	// Lookup the Job by ID
 	var err error
 	s.job, err = s.state.JobByID(s.eval.JobID)
@@ -187,6 +186,11 @@ func (s *GenericScheduler) process() (bool, error) {
 		return false, fmt.Errorf("failed to get job '%s': %v",
 			s.eval.JobID, err)
 	}
+	numTaskGroups := 0
+	if s.job != nil {
+		numTaskGroups = len(s.job.TaskGroups)
+	}
+	s.queuedAllocs = make(map[string]int, numTaskGroups)
 
 	// Create a plan
 	s.plan = s.eval.MakePlan(s.job)
diff --git a/scheduler/system_sched.go b/scheduler/system_sched.go
index c705928da..812a8b16e 100644
--- a/scheduler/system_sched.go
+++ b/scheduler/system_sched.go
@@ -85,7 +85,6 @@ func (s *SystemScheduler) Process(eval *structs.Evaluation) error {
 // process is wrapped in retryMax to iteratively run the handler until we have no
 // further work or we've made the maximum number of attempts.
 func (s *SystemScheduler) process() (bool, error) {
-	s.queuedAllocs = make(map[string]int)
 	// Lookup the Job by ID
 	var err error
 	s.job, err = s.state.JobByID(s.eval.JobID)
@@ -93,6 +92,11 @@ func (s *SystemScheduler) process() (bool, error) {
 		return false, fmt.Errorf("failed to get job '%s': %v",
 			s.eval.JobID, err)
 	}
+	numTaskGroups := 0
+	if s.job != nil {
+		numTaskGroups = len(s.job.TaskGroups)
+	}
+	s.queuedAllocs = make(map[string]int, numTaskGroups)
 
 	// Get the ready nodes in the required datacenters
 	if s.job != nil {