From 6260e2b9af9d694fa78110027c54af84835cb6fd Mon Sep 17 00:00:00 2001
From: qiuming520 <qiuming_yewu@cmss.chinamobile.com>
Date: Tue, 21 May 2024 19:18:00 +0800
Subject: [PATCH] fix: Modify PromQL to support IPv6

Signed-off-by: qiuming520 <qiuming_yewu@cmss.chinamobile.com>
---
 pkg/controller/annotator/node.go        | 29 +++++++++++++++++++------
 pkg/controller/prometheus/prometheus.go |  9 +++++++-
 pkg/utils/utils.go                      | 13 +++++++++++
 3 files changed, 43 insertions(+), 8 deletions(-)

diff --git a/pkg/controller/annotator/node.go b/pkg/controller/annotator/node.go
index b83183a..0ddf11a 100644
--- a/pkg/controller/annotator/node.go
+++ b/pkg/controller/annotator/node.go
@@ -14,10 +14,9 @@ import (
 	"k8s.io/client-go/util/workqueue"
 	"k8s.io/klog/v2"
 
-	policy "github.com/gocrane/crane-scheduler/pkg/plugins/apis/policy"
-
 	prom "github.com/gocrane/crane-scheduler/pkg/controller/prometheus"
-	utils "github.com/gocrane/crane-scheduler/pkg/utils"
+	"github.com/gocrane/crane-scheduler/pkg/plugins/apis/policy"
+	"github.com/gocrane/crane-scheduler/pkg/utils"
 )
 
 const (
@@ -99,10 +98,16 @@ func (n *nodeController) syncNode(key string) (bool, error) {
 }
 
 func annotateNodeLoad(promClient prom.PromClient, kubeClient clientset.Interface, node *v1.Node, key string) error {
-	value, err := promClient.QueryByNodeIP(key, getNodeInternalIP(node))
+	value, err := promClient.QueryByNodeIP(key, getNodeInternalIP(node, utils.IPv4))
+	if err == nil && len(value) > 0 {
+		return patchNodeAnnotation(kubeClient, node, key, value)
+	}
+
+	value, err = promClient.QueryByNodeIP(key, getNodeInternalIP(node, utils.IPv6))
 	if err == nil && len(value) > 0 {
 		return patchNodeAnnotation(kubeClient, node, key, value)
 	}
+
 	value, err = promClient.QueryByNodeName(key, getNodeName(node))
 	if err == nil && len(value) > 0 {
 		return patchNodeAnnotation(kubeClient, node, key, value)
@@ -176,13 +181,23 @@ func (n *nodeController) CreateMetricSyncTicker(stopCh <-chan struct{}) {
 	}
 }
 
-func getNodeInternalIP(node *v1.Node) string {
+func getNodeInternalIP(node *v1.Node, flag string) string {
 	for _, addr := range node.Status.Addresses {
 		if addr.Type == v1.NodeInternalIP {
-			return addr.Address
+			switch flag {
+			case utils.IPv6:
+				if utils.IsValidIPv6(addr.Address) {
+					return addr.Address
+				}
+			case utils.IPv4:
+				if utils.IsValidIPv4(addr.Address) {
+					return addr.Address
+				}
+			default:
+				return addr.Address
+			}
 		}
 	}
-
 	return node.Name
 }
 
diff --git a/pkg/controller/prometheus/prometheus.go b/pkg/controller/prometheus/prometheus.go
index 09421f2..6f59148 100644
--- a/pkg/controller/prometheus/prometheus.go
+++ b/pkg/controller/prometheus/prometheus.go
@@ -11,6 +11,8 @@ import (
 	v1 "github.com/prometheus/client_golang/api/prometheus/v1"
 	"github.com/prometheus/common/model"
 	"k8s.io/klog/v2"
+
+	"github.com/gocrane/crane-scheduler/pkg/utils"
 )
 
 const (
@@ -57,7 +59,12 @@ func (p *promClient) QueryByNodeIP(metricName, ip string) (string, error) {
 		return result, nil
 	}
 
-	querySelector = fmt.Sprintf("%s{instance=~\"%s:.+\"} /100", metricName, ip)
+	if utils.IsValidIPv6(ip) {
+		querySelector = fmt.Sprintf("%s{instance=~\"\\\\[%s\\\\]:.+\"} /100", metricName, ip)
+	} else {
+		querySelector = fmt.Sprintf("%s{instance=~\"%s:.+\"} /100", metricName, ip)
+	}
+
 	result, err = p.query(querySelector)
 	if result != "" && err == nil {
 		return result, nil
diff --git a/pkg/utils/utils.go b/pkg/utils/utils.go
index 22a15f1..72e4efd 100644
--- a/pkg/utils/utils.go
+++ b/pkg/utils/utils.go
@@ -1,6 +1,7 @@
 package utils
 
 import (
+	"net"
 	"os"
 	"time"
 
@@ -11,6 +12,8 @@ const (
 	TimeFormat       = "2006-01-02T15:04:05Z"
 	DefaultTimeZone  = "Asia/Shanghai"
 	DefaultNamespace = "crane-system"
+	IPv4             = "IPv4"
+	IPv6             = "IPv6"
 )
 
 // IsDaemonsetPod judges if this pod belongs to one daemonset workload.
@@ -66,3 +69,13 @@ func NormalizeScore(value, max, min int64) int64 {
 
 	return value
 }
+
+func IsValidIPv4(ip string) bool {
+	addr := net.ParseIP(ip)
+	return addr != nil && addr.To4() != nil
+}
+
+func IsValidIPv6(ip string) bool {
+	addr := net.ParseIP(ip)
+	return addr != nil && addr.To4() == nil
+}