scheduler: take into account posstart task to prevent overlapping cpu…

…sets
hashicorp · Oct 25, 2024 · 77b60ba · 77b60ba
1 parent 6236f35
commit 77b60ba
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 0 deletions.
diff --git a/nomad/structs/structs.go b/nomad/structs/structs.go
@@ -3856,6 +3856,7 @@ func (a *AllocatedResources) Comparable() *ComparableResources {
 	prestartSidecarTasks := &AllocatedTaskResources{}
 	prestartEphemeralTasks := &AllocatedTaskResources{}
 	main := &AllocatedTaskResources{}
+	poststartTasks := &AllocatedTaskResources{}
 	poststopTasks := &AllocatedTaskResources{}
 
 	for taskName, r := range a.Tasks {
@@ -3868,12 +3869,15 @@ func (a *AllocatedResources) Comparable() *ComparableResources {
 			} else {
 				prestartEphemeralTasks.Add(r)
 			}
+		} else if lc.Hook == TaskLifecycleHookPoststart {
+			poststartTasks.Add(r)
 		} else if lc.Hook == TaskLifecycleHookPoststop {
 			poststopTasks.Add(r)
 		}
 	}
 
 	// update this loop to account for lifecycle hook
+	main.Add(poststartTasks)
 	prestartEphemeralTasks.Max(main)
 	prestartEphemeralTasks.Max(poststopTasks)
 	prestartSidecarTasks.Add(prestartEphemeralTasks)

diff --git a/nomad/structs/structs_test.go b/nomad/structs/structs_test.go
@@ -7759,6 +7759,66 @@ func TestComparableResources_Superset(t *testing.T) {
 	}
 }
 
+func TestAllocatedResources_Comparable_Flattened(t *testing.T) {
+	ci.Parallel(t)
+
+	allocationResources := AllocatedResources{
+		TaskLifecycles: map[string]*TaskLifecycleConfig{
+			"prestart-task": {
+				Hook:    TaskLifecycleHookPrestart,
+				Sidecar: false,
+			},
+			"prestart-sidecar-task": {
+				Hook:    TaskLifecycleHookPrestart,
+				Sidecar: true,
+			},
+			"poststart-task": {
+				Hook:    TaskLifecycleHookPoststart,
+				Sidecar: false,
+			},
+			"poststop-task": {
+				Hook:    TaskLifecycleHookPoststop,
+				Sidecar: false,
+			},
+		},
+		Tasks: map[string]*AllocatedTaskResources{
+			"prestart-task": {
+				Cpu: AllocatedCpuResources{
+					CpuShares:     2000,
+					ReservedCores: []uint16{0, 1},
+				},
+			},
+			"prestart-sidecar-task": {
+				Cpu: AllocatedCpuResources{
+					CpuShares:     2000,
+					ReservedCores: []uint16{2, 3},
+				},
+			},
+			"main-task": {
+				Cpu: AllocatedCpuResources{
+					CpuShares:     2000,
+					ReservedCores: []uint16{0, 1},
+				},
+			},
+			"poststart-task": {
+				Cpu: AllocatedCpuResources{
+					CpuShares:     2000,
+					ReservedCores: []uint16{4, 5},
+				},
+			},
+			"poststop-task": {
+				Cpu: AllocatedCpuResources{
+					CpuShares:     2000,
+					ReservedCores: []uint16{0, 1},
+				},
+			},
+		},
+	}
+
+	// prestart-sidecar-task + (main-task + poststart-task)
+	must.Len(t, 6, allocationResources.Comparable().Flattened.Cpu.ReservedCores)
+}
+
 func requireErrors(t *testing.T, err error, expected ...string) {
 	t.Helper()
 	require.Error(t, err)