chore: use generic k8s pod resource function

flanksource · Apr 5, 2024 · 5deba07 · 5deba07
1 parent 2b84774
commit 5deba07
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 42 deletions.
diff --git a/kubernetes/cel_export.go b/kubernetes/cel_export.go
@@ -16,7 +16,7 @@ func Library() []cel.EnvOption {
 		celPodProperties(),
 		celNodeProperties(),
 		celk8sLabels(),
-		celPodMaxMemoryBytes(),
-		celPodMaxCPUMillicores(),
+		celPodResourceLimits(),
+		celPodResourceRequests(),
 	}
 }
diff --git a/kubernetes/topology.go b/kubernetes/topology.go
@@ -1,6 +1,7 @@
 package kubernetes
 
 import (
+	"fmt"
 	"strings"
 
 	"github.com/flanksource/gomplate/v3/conv"
@@ -59,54 +60,33 @@ func celPodProperties() cel.EnvOption {
 	)
 }
 
-func celPodMaxMemoryBytes() cel.EnvOption {
-	return cel.Function("k8s.podMaxMemoryBytes",
-		cel.Overload("k8s.podMaxMemoryBytes_obj_int",
-			[]*cel.Type{cel.AnyType},
+func celPodResourceLimits() cel.EnvOption {
+	return cel.Function("k8s.getResourcesLimit",
+		cel.Overload("k8s.getResourcesLimit_obj_str_int",
+			[]*cel.Type{cel.AnyType, cel.StringType},
 			cel.AnyType,
-			cel.UnaryBinding(func(obj ref.Val) ref.Val {
-				val := conv.ToInt(podMaxMemory(obj.Value()))
+			cel.BinaryBinding(func(obj, resourceType ref.Val) ref.Val {
+				val := getPodResources(obj.Value(), fmt.Sprint(resourceType.Value()), "limits")
 				return types.Int(val)
 			}),
 		),
 	)
 }
 
-func podMaxMemory(input any) int64 {
-	obj := GetUnstructured(input)
-	if obj == nil {
-		return 0
-	}
-
-	var pod corev1.Pod
-	err := runtime.DefaultUnstructuredConverter.FromUnstructured(obj.Object, &pod)
-	if err != nil {
-		return 0
-	}
-	var totalMemBytes int64
-	for _, container := range pod.Spec.Containers {
-		mem := container.Resources.Limits.Memory()
-		if mem != nil {
-			totalMemBytes += _k8sMemoryAsBytes(mem.String())
-		}
-	}
-	return totalMemBytes
-}
-
-func celPodMaxCPUMillicores() cel.EnvOption {
-	return cel.Function("k8s.podMaxCPUMillicores",
-		cel.Overload("k8s.podMaxCPUMillicores_obj_int",
-			[]*cel.Type{cel.AnyType},
+func celPodResourceRequests() cel.EnvOption {
+	return cel.Function("k8s.getResourcesRequests",
+		cel.Overload("k8s.getResourcesRequests_obj_str_int",
+			[]*cel.Type{cel.AnyType, cel.StringType},
 			cel.AnyType,
-			cel.UnaryBinding(func(obj ref.Val) ref.Val {
-				val := conv.ToInt(podMaxCPU(obj.Value()))
+			cel.BinaryBinding(func(obj, resourceType ref.Val) ref.Val {
+				val := getPodResources(obj.Value(), fmt.Sprint(resourceType.Value()), "requests")
 				return types.Int(val)
 			}),
 		),
 	)
 }
 
-func podMaxCPU(input any) int64 {
+func getPodResources(input any, resourceType string, allocType string) int64 {
 	obj := GetUnstructured(input)
 	if obj == nil {
 		return 0
@@ -117,14 +97,53 @@ func podMaxCPU(input any) int64 {
 	if err != nil {
 		return 0
 	}
-	var totalCPU int64
-	for _, container := range pod.Spec.Containers {
-		cpu := container.Resources.Limits.Cpu()
-		if cpu != nil {
-			totalCPU += _k8sCPUAsMillicores(cpu.String())
+
+	if resourceType == "memory" && allocType == "limits" {
+		var totalMemBytes int64
+		for _, container := range pod.Spec.Containers {
+			mem := container.Resources.Limits.Memory()
+			if mem != nil {
+				totalMemBytes += _k8sMemoryAsBytes(mem.String())
+			}
+		}
+		return totalMemBytes
+	}
+
+	if resourceType == "memory" && allocType == "requests" {
+		var totalMemBytes int64
+		for _, container := range pod.Spec.Containers {
+			mem := container.Resources.Requests.Memory()
+			if mem != nil {
+				totalMemBytes += _k8sMemoryAsBytes(mem.String())
+			}
+		}
+		return totalMemBytes
+	}
+
+	if resourceType == "cpu" && allocType == "limits" {
+		var totalCPU int64
+		for _, container := range pod.Spec.Containers {
+			cpu := container.Resources.Limits.Cpu()
+			if cpu != nil {
+				totalCPU += _k8sCPUAsMillicores(cpu.String())
+			}
 		}
+		return totalCPU
+
 	}
-	return totalCPU
+
+	if resourceType == "cpu" && allocType == "requests" {
+		var totalCPU int64
+		for _, container := range pod.Spec.Containers {
+			cpu := container.Resources.Requests.Cpu()
+			if cpu != nil {
+				totalCPU += _k8sCPUAsMillicores(cpu.String())
+			}
+		}
+		return totalCPU
+	}
+
+	return 0
 }
 
 func PodComponentProperties(input any) []map[string]any {