Skip to content

Commit

Permalink
Set unique labels string on prometheus metrics (flyteorg#136)
Browse files Browse the repository at this point in the history
* computing and setting unique labels string on metrics

Signed-off-by: Daniel Rammer <daniel@union.ai>

* added unit tests to counter

Signed-off-by: Daniel Rammer <daniel@union.ai>

* added unit tests to gauge

Signed-off-by: Daniel Rammer <daniel@union.ai>

* added unit tests to stopwatch

Signed-off-by: Daniel Rammer <daniel@union.ai>

* added unit tests to summary

Signed-off-by: Daniel Rammer <daniel@union.ai>

* fixed lint issue

Signed-off-by: Daniel Rammer <daniel@union.ai>

* fixed another lint issue

Signed-off-by: Daniel Rammer <daniel@union.ai>
  • Loading branch information
hamersaw committed Jun 9, 2022
1 parent 70688ae commit 1279ea2
Show file tree
Hide file tree
Showing 8 changed files with 537 additions and 248 deletions.
30 changes: 11 additions & 19 deletions promutils/labeled/counter.go
Original file line number Diff line number Diff line change
Expand Up @@ -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())
}
Expand All @@ -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())
}
Expand All @@ -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 {
Expand All @@ -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
Expand Down
117 changes: 96 additions & 21 deletions promutils/labeled/counter_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
})
}
25 changes: 16 additions & 9 deletions promutils/labeled/gauge.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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())
}
Expand All @@ -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())
}
Expand All @@ -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())
}
Expand All @@ -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())
}
Expand All @@ -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())
}
Expand Down Expand Up @@ -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
Expand Down
Loading

0 comments on commit 1279ea2

Please sign in to comment.