From 9dc002a9b3bc5097defb0b17afeceb5cc3b08f54 Mon Sep 17 00:00:00 2001
From: Alex Dadgar <alex.dadgar@gmail.com>
Date: Mon, 1 May 2017 14:49:57 -0700
Subject: [PATCH] Proper denormalization in optimistic state store

---
 nomad/fsm.go           |  8 +-------
 nomad/plan_apply.go    |  5 +++++
 nomad/structs/funcs.go | 13 +++++++++++++
 3 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/nomad/fsm.go b/nomad/fsm.go
index 0ece2bb99..05fe987f1 100644
--- a/nomad/fsm.go
+++ b/nomad/fsm.go
@@ -422,13 +422,7 @@ func (n *nomadFSM) applyAllocUpdate(buf []byte, index uint64) interface{} {
 	// Attach the job to all the allocations. It is pulled out in the
 	// payload to avoid the redundancy of encoding, but should be denormalized
 	// prior to being inserted into MemDB.
-	if j := req.Job; j != nil {
-		for _, alloc := range req.Alloc {
-			if alloc.Job == nil && !alloc.TerminalStatus() {
-				alloc.Job = j
-			}
-		}
-	}
+	structs.DenormalizeAllocationJobs(req.Job, req.Alloc)
 
 	// Calculate the total resources of allocations. It is pulled out in the
 	// payload to avoid encoding something that can be computed, but should be
diff --git a/nomad/plan_apply.go b/nomad/plan_apply.go
index d7449e17b..5512dff63 100644
--- a/nomad/plan_apply.go
+++ b/nomad/plan_apply.go
@@ -155,6 +155,11 @@ func (s *Server) applyPlan(job *structs.Job, result *structs.PlanResult, snap *s
 
 	// Optimistically apply to our state view
 	if snap != nil {
+		// Attach the job to all the allocations. It is pulled out in the
+		// payload to avoid the redundancy of encoding, but should be denormalized
+		// prior to being inserted into MemDB.
+		structs.DenormalizeAllocationJobs(req.Job, req.Alloc)
+
 		nextIdx := s.raft.AppliedIndex() + 1
 		if err := snap.UpsertAllocs(nextIdx, req.Alloc); err != nil {
 			return future, err
diff --git a/nomad/structs/funcs.go b/nomad/structs/funcs.go
index 479e61389..ab8a980a2 100644
--- a/nomad/structs/funcs.go
+++ b/nomad/structs/funcs.go
@@ -201,3 +201,16 @@ func VaultPoliciesSet(policies map[string]map[string]*Vault) []string {
 	}
 	return flattened
 }
+
+// DenormalizeAllocationJobs is used to attach a job to all allocations that are
+// non-terminal and do not have a job already. This is useful in cases where the
+// job is normalized.
+func DenormalizeAllocationJobs(job *Job, allocs []*Allocation) {
+	if job != nil {
+		for _, alloc := range allocs {
+			if alloc.Job == nil && !alloc.TerminalStatus() {
+				alloc.Job = job
+			}
+		}
+	}
+}