From c520d1aa88ef0d36290c40f4fa42e23a8500af1f Mon Sep 17 00:00:00 2001
From: Filipe Roque <froque@premium-minds.com>
Date: Tue, 7 Feb 2023 13:39:40 +0000
Subject: [PATCH] Adds handling of Datadog reserved tag keys

Prepends "fabio-" to reserved tag keys

https://docs.datadoghq.com/getting_started/tagging/#overview
---
 metrics/provider_dogstatsd.go      | 53 ++++++++++++++++++++++++++++--
 metrics/provider_dogstatsd_test.go | 32 +++++++++---------
 2 files changed, 68 insertions(+), 17 deletions(-)

diff --git a/metrics/provider_dogstatsd.go b/metrics/provider_dogstatsd.go
index 22d562e2a..c9a24c7c9 100644
--- a/metrics/provider_dogstatsd.go
+++ b/metrics/provider_dogstatsd.go
@@ -16,11 +16,11 @@ type DogstatsdProvider struct {
 }
 
 func (dp *DogstatsdProvider) NewCounter(name string, labels ...string) gkm.Counter {
-	return dp.D.NewCounter(name, 1)
+	return &dogstatsdCounter{dp.D.NewCounter(name, 1)}
 }
 
 func (dp *DogstatsdProvider) NewGauge(name string, labels ...string) gkm.Gauge {
-	return dp.D.NewGauge(name)
+	return &dogstatsdGauge{dp.D.NewGauge(name)}
 }
 
 func (dp *DogstatsdProvider) NewHistogram(name string, labels ...string) gkm.Histogram {
@@ -42,6 +42,12 @@ func NewDogstatsdProvider(prefix, addr string, interval time.Duration) (*Dogstat
 	return d, nil
 }
 
+type dogstatsdCounter struct {
+	gkm.Counter
+}
+type dogstatsdGauge struct {
+	gkm.Gauge
+}
 type dogstatsdHistogram struct {
 	gkm.Histogram
 }
@@ -49,3 +55,46 @@ type dogstatsdHistogram struct {
 func (dh *dogstatsdHistogram) Observe(value float64) {
 	dh.Histogram.Observe(value * 1000.0)
 }
+
+func (dh *dogstatsdCounter) With(labelValues ...string) gkm.Counter {
+	return dh.Counter.With(correctReservedTagKeys(labelValues)...)
+}
+
+func (dh *dogstatsdGauge) With(labelValues ...string) gkm.Gauge {
+	return dh.Gauge.With(correctReservedTagKeys(labelValues)...)
+}
+
+func (dh *dogstatsdHistogram) With(labelValues ...string) gkm.Histogram {
+	return dh.Histogram.With(correctReservedTagKeys(labelValues)...)
+}
+
+func correctReservedTagKeys(labelValues []string) []string {
+	var rval []string
+	for i, v := range labelValues {
+		if i%2 == 0 {
+			rval = append(rval, correctReservedTagKey(v))
+		} else {
+			rval = append(rval, v)
+		}
+	}
+	return rval
+}
+
+func correctReservedTagKey(label string) string {
+	switch label {
+	case "host":
+		return "fabio-host"
+	case "device":
+		return "fabio-device"
+	case "source":
+		return "fabio-source"
+	case "service":
+		return "fabio-service"
+	case "env":
+		return "fabio-env"
+	case "version":
+		return "fabio-version"
+	default:
+		return label
+	}
+}
diff --git a/metrics/provider_dogstatsd_test.go b/metrics/provider_dogstatsd_test.go
index f673e87eb..dd169937d 100644
--- a/metrics/provider_dogstatsd_test.go
+++ b/metrics/provider_dogstatsd_test.go
@@ -14,22 +14,24 @@ func TestDogstatsdProvider(t *testing.T) {
 	d := dogstatsd.New(prefix, log.NewNopLogger())
 	provider := &DogstatsdProvider{D: d}
 	for _, tst := range []struct {
-		name     string
-		labels   []string
-		values   []string
-		prefix   string
-		countval float64
-		gaugeval float64
-		histoval float64
+		name            string
+		labels          []string
+		values          []string
+		expected_values []string
+		prefix          string
+		countval        float64
+		gaugeval        float64
+		histoval        float64
 	}{
 		{
-			name:     "simpleTest",
-			labels:   []string{"service", "host", "path", "target", "other"},
-			values:   []string{"service", "foo", "host", "bar", "path", "/asdf", "target", "http://jkl.org:1234", "other", "trailer"},
-			prefix:   "tst",
-			countval: 20,
-			gaugeval: 30,
-			histoval: (time.Microsecond * 50).Seconds(),
+			name:            "simpleTest",
+			labels:          []string{"service", "host", "path", "target", "other"},
+			values:          []string{"service", "foo", "host", "bar", "path", "/asdf", "target", "http://jkl.org:1234", "other", "trailer"},
+			expected_values: []string{"fabio-service", "foo", "fabio-host", "bar", "path", "/asdf", "target", "http://jkl.org:1234", "other", "trailer"},
+			prefix:          "tst",
+			countval:        20,
+			gaugeval:        30,
+			histoval:        (time.Microsecond * 50).Seconds(),
 		},
 	} {
 		t.Run(tst.name, func(t *testing.T) {
@@ -59,7 +61,7 @@ func TestDogstatsdProvider(t *testing.T) {
 					if se.value != v.v {
 						t.Errorf("%s failed: expected: %.02f, got %02f", v.n, v.v, se.value)
 					}
-					if len(tst.values) > 0 && !reflect.DeepEqual(se.tags, tst.values) {
+					if len(tst.expected_values) > 0 && !reflect.DeepEqual(se.tags, tst.expected_values) {
 						t.Errorf("tags did not survive round trip parsing")
 					}
 				} else {