support custom GetCPUQuantity and GetMemoryQuantity func

pkg/agent/qrm-plugins/cpu/dynamicpolicy/cpueviction/strategy/pressure_suppression.go

@@ -102,7 +102,7 @@ func (p *CPUPressureSuppression) GetEvictPods(_ context.Context, request *plugin
 	// sum all pod cpu request
 	totalCPURequest := resource.Quantity{}
 	for _, pod := range filteredPods {
-		totalCPURequest.Add(native.GetCPUQuantity(native.SumUpPodRequestResources(pod)))
+		totalCPURequest.Add(native.CPUQuantityGetter()(native.SumUpPodRequestResources(pod)))
 	}
 	general.Infof("total cpu request is %v, reclaim pool size is %v", totalCPURequest.String(), poolSize)
 
@@ -126,7 +126,7 @@ func (p *CPUPressureSuppression) GetEvictPods(_ context.Context, request *plugin
 					Reason: fmt.Sprintf("current pool suppression rate %.2f is over than the "+
 						"pod suppression tolerance rate %.2f", poolSuppressionRate, podToleranceRate),
 				})
-				totalCPURequest.Sub(native.GetCPUQuantity(native.SumUpPodRequestResources(pod)))
+				totalCPURequest.Sub(native.CPUQuantityGetter()(native.SumUpPodRequestResources(pod)))
 			}
 		} else {
 			p.lastToleranceTime.Delete(key)

pkg/agent/qrm-plugins/cpu/dynamicpolicy/policy.go

@@ -1039,7 +1039,7 @@ func (p *DynamicPolicy) getContainerRequestedCores(allocationInfo *state.Allocat
 			return 0
 		}
 
-		cpuQuantity := native.GetCPUQuantity(container.Resources.Requests)
+		cpuQuantity := native.CPUQuantityGetter()(container.Resources.Requests)
 		allocationInfo.RequestQuantity = general.Max(int(cpuQuantity.Value()), 0)
 		general.Infof("get cpu request quantity: %d for pod: %s/%s container: %s from podWatcher",
 			allocationInfo.RequestQuantity, allocationInfo.PodNamespace, allocationInfo.PodName, allocationInfo.ContainerName)

pkg/agent/qrm-plugins/cpu/nativepolicy/policy.go

@@ -642,7 +642,7 @@ func (p *NativePolicy) getContainerRequestedCores(allocationInfo *state.Allocati
 			return 0
 		}
 
-		cpuQuantity := native.GetCPUQuantity(container.Resources.Requests)
+		cpuQuantity := native.CPUQuantityGetter()(container.Resources.Requests)
 		allocationInfo.RequestQuantity = general.Max(int(cpuQuantity.Value()), 0)
 		general.Infof("get cpu request quantity: %d for pod: %s/%s container: %s from podWatcher",
 			allocationInfo.RequestQuantity, allocationInfo.PodNamespace, allocationInfo.PodName, allocationInfo.ContainerName)

pkg/agent/qrm-plugins/memory/dynamicpolicy/policy.go

@@ -850,7 +850,7 @@ func (p *DynamicPolicy) getContainerRequestedMemoryBytes(allocationInfo *state.A
 		return 0
 	}
 
-	memoryQuantity := native.GetMemoryQuantity(container.Resources.Requests)
+	memoryQuantity := native.MemoryQuantityGetter()(container.Resources.Requests)
 	requestBytes := general.Max(int(memoryQuantity.Value()), 0)
 
 	general.Infof("get memory request bytes: %d for pod: %s/%s container: %s from podWatcher",

pkg/util/native/pod_sorter.go

@@ -61,8 +61,8 @@ func PodCPURequestCmpFunc(i1, i2 interface{}) int {
 	p1Request := SumUpPodRequestResources(i1.(*v1.Pod))
 	p2Request := SumUpPodRequestResources(i2.(*v1.Pod))
 
-	p1CPUQuantity := GetCPUQuantity(p1Request)
-	p2CPUQuantity := GetCPUQuantity(p2Request)
+	p1CPUQuantity := CPUQuantityGetter()(p1Request)
+	p2CPUQuantity := CPUQuantityGetter()(p2Request)
 
 	return p1CPUQuantity.Cmp(p2CPUQuantity)
 }

pkg/util/native/resources.go

@@ -21,6 +21,7 @@ import (
 	"sort"
 	"strconv"
 	"strings"
+	"sync"
 
 	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
@@ -30,6 +31,72 @@ import (
 	"github.com/kubewharf/katalyst-core/pkg/metrics"
 )
 
+type QuantityGetter func(resourceList v1.ResourceList) resource.Quantity
+
+var (
+	quantityGetterMutex  = sync.RWMutex{}
+	cpuQuantityGetter    = DefaultCPUQuantityGetter
+	memoryQuantityGetter = DefaultMemoryQuantityGetter
+)
+
+func CPUQuantityGetter() QuantityGetter {
+	quantityGetterMutex.RLock()
+	defer quantityGetterMutex.RUnlock()
+
+	return cpuQuantityGetter
+}
+
+func SetCPUQuantityGetter(getter QuantityGetter) {
+	quantityGetterMutex.Lock()
+	defer quantityGetterMutex.Unlock()
+
+	cpuQuantityGetter = getter
+}
+
+func MemoryQuantityGetter() QuantityGetter {
+	quantityGetterMutex.RLock()
+	defer quantityGetterMutex.RUnlock()
+
+	return memoryQuantityGetter
+}
+
+func SetMemoryQuantityGetter(getter QuantityGetter) {
+	quantityGetterMutex.Lock()
+	defer quantityGetterMutex.Unlock()
+
+	memoryQuantityGetter = getter
+}
+
+// DefaultCPUQuantityGetter returns cpu quantity for resourceList. since we may have
+// different representations for cpu resource name, the prioritizes will be:
+// native cpu name -> reclaimed milli cpu name
+func DefaultCPUQuantityGetter(resourceList v1.ResourceList) resource.Quantity {
+	if quantity, ok := resourceList[v1.ResourceCPU]; ok {
+		return quantity
+	}
+
+	if quantity, ok := resourceList[consts.ReclaimedResourceMilliCPU]; ok {
+		return *resource.NewMilliQuantity(quantity.Value(), quantity.Format)
+	}
+
+	return resource.Quantity{}
+}
+
+// DefaultMemoryQuantityGetter returns memory quantity for resourceList. since we may have
+// different representations for memory resource name, the prioritizes will be:
+// native memory name -> reclaimed memory name
+func DefaultMemoryQuantityGetter(resourceList v1.ResourceList) resource.Quantity {
+	if quantity, ok := resourceList[v1.ResourceMemory]; ok {
+		return quantity
+	}
+
+	if quantity, ok := resourceList[consts.ReclaimedResourceMemory]; ok {
+		return quantity
+	}
+
+	return resource.Quantity{}
+}
+
 // ResourceThreshold is map of resource name to threshold of water level
 type ResourceThreshold map[v1.ResourceName]float64
 
@@ -95,36 +162,6 @@ func AddResources(a, b v1.ResourceList) v1.ResourceList {
 	return res
 }
 
-// GetCPUQuantity returns cpu quantity for resourceList. since we may have
-// different representations for cpu resource name, the prioritizes will be:
-// native cpu name -> reclaimed milli cpu name
-func GetCPUQuantity(resourceList v1.ResourceList) resource.Quantity {
-	if quantity, ok := resourceList[v1.ResourceCPU]; ok {
-		return quantity
-	}
-
-	if quantity, ok := resourceList[consts.ReclaimedResourceMilliCPU]; ok {
-		return *resource.NewMilliQuantity(quantity.Value(), quantity.Format)
-	}
-
-	return resource.Quantity{}
-}
-
-// GetMemoryQuantity returns memory quantity for resourceList. since we may have
-// different representations for memory resource name, the prioritizes will be:
-// native memory name -> reclaimed memory name
-func GetMemoryQuantity(resourceList v1.ResourceList) resource.Quantity {
-	if quantity, ok := resourceList[v1.ResourceMemory]; ok {
-		return quantity
-	}
-
-	if quantity, ok := resourceList[consts.ReclaimedResourceMemory]; ok {
-		return quantity
-	}
-
-	return resource.Quantity{}
-}
-
 // MergeResources merge multi ResourceList into one ResourceList, the resource of
 // same resource name in all ResourceList we only use the first merged one.
 func MergeResources(updateList ...*v1.ResourceList) *v1.ResourceList {