updates based on review.

hashicorp · Sep 29, 2023 · 8b95ad6 · 8b95ad6
1 parent 729421a
commit 8b95ad6
Show file tree

Hide file tree

Showing 2 changed files with 113 additions and 6 deletions.
diff --git a/nomad/plan_apply.go b/nomad/plan_apply.go
@@ -832,7 +832,7 @@ func evaluatePlanAllocIndexes(stateSnap *state.StateSnapshot, plan *structs.Plan
 	// allocation name is formed using the jobID, taskgroup name, and alloc
 	// index. We can therefore use this to identify alloc index duplicates and
 	// save the computational overhead of extracting the index from the name.
-	allocIndexMap := make(map[string]string)
+	allocIndexMap := make(map[string]*set.Set[string])
 
 	stateAllocs, err := stateSnap.AllocsByJob(nil, plan.Job.Namespace, plan.Job.ID, true)
 	if err != nil {
@@ -846,7 +846,10 @@ func evaluatePlanAllocIndexes(stateSnap *state.StateSnapshot, plan *structs.Plan
 		switch alloc.Job.Version {
 		case plan.Job.Version:
 			if !alloc.TerminalStatus() && alloc.ClientStatus != structs.AllocClientStatusUnknown {
-				allocIndexMap[alloc.GetName()] = alloc.GetID()
+				if _, ok := allocIndexMap[alloc.GetName()]; !ok {
+					allocIndexMap[alloc.GetName()] = set.New[string](0)
+				}
+				allocIndexMap[alloc.GetName()].Insert(alloc.GetID())
 			}
 		default:
 		}
@@ -858,8 +861,10 @@ func evaluatePlanAllocIndexes(stateSnap *state.StateSnapshot, plan *structs.Plan
 	// and a plan is generated to replace the failed allocations.
 	for _, nodeUpdate := range plan.NodeUpdate {
 		for _, allocUpdate := range nodeUpdate {
-			if allocID, ok := allocIndexMap[allocUpdate.GetName()]; ok && allocID == allocUpdate.GetID() {
-				delete(allocIndexMap, allocUpdate.GetName())
+			if allocIDSet, ok := allocIndexMap[allocUpdate.GetName()]; ok {
+				if allocIDSet.Remove(allocUpdate.GetID()) && allocIDSet.Size() == 0 {
+					delete(allocIndexMap, allocUpdate.GetName())
+				}
 			}
 		}
 	}
@@ -888,8 +893,10 @@ func evaluatePlanAllocIndexes(stateSnap *state.StateSnapshot, plan *structs.Plan
 			// plan.
 			switch nodeAlloc.ClientStatus {
 			case structs.AllocClientStatusUnknown:
-				if allocID, ok := allocIndexMap[nodeAlloc.GetName()]; ok && allocID == nodeAlloc.GetID() {
-					delete(allocIndexMap, nodeAlloc.GetName())
+				if allocIDSet, ok := allocIndexMap[nodeAlloc.GetName()]; ok {
+					if allocIDSet.Remove(nodeAlloc.GetID()) && allocIDSet.Size() == 0 {
+						delete(allocIndexMap, nodeAlloc.GetName())
+					}
 				}
 			default:
 				if !nodeAllocSet.Insert(nodeAlloc.Name) {

diff --git a/nomad/plan_apply_test.go b/nomad/plan_apply_test.go
@@ -1609,4 +1609,104 @@ func Test_evaluatePlanAllocIndexes(t *testing.T) {
 		}
 		must.NoError(t, evaluatePlanAllocIndexes(testStateSnapshotFn(t, testState), &testPlan))
 	})
+
+	// Perform a test where duplicates already exist within state. This is
+	// possible on all but fresh clusters.
+	t.Run("duplicate exist in state", func(t *testing.T) {
+
+		// Generate and insert 2 test allocations into state. These have
+		// duplicate alloc indexes, to mimic Nomad upgrading on a cluster where
+		// these are already persisted to state.
+		testAllocs := []*structs.Allocation{mock.Alloc(), mock.Alloc()}
+
+		for _, testAlloc := range testAllocs {
+			testAlloc.Job = testAllocs[0].Job
+			testAlloc.JobID = testAllocs[0].Job.ID
+			testAlloc.ClientStatus = structs.AllocClientStatusRunning
+			testAlloc.Name = fmt.Sprintf("%s.%s[%v]", testAlloc.JobID, testAlloc.TaskGroup, 0)
+		}
+
+		must.NoError(t, testState.UpsertAllocs(structs.MsgTypeTestSetup, 70, testAllocs))
+
+		testPlan := structs.Plan{
+			Job: &structs.Job{
+				ID:        testAllocs[0].JobID,
+				Namespace: testAllocs[0].Namespace,
+				Version:   1,
+				Type:      structs.JobTypeService,
+			},
+			NodeUpdate: map[string][]*structs.Allocation{
+				"8bdb21db-9445-3650-ca9d-0d7883cc8a73": {
+					{Name: testAllocs[0].Name, ID: testAllocs[0].ID},
+				},
+				"52b3508a-c88e-fc83-74c9-829d5ef1103a": {
+					{Name: testAllocs[1].Name, ID: testAllocs[1].ID},
+				},
+			},
+			NodeAllocation: map[string][]*structs.Allocation{
+				"8bdb21db-9445-3650-ca9d-0d7883cc8a73": {
+					{Name: fmt.Sprintf("%s.%s[%v]",
+						testAllocs[0].JobID, testAllocs[0].TaskGroup, 0),
+						ID: "bd8cf1bc-5ded-5942-5b39-078878e74e15"},
+				},
+				"52b3508a-c88e-fc83-74c9-829d5ef1103a": {
+					{Name: fmt.Sprintf("%s.%s[%v]",
+						testAllocs[1].JobID, testAllocs[1].TaskGroup, 1),
+						ID: "bd8cf1bc-5ded-5942-5b39-078878e74e15"},
+				},
+			},
+		}
+		must.NoError(t, evaluatePlanAllocIndexes(testStateSnapshotFn(t, testState), &testPlan))
+
+		// Delete the alloc state entry, so it doesn't impact any subsequent
+		// tests.
+		must.NoError(t, testState.DeleteEval(
+			80, []string{}, []string{testAllocs[0].ID, testAllocs[1].ID}, false))
+	})
+
+	t.Run("duplicate exist in state conflict", func(t *testing.T) {
+
+		// Generate and insert 2 test allocations into state. These have
+		// duplicate alloc indexes, to mimic Nomad upgrading on a cluster where
+		// these are already persisted to state.
+		testAllocs := []*structs.Allocation{mock.Alloc(), mock.Alloc()}
+
+		for _, testAlloc := range testAllocs {
+			testAlloc.Job = testAllocs[0].Job
+			testAlloc.JobID = testAllocs[0].Job.ID
+			testAlloc.ClientStatus = structs.AllocClientStatusRunning
+			testAlloc.Name = fmt.Sprintf("%s.%s[%v]", testAlloc.JobID, testAlloc.TaskGroup, 0)
+		}
+
+		must.NoError(t, testState.UpsertAllocs(structs.MsgTypeTestSetup, 90, testAllocs))
+
+		testPlan := structs.Plan{
+			Job: &structs.Job{
+				ID:        testAllocs[0].JobID,
+				Namespace: testAllocs[0].Namespace,
+				Version:   0,
+				Type:      structs.JobTypeService,
+			},
+			NodeUpdate: map[string][]*structs.Allocation{
+				"8bdb21db-9445-3650-ca9d-0d7883cc8a73": {
+					{Name: testAllocs[0].Name, ID: testAllocs[0].ID},
+				},
+			},
+			NodeAllocation: map[string][]*structs.Allocation{
+				"8bdb21db-9445-3650-ca9d-0d7883cc8a73": {
+					{Name: fmt.Sprintf("%s.%s[%v]",
+						testAllocs[0].JobID, testAllocs[0].TaskGroup, 0),
+						ID: "bd8cf1bc-5ded-5942-5b39-078878e74e15"},
+				},
+			},
+		}
+		must.ErrorContains(t,
+			evaluatePlanAllocIndexes(testStateSnapshotFn(t, testState), &testPlan),
+			duplicateAllocIndexErrorString)
+
+		// Delete the alloc state entry, so it doesn't impact any subsequent
+		// tests.
+		must.NoError(t, testState.DeleteEval(
+			100, []string{}, []string{testAllocs[0].ID, testAllocs[1].ID}, false))
+	})
 }