address comments

Author: yulong-db

go.mod

@@ -114,7 +114,6 @@ require (
 require (
 	capnproto.org/go/capnp/v3 v3.0.1-alpha.2.0.20240830165715-46ccd63a72af
 	github.com/cortexproject/promqlsmith v0.0.0-20240506042652-6cfdd9739a5e
-	github.com/gobwas/glob v0.2.3
 	github.com/grpc-ecosystem/go-grpc-middleware/providers/prometheus v1.0.1
 	github.com/hashicorp/golang-lru v0.6.0
 	github.com/hashicorp/golang-lru/v2 v2.0.7

go.sum

@@ -1670,8 +1670,6 @@ github.com/go-task/slim-sprig/v3 v3.0.0 h1:sUs3vkvUymDpBKi3qH1YSqBQk9+9D/8M2mN1v
 github.com/go-task/slim-sprig/v3 v3.0.0/go.mod h1:W848ghGpv3Qj3dhTPRyJypKRiqCdHZiAzKg9hl15HA8=
 github.com/go-zookeeper/zk v1.0.3 h1:7M2kwOsc//9VeeFiPtf+uSJlVpU66x9Ba5+8XK7/TDg=
 github.com/go-zookeeper/zk v1.0.3/go.mod h1:nOB03cncLtlp4t+UAkGSV+9beXP/akpekBwL+UX1Qcw=
-github.com/gobwas/glob v0.2.3 h1:A4xDbljILXROh+kObIiy5kIaPYD8e96x1tgBhUI5J+Y=
-github.com/gobwas/glob v0.2.3/go.mod h1:d3Ez4x06l9bZtSvzIay5+Yzi0fmZzPgnTbPcKjJAkT8=
 github.com/gobwas/httphead v0.1.0 h1:exrUm0f4YX0L7EBwZHuCF4GDp8aJfVeBrlLQrs6NqWU=
 github.com/gobwas/httphead v0.1.0/go.mod h1:O/RXo79gxV8G+RqlR/otEwx4Q36zl9rqC5u12GKvMCM=
 github.com/gobwas/pool v0.2.1 h1:xfeeEhW7pwmX8nuLVlqbzVc7udMDrwetjEv+TZIz1og=

pkg/query/aggregation_label_rewriter.go

@@ -4,19 +4,18 @@
 package query
 
 import (
+	"strings"
 	"time"
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
-	"github.com/gobwas/glob"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
 	"github.com/prometheus/prometheus/model/labels"
 )
 
 const (
-	aggregationLabelName            = "__rollup__"
-	aggregatedMetricNameGlobPattern = "{*:aggr,*:sum,*:count,*:avg,*:min,*:max}"
+	aggregationLabelName = "__rollup__"
 )
 
 type AggregationLabelRewriter struct {
@@ -25,8 +24,6 @@ type AggregationLabelRewriter struct {
 
 	enabled           bool
 	desiredLabelValue string
-
-	aggregatedMetricNameGlob glob.Glob
 }
 
 type aggregationLabelRewriterMetrics struct {
@@ -36,6 +33,15 @@ type aggregationLabelRewriterMetrics struct {
 	aggregationLabelRewriterRuntimeSeconds *prometheus.HistogramVec
 }
 
+func isAggregatedMetric(metricName string) bool {
+	return strings.HasSuffix(metricName, ":aggr") ||
+		strings.HasSuffix(metricName, ":avg") ||
+		strings.HasSuffix(metricName, ":count") ||
+		strings.HasSuffix(metricName, ":sum") ||
+		strings.HasSuffix(metricName, ":min") ||
+		strings.HasSuffix(metricName, ":max")
+}
+
 func newAggregationLabelRewriterMetrics(reg prometheus.Registerer, desiredLabelValue string) *aggregationLabelRewriterMetrics {
 	m := &aggregationLabelRewriterMetrics{}
 
@@ -54,8 +60,9 @@ func newAggregationLabelRewriterMetrics(reg prometheus.Registerer, desiredLabelV
 		ConstLabels: prometheus.Labels{"new_value": desiredLabelValue},
 	})
 	m.aggregationLabelRewriterRuntimeSeconds = promauto.With(reg).NewHistogramVec(prometheus.HistogramOpts{
-		Name: "total_runtime_seconds",
-		Help: "Runtime of aggregation-label-rewriter in seconds",
+		Name:    "total_runtime_seconds",
+		Help:    "Runtime of aggregation-label-rewriter in seconds",
+		Buckets: []float64{0.0001, 0.00025, 0.0005, 0.001, 0.005, 0.01, 0.1, 1.0, 10.0},
 	}, []string{"is_rewritten"})
 
 	return m
@@ -71,67 +78,75 @@ func NewAggregationLabelRewriter(logger log.Logger, reg prometheus.Registerer, d
 	if logger == nil {
 		logger = log.NewNopLogger()
 	}
-	g := glob.MustCompile(aggregatedMetricNameGlobPattern)
 	return &AggregationLabelRewriter{
-		enabled:                  desiredLabelValue != "",
-		logger:                   logger,
-		metrics:                  newAggregationLabelRewriterMetrics(reg, desiredLabelValue),
-		desiredLabelValue:        desiredLabelValue,
-		aggregatedMetricNameGlob: g,
+		enabled:           desiredLabelValue != "",
+		logger:            logger,
+		metrics:           newAggregationLabelRewriterMetrics(reg, desiredLabelValue),
+		desiredLabelValue: desiredLabelValue,
 	}
 }
 
 func (a *AggregationLabelRewriter) Rewrite(ms []*labels.Matcher) []*labels.Matcher {
-	if a.enabled {
-		startOfRun := time.Now()
-		needsRewrite := false
-		skipReason := "no-name-matcher"
-		var aggregationLabelMatcher *labels.Matcher
-		aggregationLabelIndex := -1
-		for i := 0; i < len(ms); i++ {
-			m := ms[i]
-			if !needsRewrite && m.Name == labels.MetricName {
-				if m.Type == labels.MatchEqual {
-					if a.aggregatedMetricNameGlob.Match(m.Value) {
-						needsRewrite = true
-					} else {
-						skipReason = "not-aggregated-metric"
-					}
+	if !a.enabled {
+		return ms
+	}
+
+	startOfRun := time.Now()
+	needsRewrite := false
+	skipReason := "no-name-matcher"
+	var aggregationLabelMatcher *labels.Matcher
+	aggregationLabelIndex := -1
+	for i := 0; i < len(ms); i++ {
+		// If we already know we need to rewrite, and we found the aggregation label
+		// we don't need to loop further, break
+		if needsRewrite && aggregationLabelMatcher != nil {
+			break
+		}
+		m := ms[i]
+		// If we are not sure if we need rewrite, and see a name matcher
+		// We check if this is an aggregated metric and decide needsRewrite
+		if !needsRewrite && m.Name == labels.MetricName {
+			if m.Type == labels.MatchEqual {
+				if isAggregatedMetric(m.Value) {
+					needsRewrite = true
 				} else {
-					skipReason = "not-using-equal-name-match"
-					level.Debug(a.logger).Log("msg", "skipped due to not using equal in name matcher", "matcher", m)
-					break
+					skipReason = "not-aggregated-metric"
 				}
-			} else if needsRewrite && aggregationLabelMatcher != nil {
+			} else {
+				skipReason = "not-using-equal-name-match"
+				level.Debug(a.logger).Log("msg", "skipped due to not using equal in name matcher", "matcher", m)
 				break
-			} else if m.Name == aggregationLabelName {
-				aggregationLabelMatcher = m
-				aggregationLabelIndex = i
 			}
+			// In any case, if we see an aggregation label, we store that for later use
+		} else if m.Name == aggregationLabelName {
+			aggregationLabelMatcher = m
+			aggregationLabelIndex = i
 		}
-		if needsRewrite {
-			newMatcher := &labels.Matcher{
-				Name:  aggregationLabelName,
-				Type:  labels.MatchEqual,
-				Value: a.desiredLabelValue,
-			}
-			if aggregationLabelMatcher != nil {
-				a.metrics.aggregationLabelRewrittenCount.WithLabelValues(aggregationLabelMatcher.Value).Inc()
-				ms[aggregationLabelIndex] = newMatcher
-			} else {
-				a.metrics.aggregationLabelAddedCount.Inc()
-				ms = append(ms, newMatcher)
-			}
-			a.metrics.aggregationLabelRewriterRuntimeSeconds.WithLabelValues("true").Observe(
-				time.Since(startOfRun).Seconds(),
-			)
+	}
+	// After the for loop, if needsRewrite is false, no need to do anything
+	// but if it is true, we either append or modify an aggregation label
+	if needsRewrite {
+		newMatcher := &labels.Matcher{
+			Name:  aggregationLabelName,
+			Type:  labels.MatchEqual,
+			Value: a.desiredLabelValue,
+		}
+		if aggregationLabelMatcher != nil {
+			a.metrics.aggregationLabelRewrittenCount.WithLabelValues(aggregationLabelMatcher.Value).Inc()
+			ms[aggregationLabelIndex] = newMatcher
 		} else {
-			a.metrics.aggregationLabelRewriteSkippedCount.WithLabelValues(skipReason).Inc()
-			a.metrics.aggregationLabelRewriterRuntimeSeconds.WithLabelValues("false").Observe(
-				time.Since(startOfRun).Seconds(),
-			)
+			a.metrics.aggregationLabelAddedCount.Inc()
+			ms = append(ms, newMatcher)
 		}
-
+		a.metrics.aggregationLabelRewriterRuntimeSeconds.WithLabelValues("true").Observe(
+			time.Since(startOfRun).Seconds(),
+		)
+	} else {
+		a.metrics.aggregationLabelRewriteSkippedCount.WithLabelValues(skipReason).Inc()
+		a.metrics.aggregationLabelRewriterRuntimeSeconds.WithLabelValues("false").Observe(
+			time.Since(startOfRun).Seconds(),
+		)
 	}
+
 	return ms
 }