thanos querier aggregation label rewrite design v2 (#154)

Author: yulong-db

cmd/thanos/query.go

@@ -245,7 +245,8 @@ func registerQuery(app *extkingpin.App) {
 	enforceTenancy := cmd.Flag("query.enforce-tenancy", "Enforce tenancy on Query APIs. Responses are returned only if the label value of the configured tenant-label-name and the value of the tenant header matches.").Default("false").Bool()
 	tenantLabel := cmd.Flag("query.tenant-label-name", "Label name to use when enforcing tenancy (if --query.enforce-tenancy is enabled).").Default(tenancy.DefaultTenantLabel).String()
 
-	rewriteAggregationLabelTo := cmd.Flag("query.aggregation-label-value-override", "The value override for __rollup__ label for aggregated metrics. If set to x, all queries on aggregated metrics will have a __rollup__=x matcher. Leave empty to disable this behavior. Default is empty.").Default("").String()
+	rewriteAggregationLabelStrategy := cmd.Flag("query.aggregation-label-strategy", "The strategy to use when rewriting aggregation labels. Used during aggregator migration only.").Default(string(query.NoopLabelRewriter)).Hidden().Enum(string(query.NoopLabelRewriter), string(query.UpsertLabelRewriter), string(query.InsertOnlyLabelRewriter))
+	rewriteAggregationLabelTo := cmd.Flag("query.aggregation-label-value-override", "The value override for aggregation label. If set to x, all queries on aggregated metrics will have a `__agg_rule_type__=x` matcher. If empty, this behavior is disabled. Default is empty.").Hidden().Default("").String()
 
 	var storeRateLimits store.SeriesSelectLimits
 	storeRateLimits.RegisterFlags(cmd)
@@ -386,6 +387,7 @@ func registerQuery(app *extkingpin.App) {
 			*enforceTenancy,
 			*tenantLabel,
 			*enableGroupReplicaPartialStrategy,
+			*rewriteAggregationLabelStrategy,
 			*rewriteAggregationLabelTo,
 		)
 	})
@@ -471,6 +473,7 @@ func runQuery(
 	enforceTenancy bool,
 	tenantLabel string,
 	groupReplicaPartialResponseStrategy bool,
+	rewriteAggregationLabelStrategy string,
 	rewriteAggregationLabelTo string,
 ) error {
 	comp := component.Query
@@ -601,6 +604,7 @@ func runQuery(
 	opts := query.Options{
 		GroupReplicaPartialResponseStrategy: groupReplicaPartialResponseStrategy,
 		DeduplicationFunc:                   queryDeduplicationFunc,
+		RewriteAggregationLabelStrategy:     rewriteAggregationLabelStrategy,
 		RewriteAggregationLabelTo:           rewriteAggregationLabelTo,
 	}
 	level.Info(logger).Log("msg", "databricks querier features", "opts", fmt.Sprintf("%+v", opts))

pkg/query/aggregation_label_rewriter.go

@@ -14,15 +14,26 @@ import (
 	"github.com/prometheus/prometheus/model/labels"
 )
 
+// RewriterStrategy defines the strategy used by the AggregationLabelRewriter.
+type RewriterStrategy string
+
 const (
-	aggregationLabelName = "__rollup__"
+	// NoopLabelRewriter is a no-op strategy that basically disables the rewriter.
+	NoopLabelRewriter RewriterStrategy = "noop"
+	// UpsertLabelRewriter is a strategy that upserts the aggregation label.
+	UpsertLabelRewriter RewriterStrategy = "upsert"
+	// InsertOnlyLabelRewriter is a strategy that only inserts the aggregation label if it does not exist.
+	InsertOnlyLabelRewriter RewriterStrategy = "insert-only"
 )
 
-type AggregationLabelRewriter struct {
-	logger  log.Logger
-	metrics *aggregationLabelRewriterMetrics
+const (
+	aggregationLabelName = "__agg_rule_type__"
+)
 
-	enabled           bool
+type AggregationLabelRewriter struct {
+	logger            log.Logger
+	metrics           *aggregationLabelRewriterMetrics
+	strategy          RewriterStrategy
 	desiredLabelValue string
 }
 
@@ -70,24 +81,27 @@ func newAggregationLabelRewriterMetrics(reg prometheus.Registerer, desiredLabelV
 
 func NewNopAggregationLabelRewriter() *AggregationLabelRewriter {
 	return &AggregationLabelRewriter{
-		enabled: false,
+		strategy: NoopLabelRewriter,
 	}
 }
 
-func NewAggregationLabelRewriter(logger log.Logger, reg prometheus.Registerer, desiredLabelValue string) *AggregationLabelRewriter {
+func NewAggregationLabelRewriter(logger log.Logger, reg prometheus.Registerer, strategy RewriterStrategy, desiredLabelValue string) *AggregationLabelRewriter {
 	if logger == nil {
 		logger = log.NewNopLogger()
 	}
+	if desiredLabelValue == "" {
+		strategy = NoopLabelRewriter
+	}
 	return &AggregationLabelRewriter{
-		enabled:           desiredLabelValue != "",
+		strategy:          strategy,
 		logger:            logger,
 		metrics:           newAggregationLabelRewriterMetrics(reg, desiredLabelValue),
 		desiredLabelValue: desiredLabelValue,
 	}
 }
 
 func (a *AggregationLabelRewriter) Rewrite(ms []*labels.Matcher) []*labels.Matcher {
-	if !a.enabled {
+	if a.strategy == NoopLabelRewriter {
 		return ms
 	}
 
@@ -123,12 +137,18 @@ func (a *AggregationLabelRewriter) Rewrite(ms []*labels.Matcher) []*labels.Match
 			aggregationLabelIndex = i
 		}
 	}
+
+	if aggregationLabelMatcher != nil && a.strategy == InsertOnlyLabelRewriter {
+		needsRewrite = false
+		skipReason = "insert-only"
+	}
+
 	// 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,
+			Type:  labels.MatchRegexp,
 			Value: a.desiredLabelValue,
 		}
 		if aggregationLabelMatcher != nil {

pkg/query/aggregation_label_rewriter_test.go

@@ -17,6 +17,7 @@ func TestAggregationLabelRewriter_Rewrite(t *testing.T) {
 	for _, tc := range []struct {
 		name               string
 		desiredLabelValue  string // Empty means disabled
+		strategy           RewriterStrategy
 		inputMatchers      []*labels.Matcher
 		expectedMatchers   []*labels.Matcher
 		expectedSkipCount  float64
@@ -25,7 +26,19 @@ func TestAggregationLabelRewriter_Rewrite(t *testing.T) {
 	}{
 		{
 			name:              "disabled rewriter should not modify label matchers",
+			desiredLabelValue: "v1",
+			strategy:          NoopLabelRewriter,
+			inputMatchers: []*labels.Matcher{
+				labels.MustNewMatcher(labels.MatchEqual, "__name__", "test:sum"),
+			},
+			expectedMatchers: []*labels.Matcher{
+				labels.MustNewMatcher(labels.MatchEqual, "__name__", "test:sum"),
+			},
+		},
+		{
+			name:              "no desired label value makes a disabled rewriter and should not modify label matchers",
 			desiredLabelValue: "",
+			strategy:          UpsertLabelRewriter,
 			inputMatchers: []*labels.Matcher{
 				labels.MustNewMatcher(labels.MatchEqual, "__name__", "test:sum"),
 			},
@@ -36,31 +49,48 @@ func TestAggregationLabelRewriter_Rewrite(t *testing.T) {
 		{
 			name:              "should add label for aggregated metric if no existing aggregation label",
 			desiredLabelValue: "5m",
+			strategy:          UpsertLabelRewriter,
 			inputMatchers: []*labels.Matcher{
 				labels.MustNewMatcher(labels.MatchEqual, "__name__", "test:sum"),
 			},
 			expectedMatchers: []*labels.Matcher{
 				labels.MustNewMatcher(labels.MatchEqual, "__name__", "test:sum"),
-				labels.MustNewMatcher(labels.MatchEqual, "__rollup__", "5m"),
+				labels.MustNewMatcher(labels.MatchRegexp, aggregationLabelName, "5m"),
 			},
 			expectedAddCount: 1,
 		},
 		{
-			name:              "should rewrite existing aggregation label for aggregated metric",
+			name:              "should rewrite existing equal aggregation label for aggregated metric",
+			desiredLabelValue: "5m",
+			strategy:          UpsertLabelRewriter,
+			inputMatchers: []*labels.Matcher{
+				labels.MustNewMatcher(labels.MatchEqual, "__name__", "test:sum"),
+				labels.MustNewMatcher(labels.MatchEqual, aggregationLabelName, "1h"),
+			},
+			expectedMatchers: []*labels.Matcher{
+				labels.MustNewMatcher(labels.MatchEqual, "__name__", "test:sum"),
+				labels.MustNewMatcher(labels.MatchRegexp, aggregationLabelName, "5m"),
+			},
+			expectedRewriteMap: map[string]float64{"1h": 1},
+		},
+		{
+			name:              "should rewrite existing regex aggregation label for aggregated metric",
 			desiredLabelValue: "5m",
+			strategy:          UpsertLabelRewriter,
 			inputMatchers: []*labels.Matcher{
 				labels.MustNewMatcher(labels.MatchEqual, "__name__", "test:sum"),
-				labels.MustNewMatcher(labels.MatchEqual, "__rollup__", "1h"),
+				labels.MustNewMatcher(labels.MatchRegexp, aggregationLabelName, "1h"),
 			},
 			expectedMatchers: []*labels.Matcher{
 				labels.MustNewMatcher(labels.MatchEqual, "__name__", "test:sum"),
-				labels.MustNewMatcher(labels.MatchEqual, "__rollup__", "5m"),
+				labels.MustNewMatcher(labels.MatchRegexp, aggregationLabelName, "5m"),
 			},
 			expectedRewriteMap: map[string]float64{"1h": 1},
 		},
 		{
 			name:              "should skip non-aggregated metric",
 			desiredLabelValue: "5m",
+			strategy:          UpsertLabelRewriter,
 			inputMatchers: []*labels.Matcher{
 				labels.MustNewMatcher(labels.MatchEqual, "__name__", "test_metric"),
 			},
@@ -72,6 +102,7 @@ func TestAggregationLabelRewriter_Rewrite(t *testing.T) {
 		{
 			name:              "should skip non-equal name matcher",
 			desiredLabelValue: "5m",
+			strategy:          UpsertLabelRewriter,
 			inputMatchers: []*labels.Matcher{
 				labels.MustNewMatcher(labels.MatchRegexp, "__name__", "test:sum"),
 			},
@@ -83,6 +114,7 @@ func TestAggregationLabelRewriter_Rewrite(t *testing.T) {
 		{
 			name:              "should skip when no name matcher",
 			desiredLabelValue: "5m",
+			strategy:          UpsertLabelRewriter,
 			inputMatchers: []*labels.Matcher{
 				labels.MustNewMatcher(labels.MatchEqual, "job", "prometheus"),
 			},
@@ -91,12 +123,27 @@ func TestAggregationLabelRewriter_Rewrite(t *testing.T) {
 			},
 			expectedSkipCount: 1,
 		},
+		{
+			name:              "if insert only, should NOT rewrite existing aggregation label for aggregated metric",
+			desiredLabelValue: "5m",
+			strategy:          InsertOnlyLabelRewriter,
+			inputMatchers: []*labels.Matcher{
+				labels.MustNewMatcher(labels.MatchEqual, "__name__", "test:sum"),
+				labels.MustNewMatcher(labels.MatchEqual, aggregationLabelName, "1h"),
+			},
+			expectedMatchers: []*labels.Matcher{
+				labels.MustNewMatcher(labels.MatchEqual, "__name__", "test:sum"),
+				labels.MustNewMatcher(labels.MatchEqual, aggregationLabelName, "1h"),
+			},
+			expectedSkipCount: 1,
+		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
 			reg := prometheus.NewRegistry()
 			rewriter := NewAggregationLabelRewriter(
 				nil,
 				reg,
+				tc.strategy,
 				tc.desiredLabelValue,
 			)
 

pkg/query/querier.go

@@ -61,6 +61,7 @@ type QueryableCreator func(
 type Options struct {
 	GroupReplicaPartialResponseStrategy bool
 	DeduplicationFunc                   string
+	RewriteAggregationLabelStrategy     string
 	RewriteAggregationLabelTo           string
 }
 
@@ -94,6 +95,7 @@ func NewQueryableCreatorWithOptions(
 	aggregationLabelRewriter := NewAggregationLabelRewriter(
 		logger,
 		extprom.WrapRegistererWithPrefix("aggregation_label_rewriter_", reg),
+		RewriterStrategy(opts.RewriteAggregationLabelStrategy),
 		opts.RewriteAggregationLabelTo,
 	)
 	return func(