From 31ea4ccbef798104ea141965f19d5c33c5d0e882 Mon Sep 17 00:00:00 2001
From: halleystar <1247920356@qq.com>
Date: Thu, 9 May 2024 14:38:23 +0800
Subject: [PATCH] feat(metrics): add additional labels to metrics

feat(metrics): change additionalLabelsArray to array and use regex to extract labels
---
 pkg/exporter/cmd/config.go   |  3 +-
 pkg/exporter/cmd/server.go   |  6 ++++
 pkg/exporter/nettop/cache.go |  7 ++++
 pkg/exporter/probe/legacy.go | 64 ++++++++++++++++++++++++++++++++++--
 4 files changed, 76 insertions(+), 4 deletions(-)

diff --git a/pkg/exporter/cmd/config.go b/pkg/exporter/cmd/config.go
index 1902f2b7..0c2deadd 100644
--- a/pkg/exporter/cmd/config.go
+++ b/pkg/exporter/cmd/config.go
@@ -16,7 +16,8 @@ type InspServerConfig struct {
 }
 
 type MetricsConfig struct {
-	Probes []ProbeConfig `yaml:"probes" mapstructure:"probes" json:"probes"`
+	Probes           []ProbeConfig `yaml:"probes" mapstructure:"probes" json:"probes"`
+	AdditionalLabels []string      `yaml:"additionalLabels" mapstructure:"additionalLabels" json:"additionalLabels"`
 }
 
 type EventConfig struct {
diff --git a/pkg/exporter/cmd/server.go b/pkg/exporter/cmd/server.go
index 3bc88989..dcb92267 100644
--- a/pkg/exporter/cmd/server.go
+++ b/pkg/exporter/cmd/server.go
@@ -289,6 +289,12 @@ func (i *inspServer) start(cfg *InspServerConfig) error {
 		var err error
 		ctx := context.TODO()
 
+		err = probe.InitAdditionalLabels(cfg.MetricsConfig.AdditionalLabels)
+		if err != nil {
+			log.Errorf("failed init additional labels: %v", err)
+			return
+		}
+
 		log.Infof("start metrics server")
 		i.metricsServer, err = NewMetricsServer()
 		if err != nil {
diff --git a/pkg/exporter/nettop/cache.go b/pkg/exporter/nettop/cache.go
index bfef46de..e202cc5e 100644
--- a/pkg/exporter/nettop/cache.go
+++ b/pkg/exporter/nettop/cache.go
@@ -163,6 +163,7 @@ type Entity struct {
 	podMeta
 	initPid int
 	pids    []int
+	labels  map[string]string
 }
 
 func (e *Entity) String() string {
@@ -177,6 +178,10 @@ func (e *Entity) GetPodNamespace() string {
 	return e.podMeta.namespace
 }
 
+func (e *Entity) GetLabels() map[string]string {
+	return e.labels
+}
+
 func (e *Entity) IsHostNetwork() bool {
 	return e.netnsMeta.isHostNetwork
 }
@@ -353,6 +358,7 @@ func cacheNetTopology(ctx context.Context) error {
 
 		namespace := sandbox.Metadata.Namespace
 		name := sandbox.Metadata.Name
+		labels := sandbox.Labels
 
 		sandboxStatus, err := criClient.PodSandboxStatus(sandbox.Id, true)
 		if err != nil {
@@ -411,6 +417,7 @@ func cacheNetTopology(ctx context.Context) error {
 				namespace: namespace,
 			},
 			initPid: info.Pid,
+			labels:  labels,
 			pids:    pids,
 		}
 
diff --git a/pkg/exporter/probe/legacy.go b/pkg/exporter/probe/legacy.go
index eedee500..4e1d701b 100644
--- a/pkg/exporter/probe/legacy.go
+++ b/pkg/exporter/probe/legacy.go
@@ -2,6 +2,8 @@ package probe
 
 import (
 	"fmt"
+	"regexp"
+	"strings"
 
 	"github.com/alibaba/kubeskoop/pkg/exporter/bpfutil"
 	"github.com/alibaba/kubeskoop/pkg/exporter/nettop"
@@ -12,9 +14,64 @@ import (
 var legacyMetricsLabels = []string{"target_node", "target_namespace", "target_pod", "node", "namespace", "pod"}
 var StandardMetricsLabels = []string{"k8s_node", "k8s_namespace", "k8s_pod"}
 var TupleMetricsLabels = []string{"protocol", "src", "src_type", "src_node", "src_namespace", "src_pod", "dst", "dst_type", "dst_node", "dst_namespace", "dst_pod", "sport", "dport"}
+var AdditionalLabelValueExpr []string
 
 func BuildStandardMetricsLabelValues(entity *nettop.Entity) []string {
-	return []string{nettop.GetNodeName(), entity.GetPodNamespace(), entity.GetPodName()}
+	metaPodLabels := []string{nettop.GetNodeName(), entity.GetPodNamespace(), entity.GetPodName()}
+	return append(metaPodLabels, BuildAdditionalLabelsValues(entity)...)
+}
+
+func InitAdditionalLabels(additionalLabels []string) error {
+	if len(additionalLabels) == 0 {
+		return nil
+	}
+
+	//append to StandardMetricsLabels and AdditionalLabelValueExpr
+	for additionalKV := range additionalLabels {
+		labelKVPair := strings.Split(additionalLabels[additionalKV], "=")
+		StandardMetricsLabels = append(StandardMetricsLabels, strings.TrimSpace(labelKVPair[0]))
+		AdditionalLabelValueExpr = append(AdditionalLabelValueExpr, strings.TrimSpace(labelKVPair[1]))
+	}
+
+	return nil
+}
+
+func BuildAdditionalLabelsValues(entity *nettop.Entity) []string {
+	if len(AdditionalLabelValueExpr) == 0 {
+		return []string{}
+	}
+
+	var values []string
+	podLabels := entity.GetLabels()
+
+	var replaceAllStringSubmatchFunc = func(re *regexp.Regexp, str string, repl func([]string) string) string {
+		result := ""
+		lastIndex := 0
+
+		for _, v := range re.FindAllSubmatchIndex([]byte(str), -1) {
+			groups := []string{}
+			for i := 0; i < len(v); i += 2 {
+				groups = append(groups, str[v[i]:v[i+1]])
+			}
+
+			result += str[lastIndex:v[0]] + repl(groups)
+			lastIndex = v[1]
+		}
+
+		return result + str[lastIndex:]
+	}
+
+	for _, labelValueExpr := range AdditionalLabelValueExpr {
+		podLabelValue := replaceAllStringSubmatchFunc(regexp.MustCompile(`\${labels:(.*?)}`), labelValueExpr, func(groups []string) string {
+			if podLabelValue, ok := podLabels[groups[1]]; ok {
+				return podLabelValue
+			}
+			return ""
+		})
+		values = append(values, podLabelValue)
+	}
+
+	return values
 }
 
 type legacyBatchMetrics struct {
@@ -89,11 +146,12 @@ func (l *legacyBatchMetrics) Collect(metrics chan<- prometheus.Metric) {
 			if err != nil || et == nil {
 				continue
 			}
-			labelValues := []string{nettop.GetNodeName(), et.GetPodNamespace(), et.GetPodName()}
+			labelValues := BuildStandardMetricsLabelValues(et)
 			// for legacy pod labels
 			emit(newMetricsName(l.module, key), labelValues, float64(value))
 
-			labelValues = append(labelValues, labelValues...)
+			var metaPodData = []string{nettop.GetNodeName(), et.GetPodNamespace(), et.GetPodName()}
+			labelValues = append(metaPodData, metaPodData...)
 			emit(legacyMetricsName(l.module, key, l.underscore), labelValues, float64(value))
 		}
 	}