From 18835f8e5e66bb120a37ad96415a97fe1a4a520d Mon Sep 17 00:00:00 2001
From: Martijn Vegter <martijn@martijnvegter.com>
Date: Tue, 5 Nov 2024 09:07:15 +0000
Subject: [PATCH] backport of commit 8545e1c79fd4331bc8bf93e1c4cd7cb71ef36db8

---
 .changelog/24297.txt          |  3 ++
 nomad/structs/structs.go      |  4 +++
 nomad/structs/structs_test.go | 61 +++++++++++++++++++++++++++++++++++
 3 files changed, 68 insertions(+)
 create mode 100644 .changelog/24297.txt

diff --git a/.changelog/24297.txt b/.changelog/24297.txt
new file mode 100644
index 00000000000..2fcc371eaa7
--- /dev/null
+++ b/.changelog/24297.txt
@@ -0,0 +1,3 @@
+```release-note:bug
+scheduler: fixed a bug where resource calculation did not account correctly for poststart tasks
+```
diff --git a/nomad/structs/structs.go b/nomad/structs/structs.go
index 2ce8d93f8fc..ba2872beb87 100644
--- a/nomad/structs/structs.go
+++ b/nomad/structs/structs.go
@@ -3859,6 +3859,7 @@ func (a *AllocatedResources) Comparable() *ComparableResources {
 	prestartSidecarTasks := &AllocatedTaskResources{}
 	prestartEphemeralTasks := &AllocatedTaskResources{}
 	main := &AllocatedTaskResources{}
+	poststartTasks := &AllocatedTaskResources{}
 	poststopTasks := &AllocatedTaskResources{}
 
 	for taskName, r := range a.Tasks {
@@ -3871,12 +3872,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..b54cb25484f 100644
--- a/nomad/structs/structs_test.go
+++ b/nomad/structs/structs_test.go
@@ -7759,6 +7759,67 @@ 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:     1000,
+					ReservedCores: []uint16{4},
+				},
+			},
+			"poststart-task": {
+				Cpu: AllocatedCpuResources{
+					CpuShares:     2000,
+					ReservedCores: []uint16{5, 6},
+				},
+			},
+			"poststop-task": {
+				Cpu: AllocatedCpuResources{
+					CpuShares:     2000,
+					ReservedCores: []uint16{7, 8},
+				},
+			},
+		},
+	}
+
+	// The output of Flattened should return the resource required during the execution of the largest lifecycle
+	must.Eq(t, 5000, allocationResources.Comparable().Flattened.Cpu.CpuShares)
+	must.Len(t, 5, allocationResources.Comparable().Flattened.Cpu.ReservedCores)
+}
+
 func requireErrors(t *testing.T, err error, expected ...string) {
 	t.Helper()
 	require.Error(t, err)