From 5d40f7966ebd38dd59a8d3f898d641f41439fc87 Mon Sep 17 00:00:00 2001
From: AllenXu93 <xuxiaozhou93@gmail.com>
Date: Wed, 8 Jan 2025 17:23:05 +0800
Subject: [PATCH] rewrite isWatchable in pod scan

---
 pkg/resourcemonitor/podresourcesscanner.go | 52 +++++++++-------------
 1 file changed, 21 insertions(+), 31 deletions(-)

diff --git a/pkg/resourcemonitor/podresourcesscanner.go b/pkg/resourcemonitor/podresourcesscanner.go
index 878784070..cbecf9105 100644
--- a/pkg/resourcemonitor/podresourcesscanner.go
+++ b/pkg/resourcemonitor/podresourcesscanner.go
@@ -19,10 +19,10 @@ package resourcemonitor
 import (
 	"context"
 	"fmt"
+	"k8s.io/kubernetes/pkg/apis/core/v1/helper/qos"
 	"strconv"
 
 	corev1 "k8s.io/api/core/v1"
-	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	client "k8s.io/client-go/kubernetes"
 	"k8s.io/klog/v2"
@@ -57,58 +57,49 @@ func NewPodResourcesScanner(namespace string, podResourceClient podresourcesapi.
 }
 
 // isWatchable tells if the the given namespace should be watched.
-func (resMon *PodResourcesScanner) isWatchable(podNamespace string, podName string, hasDevice bool) (bool, bool, error) {
-	pod, err := resMon.k8sClient.CoreV1().Pods(podNamespace).Get(context.TODO(), podName, metav1.GetOptions{})
+// In Scan(), if watchable is false, this pods scan will skip
+// so we can return directly if pod's namespace is not watchable
+func (resMon *PodResourcesScanner) isWatchable(podResource *podresourcesapi.PodResources) (bool, bool, error) {
+	if resMon.namespace != "*" && resMon.namespace != podResource.Namespace {
+		return false, false, nil
+	}
+
+	pod, err := resMon.k8sClient.CoreV1().Pods(podResource.Namespace).Get(context.TODO(), podResource.Name, metav1.GetOptions{})
 	if err != nil {
 		return false, false, err
 	}
 
-	isIntegralGuaranteed := pod.Status.QOSClass == corev1.PodQOSGuaranteed && hasExclusiveCPUs(pod)
+	podHasExclusiveCPUs := hasExclusiveCPUs(pod)
+	isPodGuaranteed := qos.GetPodQOS(pod) == corev1.PodQOSGuaranteed
 
-	if resMon.namespace == "*" && (isIntegralGuaranteed || hasDevice) {
-		return true, isIntegralGuaranteed, nil
-	}
-	// TODO:  add an explicit check for guaranteed pods and pods with devices
-	return resMon.namespace == podNamespace && (isIntegralGuaranteed || hasDevice), isIntegralGuaranteed, nil
+	return isPodGuaranteed || hasDevice(podResource), podHasExclusiveCPUs, nil
 }
 
 // hasExclusiveCPUs returns true if a guaranteed pod is allocated exclusive CPUs else returns false.
 // In isWatchable() function we check for the pod QoS and proceed if it is guaranteed (i.e. request == limit)
 // and hence we only check for request in the function below.
 func hasExclusiveCPUs(pod *corev1.Pod) bool {
-	var totalCPU int64
-	var cpuQuantity resource.Quantity
 	for _, container := range pod.Spec.InitContainers {
-
-		var ok bool
-		if cpuQuantity, ok = container.Resources.Requests[corev1.ResourceCPU]; !ok {
-			continue
-		}
-		totalCPU += cpuQuantity.Value()
-		isInitContainerGuaranteed := hasIntegralCPUs(&container)
-		if isInitContainerGuaranteed {
+		if hasIntegralCPUs(&container) {
 			return true
 		}
 	}
 	for _, container := range pod.Spec.Containers {
-		var ok bool
-		if cpuQuantity, ok = container.Resources.Requests[corev1.ResourceCPU]; !ok {
-			continue
-		}
-		totalCPU += cpuQuantity.Value()
-		isAppContainerGuaranteed := hasIntegralCPUs(&container)
-		if isAppContainerGuaranteed {
+		if hasIntegralCPUs(&container) {
 			return true
 		}
 	}
 
-	//No CPUs requested in all the containers in the pod
+	//No integralCPUs requested in all the containers of the pod
 	return false
 }
 
 // hasIntegralCPUs returns true if a container in pod is requesting integral CPUs else returns false
 func hasIntegralCPUs(container *corev1.Container) bool {
-	cpuQuantity := container.Resources.Requests[corev1.ResourceCPU]
+	cpuQuantity, ok := container.Resources.Requests[corev1.ResourceCPU]
+	if !ok {
+		return false
+	}
 	return cpuQuantity.Value()*1000 == cpuQuantity.MilliValue()
 }
 
@@ -146,8 +137,7 @@ func (resMon *PodResourcesScanner) Scan() (ScanResponse, error) {
 
 	for _, podResource := range respPodResources {
 		klog.InfoS("scanning pod", "podName", podResource.GetName())
-		hasDevice := hasDevice(podResource)
-		isWatchable, isIntegralGuaranteed, err := resMon.isWatchable(podResource.GetNamespace(), podResource.GetName(), hasDevice)
+		isWatchable, isExclusiveCPUs, err := resMon.isWatchable(podResource)
 		if err != nil {
 			return ScanResponse{}, fmt.Errorf("checking if pod in a namespace is watchable, namespace:%v, pod name %v: %w", podResource.GetNamespace(), podResource.GetName(), err)
 		}
@@ -166,7 +156,7 @@ func (resMon *PodResourcesScanner) Scan() (ScanResponse, error) {
 			}
 
 			cpuIDs := container.GetCpuIds()
-			if len(cpuIDs) > 0 && isIntegralGuaranteed {
+			if len(cpuIDs) > 0 && isExclusiveCPUs {
 				var resCPUs []string
 				for _, cpuID := range container.GetCpuIds() {
 					resCPUs = append(resCPUs, strconv.FormatInt(cpuID, 10))