Implement label gatherer (#3074)

api/metrics/gatherer_test.go

@@ -4,14 +4,15 @@
 package metrics
 
 import (
+	"github.com/prometheus/client_golang/prometheus"
+
 	dto "github.com/prometheus/client_model/go"
 )
 
-var (
-	hello      = "hello"
-	world      = "world"
-	helloWorld = "hello_world"
-)
+var counterOpts = prometheus.CounterOpts{
+	Name: "counter",
+	Help: "help",
+}
 
 type testGatherer struct {
 	mfs []*dto.MetricFamily

api/metrics/label_gatherer.go

@@ -0,0 +1,76 @@
+// Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package metrics
+
+import (
+	"errors"
+	"fmt"
+	"slices"
+
+	"github.com/prometheus/client_golang/prometheus"
+
+	dto "github.com/prometheus/client_model/go"
+)
+
+var (
+	_ MultiGatherer = (*prefixGatherer)(nil)
+
+	errDuplicateGatherer = errors.New("attempt to register duplicate gatherer")
+)
+
+// NewLabelGatherer returns a new MultiGatherer that merges metrics by adding a
+// new label.
+func NewLabelGatherer(labelName string) MultiGatherer {
+	return &labelGatherer{
+		labelName: labelName,
+	}
+}
+
+type labelGatherer struct {
+	multiGatherer
+
+	labelName string
+}
+
+func (g *labelGatherer) Register(labelValue string, gatherer prometheus.Gatherer) error {
+	g.lock.Lock()
+	defer g.lock.Unlock()
+
+	if slices.Contains(g.names, labelValue) {
+		return fmt.Errorf("%w: for %q with label %q",
+			errDuplicateGatherer,
+			g.labelName,
+			labelValue,
+		)
+	}
+
+	g.names = append(g.names, labelValue)
+	g.gatherers = append(g.gatherers, &labeledGatherer{
+		labelName:  g.labelName,
+		labelValue: labelValue,
+		gatherer:   gatherer,
+	})
+	return nil
+}
+
+type labeledGatherer struct {
+	labelName  string
+	labelValue string
+	gatherer   prometheus.Gatherer
+}
+
+func (g *labeledGatherer) Gather() ([]*dto.MetricFamily, error) {
+	// Gather returns partially filled metrics in the case of an error. So, it
+	// is expected to still return the metrics in the case an error is returned.
+	metricFamilies, err := g.gatherer.Gather()
+	for _, metricFamily := range metricFamilies {
+		for _, metric := range metricFamily.Metric {
+			metric.Label = append(metric.Label, &dto.LabelPair{
+				Name:  &g.labelName,
+				Value: &g.labelValue,
+			})
+		}
+	}
+	return metricFamilies, err
+}

api/metrics/label_gatherer_test.go

@@ -0,0 +1,217 @@
+// Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package metrics
+
+import (
+	"testing"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/protobuf/proto"
+
+	dto "github.com/prometheus/client_model/go"
+)
+
+func TestLabelGatherer_Gather(t *testing.T) {
+	const (
+		labelName         = "smith"
+		labelValueA       = "rick"
+		labelValueB       = "morty"
+		customLabelName   = "tag"
+		customLabelValueA = "a"
+		customLabelValueB = "b"
+	)
+	tests := []struct {
+		name            string
+		labelName       string
+		expectedMetrics []*dto.Metric
+		expectErr       bool
+	}{
+		{
+			name:      "no overlap",
+			labelName: customLabelName,
+			expectedMetrics: []*dto.Metric{
+				{
+					Label: []*dto.LabelPair{
+						{
+							Name:  proto.String(labelName),
+							Value: proto.String(labelValueB),
+						},
+						{
+							Name:  proto.String(customLabelName),
+							Value: proto.String(customLabelValueB),
+						},
+					},
+					Counter: &dto.Counter{
+						Value: proto.Float64(1),
+					},
+				},
+				{
+					Label: []*dto.LabelPair{
+						{
+							Name:  proto.String(labelName),
+							Value: proto.String(labelValueA),
+						},
+						{
+							Name:  proto.String(customLabelName),
+							Value: proto.String(customLabelValueA),
+						},
+					},
+					Counter: &dto.Counter{
+						Value: proto.Float64(0),
+					},
+				},
+			},
+			expectErr: false,
+		},
+		{
+			name:      "has overlap",
+			labelName: labelName,
+			expectedMetrics: []*dto.Metric{
+				{
+					Label: []*dto.LabelPair{
+						{
+							Name:  proto.String(labelName),
+							Value: proto.String(labelValueB),
+						},
+						{
+							Name:  proto.String(customLabelName),
+							Value: proto.String(customLabelValueB),
+						},
+					},
+					Counter: &dto.Counter{
+						Value: proto.Float64(1),
+					},
+				},
+			},
+			expectErr: true,
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			require := require.New(t)
+
+			gatherer := NewLabelGatherer(labelName)
+			require.NotNil(gatherer)
+
+			registerA := prometheus.NewRegistry()
+			require.NoError(gatherer.Register(labelValueA, registerA))
+			{
+				counterA := prometheus.NewCounterVec(
+					counterOpts,
+					[]string{test.labelName},
+				)
+				counterA.With(prometheus.Labels{test.labelName: customLabelValueA})
+				require.NoError(registerA.Register(counterA))
+			}
+
+			registerB := prometheus.NewRegistry()
+			require.NoError(gatherer.Register(labelValueB, registerB))
+			{
+				counterB := prometheus.NewCounterVec(
+					counterOpts,
+					[]string{customLabelName},
+				)
+				counterB.With(prometheus.Labels{customLabelName: customLabelValueB}).Inc()
+				require.NoError(registerB.Register(counterB))
+			}
+
+			metrics, err := gatherer.Gather()
+			if test.expectErr {
+				require.Error(err) //nolint:forbidigo // the error is not exported
+			} else {
+				require.NoError(err)
+			}
+			require.Equal(
+				[]*dto.MetricFamily{
+					{
+						Name:   proto.String(counterOpts.Name),
+						Help:   proto.String(counterOpts.Help),
+						Type:   dto.MetricType_COUNTER.Enum(),
+						Metric: test.expectedMetrics,
+					},
+				},
+				metrics,
+			)
+		})
+	}
+}
+
+func TestLabelGatherer_Register(t *testing.T) {
+	firstLabeledGatherer := &labeledGatherer{
+		labelValue: "first",
+		gatherer:   &testGatherer{},
+	}
+	firstLabelGatherer := func() *labelGatherer {
+		return &labelGatherer{
+			multiGatherer: multiGatherer{
+				names: []string{firstLabeledGatherer.labelValue},
+				gatherers: prometheus.Gatherers{
+					firstLabeledGatherer,
+				},
+			},
+		}
+	}
+	secondLabeledGatherer := &labeledGatherer{
+		labelValue: "second",
+		gatherer: &testGatherer{
+			mfs: []*dto.MetricFamily{{}},
+		},
+	}
+	secondLabelGatherer := &labelGatherer{
+		multiGatherer: multiGatherer{
+			names: []string{
+				firstLabeledGatherer.labelValue,
+				secondLabeledGatherer.labelValue,
+			},
+			gatherers: prometheus.Gatherers{
+				firstLabeledGatherer,
+				secondLabeledGatherer,
+			},
+		},
+	}
+
+	tests := []struct {
+		name                  string
+		labelGatherer         *labelGatherer
+		labelValue            string
+		gatherer              prometheus.Gatherer
+		expectedErr           error
+		expectedLabelGatherer *labelGatherer
+	}{
+		{
+			name:                  "first registration",
+			labelGatherer:         &labelGatherer{},
+			labelValue:            "first",
+			gatherer:              firstLabeledGatherer.gatherer,
+			expectedErr:           nil,
+			expectedLabelGatherer: firstLabelGatherer(),
+		},
+		{
+			name:                  "second registration",
+			labelGatherer:         firstLabelGatherer(),
+			labelValue:            "second",
+			gatherer:              secondLabeledGatherer.gatherer,
+			expectedErr:           nil,
+			expectedLabelGatherer: secondLabelGatherer,
+		},
+		{
+			name:                  "conflicts with previous registration",
+			labelGatherer:         firstLabelGatherer(),
+			labelValue:            "first",
+			gatherer:              secondLabeledGatherer.gatherer,
+			expectedErr:           errDuplicateGatherer,
+			expectedLabelGatherer: firstLabelGatherer(),
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			require := require.New(t)
+
+			err := test.labelGatherer.Register(test.labelValue, test.gatherer)
+			require.ErrorIs(err, test.expectedErr)
+			require.Equal(test.expectedLabelGatherer, test.labelGatherer)
+		})
+	}
+}