From 06372e12f1a2fd763c2a8a9d465f5b80511de614 Mon Sep 17 00:00:00 2001
From: Wade Simmons <wsimmons@slack-corp.com>
Date: Wed, 27 Mar 2024 14:33:39 -0400
Subject: [PATCH] stats: add ability to set static prometheus labels

This changes lets you define static Prometheus labels in the config.
This can be useful for tagging your stats with the custom labels you
need for your environment.

    stats:
      type: prometheus
      labels:
        myLabelOne: value1
        myLabelTwo: value2
---
 examples/config.yml |  4 ++++
 stats.go            | 19 +++++++++++++++----
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/examples/config.yml b/examples/config.yml
index ff5b403..fbb7c08 100644
--- a/examples/config.yml
+++ b/examples/config.yml
@@ -272,6 +272,10 @@ logging:
   #namespace: prometheusns
   #subsystem: nebula
   #interval: 10s
+  # You can optionally set static labels to include with all metrics
+  #labels:
+  #  myStaticLabel1: value1
+  #  myStaticLabel2: value2
 
   # enables counter metrics for meta packets
   #   e.g.: `messages.tx.handshake`
diff --git a/stats.go b/stats.go
index c88c45c..b5e4a27 100644
--- a/stats.go
+++ b/stats.go
@@ -93,8 +93,19 @@ func startPrometheusStats(l *logrus.Logger, i time.Duration, c *config.C, buildV
 		return nil, fmt.Errorf("stats.path should not be empty")
 	}
 
-	pr := prometheus.NewRegistry()
-	pClient := mp.NewPrometheusProvider(metrics.DefaultRegistry, namespace, subsystem, pr, i)
+	pry := prometheus.NewRegistry()
+	var prr prometheus.Registerer = pry
+
+	labelsRaw := c.GetMap("stats.labels", nil)
+	if labelsRaw != nil {
+		labels := prometheus.Labels{}
+		for k, v := range labelsRaw {
+			labels[fmt.Sprintf("%v", k)] = fmt.Sprintf("%v", v)
+		}
+		prr = prometheus.WrapRegistererWith(labels, prr)
+	}
+
+	pClient := mp.NewPrometheusProvider(metrics.DefaultRegistry, namespace, subsystem, prr, i)
 	if !configTest {
 		go pClient.UpdatePrometheusMetrics()
 	}
@@ -111,14 +122,14 @@ func startPrometheusStats(l *logrus.Logger, i time.Duration, c *config.C, buildV
 			"boringcrypto": strconv.FormatBool(boringEnabled()),
 		},
 	})
-	pr.MustRegister(g)
+	prr.MustRegister(g)
 	g.Set(1)
 
 	var startFn func()
 	if !configTest {
 		startFn = func() {
 			l.Infof("Prometheus stats listening on %s at %s", listen, path)
-			http.Handle(path, promhttp.HandlerFor(pr, promhttp.HandlerOpts{ErrorLog: l}))
+			http.Handle(path, promhttp.HandlerFor(pry, promhttp.HandlerOpts{ErrorLog: l}))
 			log.Fatal(http.ListenAndServe(listen, nil))
 		}
 	}