Append the labels that don't exist in metricStringKeys (#124)

promutils/labeled/counter.go

@@ -3,6 +3,8 @@ package labeled
 import (
 	"context"
 
+	"k8s.io/apimachinery/pkg/util/sets"
+
 	"github.com/flyteorg/flytestdlib/contextutils"
 	"github.com/flyteorg/flytestdlib/promutils"
 	"github.com/prometheus/client_golang/prometheus"
@@ -45,6 +47,18 @@ func (c Counter) Add(ctx context.Context, v float64) {
 	}
 }
 
+// GetUniqueLabels Remove labels from additionalLabels that already exist in metricStringKeys
+func GetUniqueLabels(metricStringKeys []string, additionalLabels []string) []string {
+	labels := make([]string, 0, len(additionalLabels))
+	metricKeysSet := sets.NewString(metricStringKeys...)
+	for _, label := range additionalLabels {
+		if !metricKeysSet.Has(label) {
+			labels = append(labels, label)
+		}
+	}
+	return labels
+}
+
 // NewCounter  creates a new labeled counter. Label keys must be set before instantiating a counter. See labeled.SetMetricsKeys for
 // information about to configure that.
 func NewCounter(name, description string, scope promutils.Scope, opts ...MetricOption) Counter {
@@ -59,8 +73,10 @@ func NewCounter(name, description string, scope promutils.Scope, opts ...MetricO
 		if _, emitUnlabeledMetric := opt.(EmitUnlabeledMetricOption); emitUnlabeledMetric {
 			c.Counter = scope.MustNewCounter(GetUnlabeledMetricName(name), description)
 		} else if additionalLabels, casted := opt.(AdditionalLabelsOption); casted {
-			c.CounterVec = scope.MustNewCounterVec(name, description, append(metricStringKeys, additionalLabels.Labels...)...)
-			c.additionalLabels = contextutils.MetricKeysFromStrings(additionalLabels.Labels)
+			labels := GetUniqueLabels(metricStringKeys, additionalLabels.Labels)
+			// Here we only append the labels that don't exist in metricStringKeys
+			c.CounterVec = scope.MustNewCounterVec(name, description, append(metricStringKeys, labels...)...)
+			c.additionalLabels = contextutils.MetricKeysFromStrings(labels)
 		}
 	}
 

promutils/labeled/counter_test.go

@@ -16,7 +16,9 @@ func TestLabeledCounter(t *testing.T) {
 	})
 
 	scope := promutils.NewTestScope()
-	c := NewCounter("lbl_counter", "help", scope)
+	// Make sure we will not register the same metrics key again.
+	option := AdditionalLabelsOption{Labels: []string{contextutils.ProjectKey.String(), contextutils.DomainKey.String()}}
+	c := NewCounter("lbl_counter", "help", scope, option)
 	assert.NotNil(t, c)
 	ctx := context.TODO()
 	c.Inc(ctx)

promutils/labeled/gauge.go

@@ -113,8 +113,9 @@ func NewGauge(name, description string, scope promutils.Scope, opts ...MetricOpt
 		if _, emitUnlabeledMetric := opt.(EmitUnlabeledMetricOption); emitUnlabeledMetric {
 			g.Gauge = scope.MustNewGauge(GetUnlabeledMetricName(name), description)
 		} else if additionalLabels, casted := opt.(AdditionalLabelsOption); casted {
-			g.GaugeVec = scope.MustNewGaugeVec(name, description, append(metricStringKeys, additionalLabels.Labels...)...)
-			g.additionalLabels = contextutils.MetricKeysFromStrings(additionalLabels.Labels)
+			labels := GetUniqueLabels(metricStringKeys, additionalLabels.Labels)
+			g.GaugeVec = scope.MustNewGaugeVec(name, description, append(metricStringKeys, labels...)...)
+			g.additionalLabels = contextutils.MetricKeysFromStrings(labels)
 		}
 	}
 

promutils/labeled/gauge_test.go

@@ -20,7 +20,9 @@ func TestLabeledGauge(t *testing.T) {
 	scope := promutils.NewScope("testscope")
 	ctx := context.Background()
 	ctx = contextutils.WithProjectDomain(ctx, "flyte", "dev")
-	g := NewGauge("unittest", "some desc", scope)
+	// Make sure we will not register the same metrics key again.
+	option := AdditionalLabelsOption{Labels: []string{contextutils.ProjectKey.String(), contextutils.DomainKey.String()}}
+	g := NewGauge("unittest", "some desc", scope, option)
 	assert.NotNil(t, g)
 
 	g.Inc(ctx)

promutils/labeled/stopwatch.go

@@ -82,9 +82,10 @@ func NewStopWatch(name, description string, scale time.Duration, scope promutils
 		if _, emitUnableMetric := opt.(EmitUnlabeledMetricOption); emitUnableMetric {
 			sw.StopWatch = scope.MustNewStopWatch(GetUnlabeledMetricName(name), description, scale)
 		} else if additionalLabels, casted := opt.(AdditionalLabelsOption); casted {
+			labels := GetUniqueLabels(metricStringKeys, additionalLabels.Labels)
 			sw.StopWatchVec = scope.MustNewStopWatchVec(name, description, scale,
-				append(metricStringKeys, additionalLabels.Labels...)...)
-			sw.additionalLabels = contextutils.MetricKeysFromStrings(additionalLabels.Labels)
+				append(metricStringKeys, labels...)...)
+			sw.additionalLabels = contextutils.MetricKeysFromStrings(labels)
 		}
 	}
 

promutils/labeled/stopwatch_test.go

@@ -18,7 +18,9 @@ func TestLabeledStopWatch(t *testing.T) {
 
 	t.Run("always labeled", func(t *testing.T) {
 		scope := promutils.NewTestScope()
-		c := NewStopWatch("lbl_counter", "help", time.Second, scope)
+		// Make sure we will not register the same metrics key again.
+		option := AdditionalLabelsOption{Labels: []string{contextutils.ProjectKey.String(), contextutils.DomainKey.String()}}
+		c := NewStopWatch("lbl_counter", "help", time.Second, scope, option)
 		assert.NotNil(t, c)
 		ctx := context.TODO()
 		w := c.Start(ctx)

promutils/labeled/summary.go

@@ -43,8 +43,9 @@ func NewSummary(name, description string, scope promutils.Scope, opts ...MetricO
 		if _, emitUnlabeledMetric := opt.(EmitUnlabeledMetricOption); emitUnlabeledMetric {
 			s.Summary = scope.MustNewSummary(GetUnlabeledMetricName(name), description)
 		} else if additionalLabels, casted := opt.(AdditionalLabelsOption); casted {
-			s.SummaryVec = scope.MustNewSummaryVec(name, description, append(metricStringKeys, additionalLabels.Labels...)...)
-			s.additionalLabels = contextutils.MetricKeysFromStrings(additionalLabels.Labels)
+			labels := GetUniqueLabels(metricStringKeys, additionalLabels.Labels)
+			s.SummaryVec = scope.MustNewSummaryVec(name, description, append(metricStringKeys, labels...)...)
+			s.additionalLabels = contextutils.MetricKeysFromStrings(labels)
 		}
 	}
 

promutils/labeled/summary_test.go

@@ -21,7 +21,9 @@ func TestLabeledSummary(t *testing.T) {
 		scope := promutils.NewScope("testscope_summary")
 		ctx := context.Background()
 		ctx = contextutils.WithProjectDomain(ctx, "flyte", "dev")
-		s := NewSummary("unittest", "some desc", scope)
+		// Make sure we will not register the same metrics key again.
+		option := AdditionalLabelsOption{Labels: []string{contextutils.ProjectKey.String(), contextutils.DomainKey.String()}}
+		s := NewSummary("unittest", "some desc", scope, option)
 		assert.NotNil(t, s)
 
 		s.Observe(ctx, 10)