diff --git a/promutils/labeled/gauge.go b/promutils/labeled/gauge.go new file mode 100644 index 0000000000..d6c3b1ca18 --- /dev/null +++ b/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 +} diff --git a/promutils/labeled/gauge_test.go b/promutils/labeled/gauge_test.go new file mode 100644 index 0000000000..7c7ff588d5 --- /dev/null +++ b/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) +}