Set unique labels string on prometheus metrics

promutils/labeled/counter.go

@@ -16,13 +16,13 @@ type Counter struct {
 	*prometheus.CounterVec
 
 	prometheus.Counter
-	additionalLabels []contextutils.Key
+	labels []contextutils.Key
 }
 
 // Inc increments the counter by 1. Use Add to increment it by arbitrary non-negative values. The data point will be
 // labeled with values from context. See labeled.SetMetricsKeys for information about how to configure that.
 func (c Counter) Inc(ctx context.Context) {
-	counter, err := c.CounterVec.GetMetricWith(contextutils.Values(ctx, append(metricKeys, c.additionalLabels...)...))
+	counter, err := c.CounterVec.GetMetricWith(contextutils.Values(ctx, c.labels...))
 	if err != nil {
 		panic(err.Error())
 	}
@@ -36,7 +36,7 @@ func (c Counter) Inc(ctx context.Context) {
 // Add adds the given value to the counter. It panics if the value is < 0.. The data point will be labeled with values
 // from context. See labeled.SetMetricsKeys for information about how to configure that.
 func (c Counter) Add(ctx context.Context, v float64) {
-	counter, err := c.CounterVec.GetMetricWith(contextutils.Values(ctx, append(metricKeys, c.additionalLabels...)...))
+	counter, err := c.CounterVec.GetMetricWith(contextutils.Values(ctx, c.labels...))
 	if err != nil {
 		panic(err.Error())
 	}
@@ -47,18 +47,6 @@ 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 {
@@ -73,15 +61,19 @@ 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 {
-			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)
+			// compute unique labels
+			labelSet := sets.NewString(metricStringKeys...)
+			labelSet.Insert(additionalLabels.Labels...)
+			labels := labelSet.List()
+
+			c.CounterVec = scope.MustNewCounterVec(name, description, labels...)
+			c.labels = contextutils.MetricKeysFromStrings(labels)
 		}
 	}
 
 	if c.CounterVec == nil {
 		c.CounterVec = scope.MustNewCounterVec(name, description, metricStringKeys...)
+		c.labels = metricKeys
 	}
 
 	return c

promutils/labeled/counter_test.go

@@ -2,37 +2,112 @@ package labeled
 
 import (
 	"context"
+	"strings"
 	"testing"
 
 	"github.com/flyteorg/flytestdlib/contextutils"
 	"github.com/flyteorg/flytestdlib/promutils"
+
+	"github.com/prometheus/client_golang/prometheus/testutil"
+
 	"github.com/stretchr/testify/assert"
 )
 
 func TestLabeledCounter(t *testing.T) {
 	UnsetMetricKeys()
 	assert.NotPanics(t, func() {
-		SetMetricKeys(contextutils.ProjectKey, contextutils.DomainKey, contextutils.WorkflowIDKey, contextutils.TaskIDKey, contextutils.LaunchPlanIDKey)
+		SetMetricKeys(contextutils.ProjectKey, contextutils.DomainKey, contextutils.WorkflowIDKey, contextutils.TaskIDKey)
+	})
+
+	t.Run("Labeled", func(t *testing.T) {
+		scope := promutils.NewScope("testscope_counter")
+		c := NewCounter("c1", "some desc", scope)
+		assert.NotNil(t, c)
+
+		ctx := context.TODO()
+		var header = `
+			# HELP testscope_counter:c1 some desc
+			# TYPE testscope_counter:c1 counter
+		`
+
+		c.Inc(ctx)
+		c.Add(ctx, 1.0)
+		var expected = `
+			testscope_counter:c1{domain="",project="",task="",wf=""} 2
+		`
+		err := testutil.CollectAndCompare(c.CounterVec, strings.NewReader(header+expected))
+		assert.NoError(t, err)
+
+		ctx = contextutils.WithProjectDomain(ctx, "project", "domain")
+		c.Inc(ctx)
+		c.Add(ctx, 1.0)
+		expected = `
+			testscope_counter:c1{domain="",project="",task="",wf=""} 2
+			testscope_counter:c1{domain="domain",project="project",task="",wf=""} 2
+		`
+		err = testutil.CollectAndCompare(c.CounterVec, strings.NewReader(header+expected))
+		assert.NoError(t, err)
+
+		ctx = contextutils.WithTaskID(ctx, "task")
+		c.Inc(ctx)
+		c.Add(ctx, 1.0)
+		expected = `
+			testscope_counter:c1{domain="",project="",task="",wf=""} 2
+			testscope_counter:c1{domain="domain",project="project",task="",wf=""} 2
+			testscope_counter:c1{domain="domain",project="project",task="task",wf=""} 2
+		`
+		err = testutil.CollectAndCompare(c.CounterVec, strings.NewReader(header+expected))
+		assert.NoError(t, err)
 	})
 
-	scope := promutils.NewTestScope()
-	// 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)
-	c.Add(ctx, 1.0)
-
-	ctx = contextutils.WithProjectDomain(ctx, "project", "domain")
-	c.Inc(ctx)
-	c.Add(ctx, 1.0)
-
-	ctx = contextutils.WithTaskID(ctx, "task")
-	c.Inc(ctx)
-	c.Add(ctx, 1.0)
-
-	ctx = contextutils.WithLaunchPlanID(ctx, "lp")
-	c.Inc(ctx)
-	c.Add(ctx, 1.0)
+	t.Run("Unlabeled", func(t *testing.T) {
+		scope := promutils.NewScope("testscope_counter")
+		c := NewCounter("c2", "some desc", scope, EmitUnlabeledMetric)
+		assert.NotNil(t, c)
+
+		ctx := context.TODO()
+		var header = `
+			# HELP testscope_counter:c2_unlabeled some desc
+			# TYPE testscope_counter:c2_unlabeled counter
+		`
+
+		c.Inc(ctx)
+		c.Add(ctx, 1.0)
+		var expected = `
+			testscope_counter:c2_unlabeled 2
+		`
+		err := testutil.CollectAndCompare(c.Counter, strings.NewReader(header+expected))
+		assert.NoError(t, err)
+	})
+
+	t.Run("AdditionalLabels", func(t *testing.T) {
+		scope := promutils.NewScope("testscope_counter")
+		opts := AdditionalLabelsOption{Labels: []string{contextutils.ProjectKey.String(), contextutils.ExecIDKey.String()}}
+		c := NewCounter("c3", "some desc", scope, opts)
+		assert.NotNil(t, c)
+
+		ctx := context.TODO()
+		var header = `
+			# HELP testscope_counter:c3 some desc
+			# TYPE testscope_counter:c3 counter
+		`
+
+		c.Inc(ctx)
+		c.Add(ctx, 1.0)
+		var expected = `
+			testscope_counter:c3{domain="",exec_id="",project="",task="",wf=""} 2
+		`
+		err := testutil.CollectAndCompare(c.CounterVec, strings.NewReader(header+expected))
+		assert.NoError(t, err)
+
+		ctx = contextutils.WithExecutionID(ctx, "exec_id")
+		c.Inc(ctx)
+		c.Add(ctx, 1.0)
+		expected = `
+			testscope_counter:c3{domain="",exec_id="",project="",task="",wf=""} 2
+			testscope_counter:c3{domain="",exec_id="exec_id",project="",task="",wf=""} 2
+		`
+		err = testutil.CollectAndCompare(c.CounterVec, strings.NewReader(header+expected))
+		assert.NoError(t, err)
+	})
 }

promutils/labeled/gauge.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"
@@ -13,13 +15,13 @@ type Gauge struct {
 	*prometheus.GaugeVec
 
 	prometheus.Gauge
-	additionalLabels []contextutils.Key
+	labels []contextutils.Key
 }
 
 // Inc increments the gauge by 1. Use Add to increment by arbitrary values. The data point will be
 // labeled with values from context. See labeled.SetMetricsKeys for information about how to configure that.
 func (g Gauge) Inc(ctx context.Context) {
-	gauge, err := g.GaugeVec.GetMetricWith(contextutils.Values(ctx, append(metricKeys, g.additionalLabels...)...))
+	gauge, err := g.GaugeVec.GetMetricWith(contextutils.Values(ctx, g.labels...))
 	if err != nil {
 		panic(err.Error())
 	}
@@ -33,7 +35,7 @@ func (g Gauge) Inc(ctx context.Context) {
 // Add adds the given value to the Gauge. (The value can be negative, resulting in a decrease of the Gauge.)
 // The data point will be labeled with values from context. See labeled.SetMetricsKeys for information about how to configure that.
 func (g Gauge) Add(ctx context.Context, v float64) {
-	gauge, err := g.GaugeVec.GetMetricWith(contextutils.Values(ctx, append(metricKeys, g.additionalLabels...)...))
+	gauge, err := g.GaugeVec.GetMetricWith(contextutils.Values(ctx, g.labels...))
 	if err != nil {
 		panic(err.Error())
 	}
@@ -47,7 +49,7 @@ func (g Gauge) Add(ctx context.Context, v float64) {
 // Set sets the Gauge to an arbitrary value.
 // The data point will be labeled with values from context. See labeled.SetMetricsKeys for information about how to configure that.
 func (g Gauge) Set(ctx context.Context, v float64) {
-	gauge, err := g.GaugeVec.GetMetricWith(contextutils.Values(ctx, append(metricKeys, g.additionalLabels...)...))
+	gauge, err := g.GaugeVec.GetMetricWith(contextutils.Values(ctx, g.labels...))
 	if err != nil {
 		panic(err.Error())
 	}
@@ -61,7 +63,7 @@ func (g Gauge) Set(ctx context.Context, v float64) {
 // Dec decrements the level by 1. Use Sub to decrement by arbitrary values. The data point will be
 // labeled with values from context. See labeled.SetMetricsKeys for information about how to configure that.
 func (g Gauge) Dec(ctx context.Context) {
-	gauge, err := g.GaugeVec.GetMetricWith(contextutils.Values(ctx, append(metricKeys, g.additionalLabels...)...))
+	gauge, err := g.GaugeVec.GetMetricWith(contextutils.Values(ctx, g.labels...))
 	if err != nil {
 		panic(err.Error())
 	}
@@ -75,7 +77,7 @@ func (g Gauge) Dec(ctx context.Context) {
 // Sub adds the given value to the Gauge. The value can be negative, resulting in an increase of the Gauge.
 // The data point will be labeled with values from context. See labeled.SetMetricsKeys for information about how to configure that.
 func (g Gauge) Sub(ctx context.Context, v float64) {
-	gauge, err := g.GaugeVec.GetMetricWith(contextutils.Values(ctx, append(metricKeys, g.additionalLabels...)...))
+	gauge, err := g.GaugeVec.GetMetricWith(contextutils.Values(ctx, g.labels...))
 	if err != nil {
 		panic(err.Error())
 	}
@@ -113,14 +115,19 @@ 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 {
-			labels := GetUniqueLabels(metricStringKeys, additionalLabels.Labels)
-			g.GaugeVec = scope.MustNewGaugeVec(name, description, append(metricStringKeys, labels...)...)
-			g.additionalLabels = contextutils.MetricKeysFromStrings(labels)
+			// compute unique labels
+			labelSet := sets.NewString(metricStringKeys...)
+			labelSet.Insert(additionalLabels.Labels...)
+			labels := labelSet.List()
+
+			g.GaugeVec = scope.MustNewGaugeVec(name, description, labels...)
+			g.labels = contextutils.MetricKeysFromStrings(labels)
 		}
 	}
 
 	if g.GaugeVec == nil {
 		g.GaugeVec = scope.MustNewGaugeVec(name, description, metricStringKeys...)
+		g.labels = metricKeys
 	}
 
 	return g