From f3de47e63dfd14971ce8acfacf36203e00bc7364 Mon Sep 17 00:00:00 2001
From: James Rasell <jrasell@users.noreply.github.com>
Date: Thu, 11 Jul 2024 14:54:25 +0100
Subject: [PATCH] quota: Allow cores to be configured within an enterprise
 quota. (#23543)

---
 command/quota_apply.go   | 1 +
 command/quota_init.go    | 2 ++
 command/quota_status.go  | 8 +++++---
 nomad/structs/structs.go | 4 +++-
 4 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/command/quota_apply.go b/command/quota_apply.go
index dddbad692..a7aef1e66 100644
--- a/command/quota_apply.go
+++ b/command/quota_apply.go
@@ -266,6 +266,7 @@ func parseQuotaResource(result *api.Resources, list *ast.ObjectList) error {
 
 	// Check for invalid keys
 	valid := []string{
+		"cores",
 		"cpu",
 		"memory",
 		"memory_max",
diff --git a/command/quota_init.go b/command/quota_init.go
index 0f706354f..badb0b8b3 100644
--- a/command/quota_init.go
+++ b/command/quota_init.go
@@ -119,6 +119,7 @@ description = "Limit the shared default namespace"
 limit {
   region = "global"
   region_limit {
+    cores      = 0
     cpu        = 2500
     memory     = 1000
     memory_max = 1000
@@ -135,6 +136,7 @@ var defaultJsonQuotaSpec = strings.TrimSpace(`
 		{
 			"Region": "global",
 			"RegionLimit": {
+				"Cores": 0,
 				"CPU": 2500,
 				"MemoryMB": 1000,
 				"MemoryMaxMB": 1000
diff --git a/command/quota_status.go b/command/quota_status.go
index b3f9e257c..7aa6d6a47 100644
--- a/command/quota_status.go
+++ b/command/quota_status.go
@@ -188,7 +188,7 @@ func formatQuotaLimits(spec *api.QuotaSpec, usages map[string]*api.QuotaUsage) s
 	sort.Sort(api.QuotaLimitSort(spec.Limits))
 
 	limits := make([]string, len(spec.Limits)+1)
-	limits[0] = "Region|CPU Usage|Memory Usage|Memory Max Usage|Variables Usage"
+	limits[0] = "Region|CPU Usage|Core Usage|Memory Usage|Memory Max Usage|Variables Usage"
 	i := 0
 	for _, specLimit := range spec.Limits {
 		i++
@@ -206,12 +206,13 @@ func formatQuotaLimits(spec *api.QuotaSpec, usages map[string]*api.QuotaUsage) s
 
 		used, ok := lookupUsage()
 		if !ok {
+			cores := fmt.Sprintf("- / %s", formatQuotaLimitInt(specLimit.RegionLimit.Cores))
 			cpu := fmt.Sprintf("- / %s", formatQuotaLimitInt(specLimit.RegionLimit.CPU))
 			memory := fmt.Sprintf("- / %s", formatQuotaLimitInt(specLimit.RegionLimit.MemoryMB))
 			memoryMax := fmt.Sprintf("- / %s", formatQuotaLimitInt(specLimit.RegionLimit.MemoryMaxMB))
 
 			vars := fmt.Sprintf("- / %s", formatQuotaLimitInt(specLimit.VariablesLimit))
-			limits[i] = fmt.Sprintf("%s|%s|%s|%s|%s", specLimit.Region, cpu, memory, memoryMax, vars)
+			limits[i] = fmt.Sprintf("%s|%s|%s|%s|%s|%s", specLimit.Region, cpu, cores, memory, memoryMax, vars)
 			continue
 		}
 
@@ -222,12 +223,13 @@ func formatQuotaLimits(spec *api.QuotaSpec, usages map[string]*api.QuotaUsage) s
 			return *v
 		}
 
+		cores := fmt.Sprintf("%d / %s", orZero(used.RegionLimit.Cores), formatQuotaLimitInt(specLimit.RegionLimit.Cores))
 		cpu := fmt.Sprintf("%d / %s", orZero(used.RegionLimit.CPU), formatQuotaLimitInt(specLimit.RegionLimit.CPU))
 		memory := fmt.Sprintf("%d / %s", orZero(used.RegionLimit.MemoryMB), formatQuotaLimitInt(specLimit.RegionLimit.MemoryMB))
 		memoryMax := fmt.Sprintf("%d / %s", orZero(used.RegionLimit.MemoryMaxMB), formatQuotaLimitInt(specLimit.RegionLimit.MemoryMaxMB))
 
 		vars := fmt.Sprintf("%d / %s", orZero(used.VariablesLimit), formatQuotaLimitInt(specLimit.VariablesLimit))
-		limits[i] = fmt.Sprintf("%s|%s|%s|%s|%s", specLimit.Region, cpu, memory, memoryMax, vars)
+		limits[i] = fmt.Sprintf("%s|%s|%s|%s|%s|%s", specLimit.Region, cpu, cores, memory, memoryMax, vars)
 	}
 
 	return formatList(limits)
diff --git a/nomad/structs/structs.go b/nomad/structs/structs.go
index fed7ade9d..44bb64f33 100644
--- a/nomad/structs/structs.go
+++ b/nomad/structs/structs.go
@@ -3800,11 +3800,13 @@ func (a *AllocatedResources) Comparable() *ComparableResources {
 	return c
 }
 
-// OldTaskResources returns the pre-0.9.0 map of task resources
+// OldTaskResources returns the pre-0.9.0 map of task resources. This
+// functionality is still used within the scheduling code.
 func (a *AllocatedResources) OldTaskResources() map[string]*Resources {
 	m := make(map[string]*Resources, len(a.Tasks))
 	for name, res := range a.Tasks {
 		m[name] = &Resources{
+			Cores:       len(res.Cpu.ReservedCores),
 			CPU:         int(res.Cpu.CpuShares),
 			MemoryMB:    int(res.Memory.MemoryMB),
 			MemoryMaxMB: int(res.Memory.MemoryMaxMB),