feat: add support for custom labels in the counter metric (#62)

* feat: add support for custom labels to counter metric * add test * revert module change * apply suggestions from code review Co-authored-by: Depado <[email protected]> --------- Co-authored-by: Depado <[email protected]>
depado · Nov 27, 2023 · 5b6bce4 · 5b6bce4
1 parent 0f5f022
commit 5b6bce4
Show file tree

Hide file tree

Showing 4 changed files with 101 additions and 4 deletions.
diff --git a/README.md b/README.md
@@ -28,6 +28,9 @@ Inspired by [github.com/zsais/go-gin-prometheus](https://github.com/zsais/go-gin
 	- [Subsystem](#subsystem)
 	- [Engine](#engine)
 	- [Prometheus Registry](#prometheus-registry)
+	- [HandlerNameFunc](#handlernamefunc)
+	- [RequestPathFunc](#requestpathfunc)
+	- [CustomCounterLabels](#customcounterlabels)
 	- [Ignore](#ignore)
 	- [Token](#token)
 	- [Bucket size](#bucket-size)
@@ -229,6 +232,24 @@ p := ginprom.New(
 r.Use(p.Instrument())
 ```
 
+### CustomCounterLabels
+
+Add custom labels to the counter metric.
+
+```go
+r := gin.Default()
+p := ginprom.New(
+  ginprom.CustomCounterLabels([]string{"client_id"}, func(c *gin.Context) map[string]string {
+    client_id := c.GetHeader("x-client-id")
+    if client_id == "" {
+      client_id = "unknown"
+    }
+    return map[string]string{"client_id": client_id}
+  }),
+)
+r.Use(p.Instrument())
+```
+
 ### Ignore
 
 Ignore allows to completely ignore some routes. Even though you can apply the

diff --git a/options.go b/options.go
@@ -151,3 +151,10 @@ func RequestPathFunc(f func(c *gin.Context) string) PrometheusOption {
 		p.RequestPathFunc = f
 	}
 }
+
+func CustomCounterLabels(labels []string, f func(c *gin.Context) map[string]string) PrometheusOption {
+	return func(p *Prometheus) {
+		p.customCounterLabelsProvider = f
+		p.customCounterLabels = labels
+	}
+}
diff --git a/prom.go b/prom.go
@@ -57,8 +57,10 @@ type Prometheus struct {
 	reqDur       *prometheus.HistogramVec
 	reqSz, resSz prometheus.Summary
 
-	customGauges   pmapGauge
-	customCounters pmapCounter
+	customGauges                pmapGauge
+	customCounters              pmapCounter
+	customCounterLabelsProvider func(c *gin.Context) map[string]string
+	customCounterLabels         []string
 
 	MetricsPath     string
 	Namespace       string
@@ -222,6 +224,7 @@ func New(options ...PrometheusOption) *Prometheus {
 	}
 	p.customGauges.values = make(map[string]prometheus.GaugeVec)
 	p.customCounters.values = make(map[string]prometheus.CounterVec)
+	p.customCounterLabels = make([]string, 0)
 
 	p.Ignored.values = make(map[string]bool)
 	for _, option := range options {
@@ -251,7 +254,7 @@ func (p *Prometheus) register() {
 			Name:      p.RequestCounterMetricName,
 			Help:      "How many HTTP requests processed, partitioned by status code and HTTP method.",
 		},
-		[]string{"code", "method", "handler", "host", "path"},
+		append([]string{"code", "method", "handler", "host", "path"}, p.customCounterLabels...),
 	)
 	p.mustRegister(p.reqCnt)
 
@@ -312,7 +315,15 @@ func (p *Prometheus) Instrument() gin.HandlerFunc {
 		elapsed := float64(time.Since(start)) / float64(time.Second)
 		resSz := float64(c.Writer.Size())
 
-		p.reqCnt.WithLabelValues(status, c.Request.Method, p.HandlerNameFunc(c), c.Request.Host, path).Inc()
+		labels := []string{status, c.Request.Method, p.HandlerNameFunc(c), c.Request.Host, path}
+		if p.customCounterLabelsProvider != nil {
+			extraLabels := p.customCounterLabelsProvider(c)
+			for _, label := range p.customCounterLabels {
+				labels = append(labels, extraLabels[label])
+			}
+		}
+
+		p.reqCnt.WithLabelValues(labels...).Inc()
 		p.reqDur.WithLabelValues(c.Request.Method, path, c.Request.Host).Observe(elapsed)
 		p.reqSz.Observe(float64(reqSz))
 		p.resSz.Observe(resSz)

diff --git a/prom_test.go b/prom_test.go
@@ -366,6 +366,64 @@ func TestEmptyRouter(t *testing.T) {
 	unregister(p)
 }
 
+func TestCustomCounterMetrics(t *testing.T) {
+	r := gin.New()
+	p := New(Engine(r), Registry(prometheus.NewRegistry()), CustomCounterLabels([]string{"client_id", "tenant_id"}, func(c *gin.Context) map[string]string {
+		clientId := c.GetHeader("X-Client-ID")
+		if clientId == "" {
+			clientId = "unknown"
+		}
+		tenantId := c.GetHeader("X-Tenant-ID")
+		if tenantId == "" {
+			tenantId = "unknown"
+		}
+		return map[string]string{
+			"client_id": clientId,
+			"tenant_id": tenantId,
+		}
+	}))
+	r.Use(p.Instrument())
+
+	r.GET("/ping", func(c *gin.Context) {
+		c.JSON(http.StatusOK, gin.H{"status": "ok"})
+	})
+
+	g := gofight.New()
+	g.GET(p.MetricsPath).Run(r, func(r gofight.HTTPResponse, rq gofight.HTTPRequest) {
+		assert.Equal(t, http.StatusOK, r.Code)
+		assert.NotContains(t, r.Body.String(), prometheus.BuildFQName(p.Namespace, p.Subsystem, "requests_total"))
+		assert.NotContains(t, r.Body.String(), "client_id")
+		assert.NotContains(t, r.Body.String(), "tenant_id")
+	})
+
+	g.GET("/ping").Run(r, func(r gofight.HTTPResponse, rq gofight.HTTPRequest) { assert.Equal(t, http.StatusOK, r.Code) })
+
+	g.GET(p.MetricsPath).Run(r, func(r gofight.HTTPResponse, rq gofight.HTTPRequest) {
+		body := r.Body.String()
+		assert.Equal(t, http.StatusOK, r.Code)
+		assert.Contains(t, body, prometheus.BuildFQName(p.Namespace, p.Subsystem, "requests_total"))
+		assert.Contains(t, r.Body.String(), "client_id=\"unknown\"")
+		assert.Contains(t, r.Body.String(), "tenant_id=\"unknown\"")
+		assert.NotContains(t, r.Body.String(), "client_id=\"client-id\"")
+		assert.NotContains(t, r.Body.String(), "tenant_id=\"tenant-id\"")
+	})
+
+	g.GET("/ping").
+		SetHeader(gofight.H{"X-Client-Id": "client-id", "X-Tenant-Id": "tenant-id"}).
+		Run(r, func(r gofight.HTTPResponse, rq gofight.HTTPRequest) { assert.Equal(t, http.StatusOK, r.Code) })
+
+	g.GET(p.MetricsPath).Run(r, func(r gofight.HTTPResponse, rq gofight.HTTPRequest) {
+		body := r.Body.String()
+		assert.Equal(t, http.StatusOK, r.Code)
+		assert.Contains(t, body, prometheus.BuildFQName(p.Namespace, p.Subsystem, "requests_total"))
+		assert.Contains(t, r.Body.String(), "client_id=\"unknown\"")
+		assert.Contains(t, r.Body.String(), "tenant_id=\"unknown\"")
+		assert.Contains(t, r.Body.String(), "client_id=\"client-id\"")
+		assert.Contains(t, r.Body.String(), "tenant_id=\"tenant-id\"")
+	})
+	unregister(p)
+}
+
 func TestIgnore(t *testing.T) {
 	r := gin.New()
 	ipath := "/ping"