redo

Author: yulong-db

go.mod

@@ -140,6 +140,7 @@ require (
 	github.com/elastic/go-licenser v0.3.1 // indirect
 	github.com/go-ini/ini v1.67.0 // indirect
 	github.com/go-openapi/runtime v0.27.1 // indirect
+	github.com/gobwas/glob v0.2.3 // indirect
 	github.com/goccy/go-json v0.10.3 // indirect
 	github.com/godbus/dbus/v5 v5.0.4 // indirect
 	github.com/golang-jwt/jwt/v5 v5.2.1 // indirect

go.sum

@@ -1670,6 +1670,8 @@ 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

@@ -0,0 +1,125 @@
+package query
+
+import (
+	"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)"
+)
+
+type AggregationLabelRewriter struct {
+	logger  log.Logger
+	metrics *aggregationLabelRewriterMetrics
+
+	enabled           bool
+	desiredLabelValue string
+
+	aggregatedMetricNameGlob glob.Glob
+}
+
+type aggregationLabelRewriterMetrics struct {
+	aggregationLabelRewriteSkippedCount    *prometheus.CounterVec
+	aggregationLabelRewrittenCount         *prometheus.CounterVec
+	aggregationLabelAddedCount             prometheus.Counter
+	aggregationLabelRewriterRuntimeSeconds *prometheus.HistogramVec
+}
+
+func newAggregationLabelRewriterMetrics(reg prometheus.Registerer, desiredLabelValue string) *aggregationLabelRewriterMetrics {
+	m := &aggregationLabelRewriterMetrics{}
+
+	m.aggregationLabelRewriteSkippedCount = promauto.With(reg).NewCounterVec(prometheus.CounterOpts{
+		Name: "skipped_total",
+		Help: "Total number of times aggregation-label-rewriter was enabled but skipped",
+	}, []string{"reason"})
+	m.aggregationLabelRewrittenCount = promauto.With(reg).NewCounterVec(prometheus.CounterOpts{
+		Name:        "label_rewritten_total",
+		Help:        "Total number of times aggregation-label-rewriter rewrote the aggregation label",
+		ConstLabels: prometheus.Labels{"new_value": desiredLabelValue},
+	}, []string{"old_value"})
+	m.aggregationLabelAddedCount = promauto.With(reg).NewCounter(prometheus.CounterOpts{
+		Name:        "label_added_total",
+		Help:        "Total number of times aggregation-label-rewriter added the aggregation label",
+		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",
+		Buckets: prometheus.DefBuckets,
+	}, []string{"is_rewritten"})
+
+	return m
+}
+
+func NewAggregationLabelRewriter(logger log.Logger, reg prometheus.Registerer, desiredLabelValue string) *AggregationLabelRewriter {
+	g := glob.MustCompile(aggregatedMetricNameGlobPattern)
+	return &AggregationLabelRewriter{
+		enabled:                  desiredLabelValue != "",
+		logger:                   logger,
+		metrics:                  newAggregationLabelRewriterMetrics(reg, desiredLabelValue),
+		desiredLabelValue:        desiredLabelValue,
+		aggregatedMetricNameGlob: g,
+	}
+}
+
+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
+		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"
+					}
+				} 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 needsRewrite && aggregationLabelMatcher != nil {
+				break
+			} else if m.Name == aggregationLabelName {
+				aggregationLabelMatcher = m
+			}
+		}
+		if needsRewrite {
+			if aggregationLabelMatcher != nil {
+				a.metrics.aggregationLabelRewrittenCount.WithLabelValues(aggregationLabelMatcher.Value).Inc()
+				aggregationLabelMatcher.Type = labels.MatchEqual
+				aggregationLabelMatcher.Value = a.desiredLabelValue
+			} else {
+				a.metrics.aggregationLabelAddedCount.Inc()
+				newMatcher := &labels.Matcher{
+					Name:  aggregationLabelName,
+					Type:  labels.MatchEqual,
+					Value: a.desiredLabelValue,
+				}
+				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
+}

pkg/query/querier.go

@@ -104,6 +104,7 @@ func NewQueryableCreatorWithOptions(
 	) storage.Queryable {
 		return &queryable{
 			logger:              logger,
+			reg:                 reg,
 			replicaLabels:       replicaLabels,
 			storeDebugMatchers:  storeDebugMatchers,
 			proxy:               proxy,
@@ -125,6 +126,7 @@ func NewQueryableCreatorWithOptions(
 
 type queryable struct {
 	logger               log.Logger
+	reg                  prometheus.Registerer
 	replicaLabels        []string
 	storeDebugMatchers   [][]*labels.Matcher
 	proxy                storepb.StoreServer
@@ -142,11 +144,12 @@ type queryable struct {
 
 // Querier returns a new storage querier against the underlying proxy store API.
 func (q *queryable) Querier(mint, maxt int64) (storage.Querier, error) {
-	return newQuerierWithOpts(q.logger, mint, maxt, q.replicaLabels, q.storeDebugMatchers, q.proxy, q.deduplicate, q.maxResolutionMillis, q.partialResponse, q.skipChunks, q.gateProviderFn(), q.selectTimeout, q.shardInfo, q.seriesStatsReporter, q.opts), nil
+	return newQuerierWithOpts(q.logger, q.reg, mint, maxt, q.replicaLabels, q.storeDebugMatchers, q.proxy, q.deduplicate, q.maxResolutionMillis, q.partialResponse, q.skipChunks, q.gateProviderFn(), q.selectTimeout, q.shardInfo, q.seriesStatsReporter, q.opts), nil
 }
 
 type querier struct {
 	logger                  log.Logger
+	reg                     prometheus.Registerer
 	mint, maxt              int64
 	deduplicationFunc       string
 	replicaLabels           []string
@@ -161,17 +164,15 @@ type querier struct {
 	shardInfo               *storepb.ShardInfo
 	seriesStatsReporter     seriesStatsReporter
 
-	enablePreAggregationLabelRewrite bool
-	aggregationStandardSuffixes      []string
-	aggregationLabelName             string
-	rewriteAggregationLabelTo        string
+	aggregationLabelRewriter *AggregationLabelRewriter
 }
 
 // newQuerier creates implementation of storage.Querier that fetches data from the proxy
 // store API endpoints.
 // nolint:unparam
 func newQuerier(
 	logger log.Logger,
+	reg prometheus.Registerer,
 	mint,
 	maxt int64,
 	replicaLabels []string,
@@ -186,11 +187,12 @@ func newQuerier(
 	shardInfo *storepb.ShardInfo,
 	seriesStatsReporter seriesStatsReporter,
 ) *querier {
-	return newQuerierWithOpts(logger, mint, maxt, replicaLabels, storeDebugMatchers, proxy, deduplicate, maxResolutionMillis, partialResponse, skipChunks, selectGate, selectTimeout, shardInfo, seriesStatsReporter, Options{})
+	return newQuerierWithOpts(logger, reg, mint, maxt, replicaLabels, storeDebugMatchers, proxy, deduplicate, maxResolutionMillis, partialResponse, skipChunks, selectGate, selectTimeout, shardInfo, seriesStatsReporter, Options{})
 }
 
 func newQuerierWithOpts(
 	logger log.Logger,
+	reg prometheus.Registerer,
 	mint,
 	maxt int64,
 	replicaLabels []string,
@@ -223,6 +225,7 @@ func newQuerierWithOpts(
 	}
 	return &querier{
 		logger:        logger,
+		reg:           reg,
 		selectGate:    selectGate,
 		selectTimeout: selectTimeout,
 
@@ -239,10 +242,11 @@ func newQuerierWithOpts(
 		shardInfo:               shardInfo,
 		seriesStatsReporter:     seriesStatsReporter,
 
-		enablePreAggregationLabelRewrite: opts.RewriteAggregationLabelTo != "",
-		aggregationStandardSuffixes:      []string{":aggr", ":count", ":sum", ":min", ":max", ":avg"},
-		aggregationLabelName:             "__rollup__",
-		rewriteAggregationLabelTo:        opts.RewriteAggregationLabelTo,
+		aggregationLabelRewriter: NewAggregationLabelRewriter(
+			logger,
+			extprom.WrapRegistererWithPrefix("aggregation_label_rewriter_", reg),
+			opts.RewriteAggregationLabelTo,
+		),
 	}
 }
 
@@ -376,51 +380,7 @@ func (q *querier) Select(ctx context.Context, _ bool, hints *storage.SelectHints
 }
 
 func (q *querier) selectFn(ctx context.Context, hints *storage.SelectHints, ms ...*labels.Matcher) (storage.SeriesSet, storepb.SeriesStatsCounter, error) {
-	if q.enablePreAggregationLabelRewrite {
-		hasNameLabel, needLabelRewrite := false, false
-		var nameLabelMatcher *labels.Matcher
-		for _, m := range ms {
-			if m.Name == labels.MetricName {
-				hasNameLabel = true
-				if m.Type == labels.MatchEqual {
-					for _, suffix := range q.aggregationStandardSuffixes {
-						if strings.HasSuffix(m.Value, suffix) {
-							needLabelRewrite = true
-							nameLabelMatcher = m
-							break
-						}
-					}
-					break
-				} else {
-					level.Info(q.logger).Log("msg", "Skipping aggregation label rewrite attempt", "reason", "non-equal name matcher", "name_matcher", m)
-				}
-			}
-		}
-		if !hasNameLabel {
-			level.Info(q.logger).Log("msg", "Skipping aggregation label rewrite", "reason", "no metric name matcher found")
-		}
-		if needLabelRewrite {
-			hasAggregationLabel := false
-			for i, m := range ms {
-				if m.Name == q.aggregationLabelName {
-					level.Info(q.logger).Log("msg", "Rewriting aggregation label", "original_matcher", m, "name_matcher", nameLabelMatcher, "new_value", q.rewriteAggregationLabelTo)
-					hasAggregationLabel = true
-					ms[i].Type = labels.MatchEqual
-					ms[i].Value = q.rewriteAggregationLabelTo
-					break
-				}
-			}
-			if !hasAggregationLabel {
-				newMatcher := &labels.Matcher{
-					Name:  q.aggregationLabelName,
-					Type:  labels.MatchEqual,
-					Value: q.rewriteAggregationLabelTo,
-				}
-				level.Info(q.logger).Log("msg", "Adding aggregation label", "new_matcher", newMatcher, "name_matcher", nameLabelMatcher)
-				ms = append(ms, newMatcher)
-			}
-		}
-	}
+	ms = q.aggregationLabelRewriter.Rewrite(ms)
 
 	sms, err := storepb.PromMatchersToMatchers(ms...)
 	if err != nil {