Merge pull request #32 from zapier/add-aggergate-options

Add aggregate options

aggregate.go

@@ -4,7 +4,9 @@ import (
 	"io"
 	"log"
 	"sort"
+	"strings"
 	"sync"
+	"time"
 
 	"github.com/gin-gonic/gin"
 	dto "github.com/prometheus/client_model/go"
@@ -19,12 +21,59 @@ type metricFamily struct {
 type aggregate struct {
 	familiesLock sync.RWMutex
 	families     map[string]*metricFamily
+	options      aggregateOptions
 }
 
-func newAggregate() *aggregate {
-	return &aggregate{
+type ignoredLabels []string
+
+type aggregateOptions struct {
+	ignoredLabels     ignoredLabels
+	metricTTLDuration *time.Duration
+}
+
+type aggregateOptionsFunc func(a *aggregate)
+
+func AddIgnoredLabels(ignoredLabels ...string) aggregateOptionsFunc {
+	return func(a *aggregate) {
+		a.options.ignoredLabels = ignoredLabels
+	}
+}
+
+func SetTTLMetricTime(duration *time.Duration) aggregateOptionsFunc {
+	return func(a *aggregate) {
+		a.options.metricTTLDuration = duration
+	}
+}
+
+func newAggregate(opts ...aggregateOptionsFunc) *aggregate {
+	a := &aggregate{
 		families: map[string]*metricFamily{},
+		options: aggregateOptions{
+			ignoredLabels: []string{},
+		},
 	}
+
+	for _, opt := range opts {
+		opt(a)
+	}
+
+	a.options.formatOptions()
+
+	return a
+}
+
+func (ao *aggregateOptions) formatOptions() {
+	ao.formatIgnoredLabels()
+}
+
+func (ao *aggregateOptions) formatIgnoredLabels() {
+	if ao.ignoredLabels != nil {
+		for i, v := range ao.ignoredLabels {
+			ao.ignoredLabels[i] = strings.ToLower(v)
+		}
+	}
+
+	sort.Strings(ao.ignoredLabels)
 }
 
 // setFamilyOrGetExistingFamily either sets a new family or returns an existing family
@@ -61,7 +110,7 @@ func (a *aggregate) parseAndMerge(r io.Reader) error {
 	for name, family := range inFamilies {
 		// Sort labels in case source sends them inconsistently
 		for _, m := range family.Metric {
-			sort.Sort(byName(m.Label))
+			a.formatLabels(m)
 		}
 
 		if err := validateFamily(family); err != nil {

aggregate_test.go

@@ -83,11 +83,11 @@ histogram_count 2
 
 	multilabel1 = `# HELP counter A counter
 # TYPE counter counter
-counter{a="a",b="b"} 1
+counter{a="a",b="b", ignore_label="ignore_value"} 1
 `
 	multilabel2 = `# HELP counter A counter
 # TYPE counter counter
-counter{a="a",b="b"} 2
+counter{a="a",b="b", ignore_label="ignore_value"} 2
 `
 	multilabelResult = `# HELP counter A counter
 # TYPE counter counter
@@ -139,26 +139,47 @@ counter{b="b",a="a"} 2
 	reorderedLabelsResult = `# HELP counter A counter
 # TYPE counter counter
 counter{a="a",b="b"} 3
+`
+
+	ignoredLabels1 = `# HELP counter A counter
+# TYPE counter counter
+counter{a="a",b="b", ignore_me="ignored"} 1
+`
+	ignoredLabels2 = `# HELP counter A counter
+# TYPE counter counter
+counter{b="b",a="a", ignore_me="ignored"} 2
+`
+	ignoredLabelsResult = `# HELP counter A counter
+# TYPE counter counter
+counter{a="a",b="b"} 3
 `
 )
 
+func TestNewAggregate(t *testing.T) {
+	_ = newAggregate()
+
+}
+
 func TestAggregate(t *testing.T) {
 	metricMiddleware := newMetricMiddleware(nil)
 	for _, c := range []struct {
-		a, b string
-		want string
-		err1 error
-		err2 error
+		a, b          string
+		want          string
+		ignoredLabels []string
+		err1          error
+		err2          error
 	}{
-		{gaugeInput, gaugeInput, gaugeOutput, nil, nil},
-		{in1, in2, want, nil, nil},
-		{multilabel1, multilabel2, multilabelResult, nil, nil},
-		{labelFields1, labelFields2, labelFieldResult, nil, nil},
-		{duplicateLabels, "", "", fmt.Errorf("%s", duplicateError), nil},
-		{reorderedLabels1, reorderedLabels2, reorderedLabelsResult, nil, nil},
+		{gaugeInput, gaugeInput, gaugeOutput, []string{}, nil, nil},
+		{in1, in2, want, []string{}, nil, nil},
+		{multilabel1, multilabel2, multilabelResult, []string{"ignore_label"}, nil, nil},
+		{labelFields1, labelFields2, labelFieldResult, []string{}, nil, nil},
+		{duplicateLabels, "", "", []string{}, fmt.Errorf("%s", duplicateError), nil},
+		{reorderedLabels1, reorderedLabels2, reorderedLabelsResult, []string{}, nil, nil},
+		{ignoredLabels1, ignoredLabels2, ignoredLabelsResult, []string{"ignore_me"}, nil, nil},
 	} {
 		rc := &RouterConfig{
 			MetricsMiddleware: &metricMiddleware,
+			Aggregate:         newAggregate(AddIgnoredLabels(c.ignoredLabels...)),
 		}
 		router := setupRouter(rc)
 
@@ -191,23 +212,26 @@ func TestAggregate(t *testing.T) {
 	}
 }
 
-var table = []struct {
+var testMetricTable = []struct {
 	inputName      string
 	input1, input2 string
+	ignoredLabels  []string
 }{
-	{"simpleGauge", gaugeInput, gaugeInput},
-	{"fullMetrics", in1, in2},
-	{"multiLabel", multilabel1, multilabel2},
-	{"labelFields", labelFields1, labelFields2},
-	{"reorderedLabels", reorderedLabels1, reorderedLabels2},
+	{"simpleGauge", gaugeInput, gaugeInput, []string{}},
+	{"fullMetrics", in1, in2, []string{}},
+	{"multiLabel", multilabel1, multilabel2, []string{}},
+	{"multiLabelIgnore", multilabel1, multilabel2, []string{"ignore_label"}},
+	{"labelFields", labelFields1, labelFields2, []string{}},
+	{"reorderedLabels", reorderedLabels1, reorderedLabels2, []string{}},
+	{"ignoredLabels", ignoredLabels1, ignoredLabels2, []string{"ignore_me"}},
 }
 
 func BenchmarkAggregate(b *testing.B) {
-	for _, v := range table {
+	a := newAggregate()
+	for _, v := range testMetricTable {
+		a.options.ignoredLabels = v.ignoredLabels
 		b.Run(fmt.Sprintf("metric_type_%s", v.inputName), func(b *testing.B) {
 			for n := 0; n < b.N; n++ {
-				a := newAggregate()
-
 				if err := a.parseAndMerge(strings.NewReader(v.input1)); err != nil {
 					b.Fatalf("unexpected error %s", err)
 				}
@@ -221,7 +245,8 @@ func BenchmarkAggregate(b *testing.B) {
 
 func BenchmarkConcurrentAggregate(b *testing.B) {
 	a := newAggregate()
-	for _, v := range table {
+	for _, v := range testMetricTable {
+		a.options.ignoredLabels = v.ignoredLabels
 		b.Run(fmt.Sprintf("metric_type_%s", v.inputName), func(b *testing.B) {
 			if err := a.parseAndMerge(strings.NewReader(v.input1)); err != nil {
 				b.Fatalf("unexpected error %s", err)

go.mod

@@ -4,17 +4,20 @@ go 1.19
 
 require (
 	github.com/gin-gonic/gin v1.8.1
+	github.com/go-playground/assert/v2 v2.0.1
 	github.com/pmezard/go-difflib v1.0.0
 	github.com/prometheus/client_golang v1.12.1
 	github.com/prometheus/client_model v0.3.0
 	github.com/prometheus/common v0.37.0
 	github.com/slok/go-http-metrics v0.10.0
+	github.com/stretchr/testify v1.8.1
 	golang.org/x/sync v0.1.0
 )
 
 require (
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/cespare/xxhash/v2 v2.1.2 // indirect
+	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/gin-contrib/sse v0.1.0 // indirect
 	github.com/go-playground/locales v0.14.0 // indirect
 	github.com/go-playground/universal-translator v0.18.0 // indirect
@@ -36,4 +39,5 @@ require (
 	golang.org/x/text v0.3.7 // indirect
 	google.golang.org/protobuf v1.28.0 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

labels.go

@@ -0,0 +1,43 @@
+package main
+
+import (
+	"sort"
+	"strings"
+
+	dto "github.com/prometheus/client_model/go"
+)
+
+func strPtr(s string) *string {
+	return &s
+
+}
+
+func (a *aggregate) formatLabels(m *dto.Metric) {
+	sort.Sort(byName(m.Label))
+
+	if len(a.options.ignoredLabels) > 0 {
+		var newLabelList []*dto.LabelPair
+		for _, l := range m.Label {
+			if !a.options.ignoredLabels.labelInIgnoredList(l) {
+				newLabelList = append(newLabelList, l)
+			}
+		}
+		m.Label = newLabelList
+	}
+}
+
+func (iL ignoredLabels) labelInIgnoredList(l *dto.LabelPair) bool {
+	if l == nil || l.Name == nil {
+		return true
+	}
+
+	for _, label := range iL {
+		if l.Name != nil {
+			if strings.ToLower(*l.Name) == label {
+				return true
+			}
+		}
+	}
+
+	return false
+}

labels_test.go

@@ -0,0 +1,90 @@
+package main
+
+import (
+	"fmt"
+	"testing"
+
+	dto "github.com/prometheus/client_model/go"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestFormatLabels(t *testing.T) {
+	a := newAggregate()
+	a.options.ignoredLabels = []string{"ignore_me"}
+
+	m := &dto.Metric{
+		Label: []*dto.LabelPair{
+			{Name: strPtr("thing2"), Value: strPtr("value2")},
+			{Name: strPtr("ignore_me"), Value: strPtr("ignored_value")},
+			{Name: strPtr("thing1"), Value: strPtr("value1")},
+			{},
+		},
+	}
+	a.formatLabels(m)
+
+	assert.Equal(t, &dto.LabelPair{Name: strPtr("thing1"), Value: strPtr("value1")}, m.Label[0])
+	assert.Equal(t, &dto.LabelPair{Name: strPtr("thing2"), Value: strPtr("value2")}, m.Label[1])
+	assert.Len(t, m.Label, 2)
+
+}
+
+var testLabelTable = []struct {
+	inputName     string
+	m             *dto.Metric
+	ignoredLabels []string
+}{
+	{"no_labels", &dto.Metric{Label: []*dto.LabelPair{}}, []string{}},
+	{"no_labels_1_ignored_label", &dto.Metric{Label: []*dto.LabelPair{}},
+		[]string{"ignore_me"}},
+	{"no_ignored_labels", &dto.Metric{Label: []*dto.LabelPair{
+		{Name: strPtr("l1"), Value: strPtr("v1")},
+	}},
+		[]string{},
+	},
+	{"no_ignored_labels_with_3_ignored_labels", &dto.Metric{Label: []*dto.LabelPair{
+		{Name: strPtr("l1"), Value: strPtr("v1")},
+	}},
+		[]string{"ignore_me", "ignore_me_1", "ignore_me_2"},
+	},
+	{"1_ignored_labels_with_1_ignores_set", &dto.Metric{Label: []*dto.LabelPair{
+		{Name: strPtr("l1"), Value: strPtr("v1")},
+		{Name: strPtr("ignore_me"), Value: strPtr("ignore")},
+	}},
+		[]string{"ignore_me"},
+	},
+	{"1_ignored_labels_with_3_ignores_set", &dto.Metric{Label: []*dto.LabelPair{
+		{Name: strPtr("l1"), Value: strPtr("v1")},
+		{Name: strPtr("ignore_me"), Value: strPtr("ignore")},
+	}},
+		[]string{"ignore_me", "ignore_me_1", "ignore_me_2"},
+	},
+	{"2_ignored_labels", &dto.Metric{Label: []*dto.LabelPair{
+		{Name: strPtr("l1"), Value: strPtr("v1")},
+		{Name: strPtr("ignore_me"), Value: strPtr("ignore")},
+		{Name: strPtr("ignore_me_1"), Value: strPtr("ignore1")},
+	}},
+		[]string{"ignore_me", "ignore_me_1", "ignore_me_2"},
+	},
+	{"2_ignored_labels_with_lots_of_labels_ignored_labels", &dto.Metric{Label: []*dto.LabelPair{
+		{Name: strPtr("l1"), Value: strPtr("v1")},
+		{Name: strPtr("ignore_me"), Value: strPtr("ignore")},
+		{Name: strPtr("ignore_me_1"), Value: strPtr("ignore1")},
+		{Name: strPtr("l3"), Value: strPtr("v3")},
+		{Name: strPtr("l2"), Value: strPtr("v2")},
+		{Name: strPtr("l5"), Value: strPtr("v5")},
+		{Name: strPtr("l4"), Value: strPtr("v5")},
+	}},
+		[]string{"ignore_me", "ignore_me_1"},
+	},
+}
+
+func BenchmarkFormatLabels(b *testing.B) {
+	for _, v := range testLabelTable {
+		a := newAggregate(AddIgnoredLabels(v.ignoredLabels...))
+		b.Run(fmt.Sprintf("metric_type_%s", v.inputName), func(b *testing.B) {
+			for n := 0; n < b.N; n++ {
+				a.formatLabels(v.m)
+			}
+		})
+	}
+}