From 6c3b3c029bbfbbe2f181d3fd3e354e223638ff99 Mon Sep 17 00:00:00 2001
From: eric <1048315650@qq.com>
Date: Fri, 22 Nov 2024 18:40:29 +0800
Subject: [PATCH] fix(computeresource/tasklevel): Fixed a bug where abnormal
 calculations caused requested resources to exceed limits when Resource type
 was Memory or EphemeralStorage

---
 .../computeresources/tasklevel/tasklevel.go   |  4 ++
 .../tasklevel/tasklevel_test.go               | 70 +++++++++++++++----
 2 files changed, 60 insertions(+), 14 deletions(-)

diff --git a/pkg/internal/computeresources/tasklevel/tasklevel.go b/pkg/internal/computeresources/tasklevel/tasklevel.go
index 87a39d56a0d..ca00d9562a1 100644
--- a/pkg/internal/computeresources/tasklevel/tasklevel.go
+++ b/pkg/internal/computeresources/tasklevel/tasklevel.go
@@ -51,6 +51,10 @@ func computeAverageRequests(requests corev1.ResourceList, steps int) corev1.Reso
 	}
 	averageRequests := corev1.ResourceList{}
 	for k, v := range requests {
+		if k == corev1.ResourceMemory || k == corev1.ResourceEphemeralStorage {
+			averageRequests[k] = *resource.NewQuantity(v.Value()/int64(steps), requests[k].Format)
+			continue
+		}
 		averageRequests[k] = *resource.NewMilliQuantity(v.MilliValue()/int64(steps), requests[k].Format)
 	}
 	return averageRequests
diff --git a/pkg/internal/computeresources/tasklevel/tasklevel_test.go b/pkg/internal/computeresources/tasklevel/tasklevel_test.go
index 3002510aec5..6487c87dd1c 100644
--- a/pkg/internal/computeresources/tasklevel/tasklevel_test.go
+++ b/pkg/internal/computeresources/tasklevel/tasklevel_test.go
@@ -46,12 +46,21 @@ func TestApplyTaskLevelResourceRequirements(t *testing.T) {
 			Command: []string{"cmd"},
 		}},
 		ComputeResources: corev1.ResourceRequirements{
-			Requests: corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("2")},
+			Requests: corev1.ResourceList{
+				corev1.ResourceCPU:    resource.MustParse("2"),
+				corev1.ResourceMemory: resource.MustParse("2Mi"),
+			},
 		},
 		expectedComputeResources: []corev1.ResourceRequirements{{
-			Requests: corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("1")},
+			Requests: corev1.ResourceList{
+				corev1.ResourceCPU:    resource.MustParse("1"),
+				corev1.ResourceMemory: resource.MustParse("1Mi"),
+			},
 		}, {
-			Requests: corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("1")},
+			Requests: corev1.ResourceList{
+				corev1.ResourceCPU:    resource.MustParse("1"),
+				corev1.ResourceMemory: resource.MustParse("1Mi"),
+			},
 		}},
 	}, {
 		desc: "only with limits",
@@ -65,14 +74,29 @@ func TestApplyTaskLevelResourceRequirements(t *testing.T) {
 			Command: []string{"cmd"},
 		}},
 		ComputeResources: corev1.ResourceRequirements{
-			Limits: corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("500m")},
+			Limits: corev1.ResourceList{
+				corev1.ResourceCPU:    resource.MustParse("500m"),
+				corev1.ResourceMemory: resource.MustParse("256Mi"),
+			},
 		},
 		expectedComputeResources: []corev1.ResourceRequirements{{
-			Requests: corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("250m")},
-			Limits:   corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("500m")},
+			Requests: corev1.ResourceList{
+				corev1.ResourceCPU:    resource.MustParse("250m"),
+				corev1.ResourceMemory: resource.MustParse("128Mi"),
+			},
+			Limits: corev1.ResourceList{
+				corev1.ResourceCPU:    resource.MustParse("500m"),
+				corev1.ResourceMemory: resource.MustParse("256Mi"),
+			},
 		}, {
-			Requests: corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("250m")},
-			Limits:   corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("500m")},
+			Requests: corev1.ResourceList{
+				corev1.ResourceCPU:    resource.MustParse("250m"),
+				corev1.ResourceMemory: resource.MustParse("128Mi"),
+			},
+			Limits: corev1.ResourceList{
+				corev1.ResourceCPU:    resource.MustParse("500m"),
+				corev1.ResourceMemory: resource.MustParse("256Mi"),
+			},
 		}},
 	}, {
 		desc: "both with requests and limits",
@@ -86,15 +110,33 @@ func TestApplyTaskLevelResourceRequirements(t *testing.T) {
 			Command: []string{"cmd"},
 		}},
 		ComputeResources: corev1.ResourceRequirements{
-			Requests: corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("1")},
-			Limits:   corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("2")},
+			Requests: corev1.ResourceList{
+				corev1.ResourceCPU:    resource.MustParse("1"),
+				corev1.ResourceMemory: resource.MustParse("2Mi"),
+			},
+			Limits: corev1.ResourceList{
+				corev1.ResourceCPU:    resource.MustParse("2"),
+				corev1.ResourceMemory: resource.MustParse("4Mi"),
+			},
 		},
 		expectedComputeResources: []corev1.ResourceRequirements{{
-			Requests: corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("500m")},
-			Limits:   corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("2")},
+			Requests: corev1.ResourceList{
+				corev1.ResourceCPU:    resource.MustParse("500m"),
+				corev1.ResourceMemory: resource.MustParse("1Mi"),
+			},
+			Limits: corev1.ResourceList{
+				corev1.ResourceCPU:    resource.MustParse("2"),
+				corev1.ResourceMemory: resource.MustParse("4Mi"),
+			},
 		}, {
-			Requests: corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("500m")},
-			Limits:   corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("2")},
+			Requests: corev1.ResourceList{
+				corev1.ResourceCPU:    resource.MustParse("500m"),
+				corev1.ResourceMemory: resource.MustParse("1Mi"),
+			},
+			Limits: corev1.ResourceList{
+				corev1.ResourceCPU:    resource.MustParse("2"),
+				corev1.ResourceMemory: resource.MustParse("4Mi"),
+			},
 		}},
 	}, {
 		desc: "steps with compute resources are overridden by task-level compute resources",