Proper denormalization in optimistic state store

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

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

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
+			}
+		}
+	}
+}