Labeled gauge (flyteorg#62)

promutils/labeled/gauge.go

@@ -0,0 +1,125 @@
+package labeled
+
+import (
+	"context"
+
+	"github.com/lyft/flytestdlib/contextutils"
+	"github.com/lyft/flytestdlib/promutils"
+	"github.com/prometheus/client_golang/prometheus"
+)
+
+// Represents a gauge labeled with values from the context. See labeled.SetMetricsKeys for more information
+type Gauge struct {
+	*prometheus.GaugeVec
+
+	prometheus.Gauge
+	additionalLabels []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 to configure that.
+func (g Gauge) Inc(ctx context.Context) {
+	gauge, err := g.GaugeVec.GetMetricWith(contextutils.Values(ctx, metricKeys...))
+	if err != nil {
+		panic(err.Error())
+	}
+	gauge.Inc()
+
+	if g.Gauge != nil {
+		g.Gauge.Inc()
+	}
+}
+
+// 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 to configure that.
+func (g Gauge) Add(ctx context.Context, v float64) {
+	gauge, err := g.GaugeVec.GetMetricWith(contextutils.Values(ctx, metricKeys...))
+	if err != nil {
+		panic(err.Error())
+	}
+	gauge.Add(v)
+
+	if g.Gauge != nil {
+		g.Gauge.Add(v)
+	}
+}
+
+// Set sets the Gauge to an arbitrary value.
+// The data point will be labeled with values from context. See labeled.SetMetricsKeys for information about to configure that.
+func (g Gauge) Set(ctx context.Context, v float64) {
+	gauge, err := g.GaugeVec.GetMetricWith(contextutils.Values(ctx, metricKeys...))
+	if err != nil {
+		panic(err.Error())
+	}
+	gauge.Set(v)
+
+	if g.Gauge != nil {
+		g.Gauge.Set(v)
+	}
+}
+
+// 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 to configure that.
+func (g Gauge) Dec(ctx context.Context) {
+	gauge, err := g.GaugeVec.GetMetricWith(contextutils.Values(ctx, metricKeys...))
+	if err != nil {
+		panic(err.Error())
+	}
+	gauge.Dec()
+
+	if g.Gauge != nil {
+		g.Gauge.Dec()
+	}
+}
+
+// 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 to configure that.
+func (g Gauge) Sub(ctx context.Context, v float64) {
+	gauge, err := g.GaugeVec.GetMetricWith(contextutils.Values(ctx, metricKeys...))
+	if err != nil {
+		panic(err.Error())
+	}
+	gauge.Sub(v)
+
+	if g.Gauge != nil {
+		g.Gauge.Sub(v)
+	}
+}
+
+// SetToCurrentTime sets the Gauge to the current Unix time in seconds.
+func (g Gauge) SetToCurrentTime(ctx context.Context) {
+	gauge, err := g.GaugeVec.GetMetricWith(contextutils.Values(ctx, metricKeys...))
+	if err != nil {
+		panic(err.Error())
+	}
+	gauge.SetToCurrentTime()
+
+	if g.Gauge != nil {
+		g.Gauge.SetToCurrentTime()
+	}
+}
+
+// Creates a new labeled gauge. Label keys must be set before instantiating. If the unlabeled option is given,
+// this object will also instantiate and emit another gauge with the given name with an _unlabeled suffix.
+// See labeled.SetMetricsKeys for information about to configure that.
+func NewGauge(name, description string, scope promutils.Scope, opts ...MetricOption) Gauge {
+	if len(metricKeys) == 0 {
+		panic(ErrNeverSet)
+	}
+
+	g := Gauge{}
+	for _, opt := range opts {
+		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)
+		}
+	}
+
+	if g.GaugeVec == nil {
+		g.GaugeVec = scope.MustNewGaugeVec(name, description, metricStringKeys...)
+	}
+
+	return g
+}

promutils/labeled/gauge_test.go

@@ -0,0 +1,67 @@
+package labeled
+
+import (
+	"context"
+	"strings"
+	"testing"
+
+	"github.com/lyft/flytestdlib/contextutils"
+	"github.com/lyft/flytestdlib/promutils"
+	"github.com/prometheus/client_golang/prometheus/testutil"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestLabeledGauge(t *testing.T) {
+	UnsetMetricKeys()
+	assert.NotPanics(t, func() {
+		SetMetricKeys(contextutils.ProjectKey, contextutils.DomainKey, contextutils.WorkflowIDKey, contextutils.TaskIDKey, contextutils.LaunchPlanIDKey)
+	})
+
+	scope := promutils.NewScope("testscope")
+	ctx := context.Background()
+	ctx = contextutils.WithProjectDomain(ctx, "flyte", "dev")
+	g := NewGauge("unittest", "some desc", scope)
+	assert.NotNil(t, g)
+
+	g.Inc(ctx)
+
+	const header = `
+		# HELP testscope:unittest some desc
+        # TYPE testscope:unittest gauge
+	`
+	var expected = `
+        testscope:unittest{domain="dev",lp="",project="flyte",task="",wf=""} 1
+	`
+	err := testutil.CollectAndCompare(g.GaugeVec, strings.NewReader(header+expected))
+	assert.NoError(t, err)
+
+	g.Set(ctx, 42)
+	expected = `
+        testscope:unittest{domain="dev",lp="",project="flyte",task="",wf=""} 42
+	`
+	err = testutil.CollectAndCompare(g.GaugeVec, strings.NewReader(header+expected))
+	assert.NoError(t, err)
+
+	g.Add(ctx, 1)
+	expected = `
+        testscope:unittest{domain="dev",lp="",project="flyte",task="",wf=""} 43
+	`
+	err = testutil.CollectAndCompare(g.GaugeVec, strings.NewReader(header+expected))
+	assert.NoError(t, err)
+
+	g.Dec(ctx)
+	expected = `
+        testscope:unittest{domain="dev",lp="",project="flyte",task="",wf=""} 42
+	`
+	err = testutil.CollectAndCompare(g.GaugeVec, strings.NewReader(header+expected))
+	assert.NoError(t, err)
+
+	g.Sub(ctx, 1)
+	expected = `
+        testscope:unittest{domain="dev",lp="",project="flyte",task="",wf=""} 41
+	`
+	err = testutil.CollectAndCompare(g.GaugeVec, strings.NewReader(header+expected))
+	assert.NoError(t, err)
+
+	g.SetToCurrentTime(ctx)
+}