From 0f7604ca14c9fd94177b06e62982a8aec9d48f22 Mon Sep 17 00:00:00 2001
From: jay-mckay <jackjaymckay@gmail.com>
Date: Tue, 3 Dec 2024 15:27:43 -0700
Subject: [PATCH] change unit to cgroup label

---
 metrics/collector.go | 28 ++++++++--------------------
 metrics/legacy.go    |  7 +------
 metrics/unified.go   |  7 +------
 3 files changed, 10 insertions(+), 32 deletions(-)

diff --git a/metrics/collector.go b/metrics/collector.go
index 7790ae0..cb4b1d5 100644
--- a/metrics/collector.go
+++ b/metrics/collector.go
@@ -22,8 +22,8 @@ const (
 
 var (
 	namespace  = "cgroup_warden"
-	labels     = []string{"username", "unit"}
-	procLabels = []string{"username", "unit", "proc"}
+	labels     = []string{"cgroup", "username"}
+	procLabels = []string{"cgroup", "username", "proc"}
 	lock       = sync.RWMutex{}
 )
 
@@ -50,7 +50,7 @@ type Collector struct {
 }
 
 type Metric struct {
-	unit        string
+	cgroup      string
 	username    string
 	memoryUsage uint64
 	cpuUsage    float64
@@ -68,12 +68,12 @@ func (c *Collector) Describe(ch chan<- *prometheus.Desc) {
 func (c *Collector) Collect(ch chan<- prometheus.Metric) {
 	stats := c.CollectMetrics()
 	for _, s := range stats {
-		ch <- prometheus.MustNewConstMetric(c.memoryUsage, prometheus.GaugeValue, float64(s.memoryUsage), s.username, s.unit)
-		ch <- prometheus.MustNewConstMetric(c.cpuUsage, prometheus.CounterValue, s.cpuUsage, s.username, s.unit)
+		ch <- prometheus.MustNewConstMetric(c.memoryUsage, prometheus.GaugeValue, float64(s.memoryUsage), s.cgroup, s.username)
+		ch <- prometheus.MustNewConstMetric(c.cpuUsage, prometheus.CounterValue, s.cpuUsage, s.cgroup, s.username)
 		for name, p := range s.processes {
-			ch <- prometheus.MustNewConstMetric(c.procCPU, prometheus.CounterValue, float64(p.cpu), s.username, s.unit, name)
-			ch <- prometheus.MustNewConstMetric(c.procMemory, prometheus.GaugeValue, float64(p.memory), s.username, s.unit, name)
-			ch <- prometheus.MustNewConstMetric(c.procCount, prometheus.GaugeValue, float64(p.count), s.username, s.unit, name)
+			ch <- prometheus.MustNewConstMetric(c.procCPU, prometheus.CounterValue, float64(p.cpu), s.cgroup, s.username, name)
+			ch <- prometheus.MustNewConstMetric(c.procMemory, prometheus.GaugeValue, float64(p.memory), s.cgroup, s.username, name)
+			ch <- prometheus.MustNewConstMetric(c.procCount, prometheus.GaugeValue, float64(p.count), s.cgroup, s.username, name)
 		}
 	}
 }
@@ -194,18 +194,6 @@ func ProcInfo(pids map[uint64]bool) map[string]Process {
 	return processes
 }
 
-var unitRe = regexp.MustCompile(`(user-\d+\.slice)`)
-
-func unitName(cgroup string) (string, error) {
-	match := unitRe.FindStringSubmatch(cgroup)
-
-	if len(match) < 1 {
-		return "", fmt.Errorf("cannot determine slice from '%s'", cgroup)
-	}
-
-	return match[0], nil
-}
-
 var uidRe = regexp.MustCompile(`user-(\d+)\.slice`)
 
 // Looks up the username associated with a user slice cgroup.
diff --git a/metrics/legacy.go b/metrics/legacy.go
index d2c70b3..6b06aa5 100644
--- a/metrics/legacy.go
+++ b/metrics/legacy.go
@@ -68,12 +68,7 @@ func (l *legacy) CreateMetric(group string, pids pidSet) *Metric {
 
 	metric.processes = ProcInfo(pids)
 
-	unit, err := unitName(group)
-	if err != nil {
-		log.Println(err)
-		return nil
-	}
-	metric.unit = unit
+	metric.cgroup = group
 
 	username, err := lookupUsername(group)
 	if err != nil {
diff --git a/metrics/unified.go b/metrics/unified.go
index 7f5c59a..ad58b27 100644
--- a/metrics/unified.go
+++ b/metrics/unified.go
@@ -72,12 +72,7 @@ func (u *unified) CreateMetric(group string, pids pidSet) *Metric {
 
 	metric.processes = ProcInfo(pids)
 
-	unit, err := unitName(group)
-	if err != nil {
-		log.Println(err)
-		return nil
-	}
-	metric.unit = unit
+	metric.cgroup = group
 
 	username, err := lookupUsername(group)
 	if err != nil {