From 77b60ba2ff4cd848c2a77a7396e4b2d27c4994cd Mon Sep 17 00:00:00 2001
From: Martijn Vegter <martijn@martijnvegter.com>
Date: Fri, 25 Oct 2024 18:50:23 +0200
Subject: [PATCH] scheduler: take into account posstart task to prevent
 overlapping cpusets

---
 nomad/structs/structs.go      |  4 +++
 nomad/structs/structs_test.go | 60 +++++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+)

diff --git a/nomad/structs/structs.go b/nomad/structs/structs.go
index ee505b46049..fd671c6ba63 100644
--- 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
index 9f739137444..1a3510e6010 100644
--- 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)