Fix logic bug, unit test for plan apply method in state store

nomad/state/state_store.go

@@ -218,7 +218,6 @@ func (s *StateStore) UpsertPlanResults(index uint64, results *structs.ApplyPlanR
 		}
 	}
 
-	// TODO(preetha) do a pass to group by jobid
 	// Prepare preempted allocs in the plan results for update
 	for _, preemptedAlloc := range results.NodePreemptions {
 		// Look for existing alloc
@@ -229,7 +228,7 @@ func (s *StateStore) UpsertPlanResults(index uint64, results *structs.ApplyPlanR
 
 		// Nothing to do if this does not exist
 		if existing == nil {
-			return nil
+			continue
 		}
 		exist := existing.(*structs.Allocation)
 
@@ -248,7 +247,6 @@ func (s *StateStore) UpsertPlanResults(index uint64, results *structs.ApplyPlanR
 	}
 
 	// Upsert followup evals for allocs that were preempted
-
 	for _, eval := range results.PreemptionEvals {
 		if err := s.nestedUpsertEval(txn, index, eval); err != nil {
 			return err

nomad/state/state_store_test.go

@@ -244,6 +244,85 @@ func TestStateStore_UpsertPlanResults_Deployment(t *testing.T) {
 	assert.EqualValues(1001, evalOut.ModifyIndex)
 }
 
+// This test checks that:
+// 1) Preempted allocations in plan results are updated
+// 2) Evals are inserted for preempted jobs
+func TestStateStore_UpsertPlanResults_PreemptedAllocs(t *testing.T) {
+	state := testStateStore(t)
+	alloc := mock.Alloc()
+	job := alloc.Job
+	alloc.Job = nil
+
+	require := require.New(t)
+
+	err := state.UpsertJob(999, job)
+	require.Nil(err)
+
+	eval := mock.Eval()
+	eval.JobID = job.ID
+
+	// Create an eval
+	err = state.UpsertEvals(1, []*structs.Evaluation{eval})
+	require.Nil(err)
+
+	// Insert alloc that'll be preempted in the plan
+	preemptedAlloc := mock.Alloc()
+	err = state.UpsertAllocs(2, []*structs.Allocation{preemptedAlloc})
+	require.Nil(err)
+
+	minimalPreemptedAlloc := &structs.Allocation{
+		ID:                 preemptedAlloc.ID,
+		Namespace:          preemptedAlloc.Namespace,
+		DesiredStatus:      structs.AllocDesiredStatusEvict,
+		DesiredDescription: fmt.Sprintf("Preempted by allocation %v", alloc.ID),
+	}
+
+	eval2 := mock.Eval()
+	eval2.JobID = preemptedAlloc.JobID
+
+	// Create a plan result
+	res := structs.ApplyPlanResultsRequest{
+		AllocUpdateRequest: structs.AllocUpdateRequest{
+			Alloc: []*structs.Allocation{alloc},
+			Job:   job,
+		},
+		EvalID:          eval.ID,
+		NodePreemptions: []*structs.Allocation{minimalPreemptedAlloc},
+		PreemptionEvals: []*structs.Evaluation{eval2},
+	}
+
+	err = state.UpsertPlanResults(1000, &res)
+	require.Nil(err)
+
+	ws := memdb.NewWatchSet()
+
+	// Verify alloc and eval created by plan
+	out, err := state.AllocByID(ws, alloc.ID)
+	require.Nil(err)
+	require.Equal(alloc, out)
+
+	index, err := state.Index("allocs")
+	require.Nil(err)
+	require.EqualValues(1000, index)
+
+	evalOut, err := state.EvalByID(ws, eval.ID)
+	require.Nil(err)
+	require.NotNil(evalOut)
+	require.EqualValues(1000, evalOut.ModifyIndex)
+
+	// Verify preempted alloc and eval for preempted job
+	preempted, err := state.AllocByID(ws, preemptedAlloc.ID)
+	require.Nil(err)
+	require.Equal(preempted.DesiredStatus, structs.AllocDesiredStatusEvict)
+	require.Equal(preempted.DesiredDescription, fmt.Sprintf("Preempted by allocation %v", alloc.ID))
+
+	preemptedJobEval, err := state.EvalByID(ws, eval2.ID)
+	require.Nil(err)
+	require.NotNil(preemptedJobEval)
+	require.EqualValues(1000, preemptedJobEval.ModifyIndex)
+
+}
+
 // This test checks that deployment updates are applied correctly
 func TestStateStore_UpsertPlanResults_DeploymentUpdates(t *testing.T) {
 	state := testStateStore(t)