diff --git a/opts.go b/opts.go index f7e710e..5554b34 100644 --- a/opts.go +++ b/opts.go @@ -3,6 +3,7 @@ package promqlsmith import ( "time" + "github.com/prometheus/prometheus/model/labels" "github.com/prometheus/prometheus/promql/parser" "golang.org/x/exp/slices" ) @@ -81,6 +82,8 @@ type options struct { enableAtModifier bool enableVectorMatching bool atModifierMaxTimestamp int64 + + enforceLabelMatchers []*labels.Matcher } func (o *options) applyDefaults() { @@ -163,3 +166,9 @@ func WithEnabledExprs(enabledExprs []ExprType) Option { o.enabledExprs = enabledExprs }) } + +func WithEnforceLabelMatchers(matchers []*labels.Matcher) Option { + return optionFunc(func(o *options) { + o.enforceLabelMatchers = matchers + }) +} diff --git a/promqlsmith.go b/promqlsmith.go index d6c9df2..436c3a5 100644 --- a/promqlsmith.go +++ b/promqlsmith.go @@ -45,8 +45,10 @@ type PromQLSmith struct { enableVectorMatching bool atModifierMaxTimestamp int64 - seriesSet []labels.Labels - labelNames []string + seriesSet []labels.Labels + labelNames []string + labelValues map[string][]string + enforceMatchers []*labels.Matcher supportedExprs []ExprType supportedAggrs []parser.ItemType @@ -65,7 +67,6 @@ func New(rnd *rand.Rand, seriesSet []labels.Labels, opts ...Option) *PromQLSmith ps := &PromQLSmith{ rnd: rnd, seriesSet: filterEmptySeries(seriesSet), - labelNames: labelNamesFromLabelSet(seriesSet), supportedExprs: options.enabledExprs, supportedAggrs: options.enabledAggrs, supportedBinops: options.enabledBinops, @@ -74,7 +75,9 @@ func New(rnd *rand.Rand, seriesSet []labels.Labels, opts ...Option) *PromQLSmith enableAtModifier: options.enableAtModifier, atModifierMaxTimestamp: options.atModifierMaxTimestamp, enableVectorMatching: options.enableVectorMatching, + enforceMatchers: options.enforceLabelMatchers, } + ps.labelNames, ps.labelValues = labelNameAndValuesFromLabelSet(seriesSet) return ps } @@ -89,6 +92,11 @@ func (s *PromQLSmith) WalkRangeQuery() parser.Expr { return s.Walk(vectorAndScalarValueTypes...) } +// WalkSelectors generates random label matchers based on the input series labels. +func (s *PromQLSmith) WalkSelectors() []*labels.Matcher { + return s.walkSelectors() +} + // Walk will walk the ast tree using one of the randomly generated expr type. func (s *PromQLSmith) Walk(valueTypes ...parser.ValueType) parser.Expr { supportedExprs := s.supportedExprs @@ -111,16 +119,24 @@ func filterEmptySeries(seriesSet []labels.Labels) []labels.Labels { return output } -func labelNamesFromLabelSet(labelSet []labels.Labels) []string { - s := make(map[string]struct{}) +func labelNameAndValuesFromLabelSet(labelSet []labels.Labels) ([]string, map[string][]string) { + labelValueSet := make(map[string]map[string]struct{}) for _, lbls := range labelSet { lbls.Range(func(lbl labels.Label) { - s[lbl.Name] = struct{}{} + if _, ok := labelValueSet[lbl.Name]; !ok { + labelValueSet[lbl.Name] = make(map[string]struct{}) + } + labelValueSet[lbl.Name][lbl.Value] = struct{}{} }) } - output := make([]string, 0, len(s)) - for name := range s { - output = append(output, name) + labelNames := make([]string, 0, len(labelValueSet)) + labelValues := make(map[string][]string) + for name, values := range labelValueSet { + labelNames = append(labelNames, name) + labelValues[name] = make([]string, 0, len(values)) + for val := range values { + labelValues[name] = append(labelValues[name], val) + } } - return output + return labelNames, labelValues } diff --git a/promqlsmith_test.go b/promqlsmith_test.go index 0231979..f261acb 100644 --- a/promqlsmith_test.go +++ b/promqlsmith_test.go @@ -20,24 +20,68 @@ var ( labels.MetricName: "http_requests_total", "job": "prometheus", "status_code": "200", + "cluster": "us-west-2", + "env": "prod", }), labels.FromMap(map[string]string{ labels.MetricName: "http_requests_total", "job": "prometheus", "status_code": "404", + "cluster": "us-west-2", + "env": "prod", }), labels.FromMap(map[string]string{ labels.MetricName: "http_requests_total", "job": "prometheus", "status_code": "500", + "cluster": "us-west-2", + "env": "prod", }), labels.FromMap(map[string]string{ labels.MetricName: "up", "job": "prometheus", + "cluster": "us-west-2", + "env": "prod", }), labels.FromMap(map[string]string{ labels.MetricName: "up", "job": "node_exporter", + "cluster": "us-west-2", + "env": "prod", + }), + + labels.FromMap(map[string]string{ + labels.MetricName: "http_requests_total", + "job": "prometheus", + "status_code": "200", + "cluster": "us-east-1", + "env": "prod", + }), + labels.FromMap(map[string]string{ + labels.MetricName: "http_requests_total", + "job": "prometheus", + "status_code": "404", + "cluster": "us-east-1", + "env": "prod", + }), + labels.FromMap(map[string]string{ + labels.MetricName: "http_requests_total", + "job": "prometheus", + "status_code": "500", + "cluster": "us-east-1", + "env": "prod", + }), + labels.FromMap(map[string]string{ + labels.MetricName: "up", + "job": "prometheus", + "cluster": "us-east-1", + "env": "prod", + }), + labels.FromMap(map[string]string{ + labels.MetricName: "up", + "job": "node_exporter", + "cluster": "us-east-1", + "env": "prod", }), } ) @@ -84,6 +128,28 @@ func TestWalk(t *testing.T) { require.NoError(t, err) } +func TestWalkSelectors(t *testing.T) { + rnd := rand.New(rand.NewSource(time.Now().UnixNano())) + ps := New(rnd, testSeriesSet) + matchers := ps.WalkSelectors() + minLen := (len(ps.labelNames) + 1) / 2 + require.True(t, len(matchers) >= minLen) + + enforcedMatcher := labels.MustNewMatcher(labels.MatchEqual, "test", "aaa") + opts := []Option{WithEnforceLabelMatchers([]*labels.Matcher{enforcedMatcher})} + psWithEnforceMatchers := New(rnd, testSeriesSet, opts...) + matchers = psWithEnforceMatchers.WalkSelectors() + minLen = (len(ps.labelNames) + 1) / 2 + require.True(t, len(matchers) >= minLen) + var found bool + for _, matcher := range matchers { + if matcher == enforcedMatcher { + found = true + } + } + require.True(t, found) +} + func TestFilterEmptySeries(t *testing.T) { for i, tc := range []struct { ss []labels.Labels diff --git a/walk.go b/walk.go index 4a2db06..71c707a 100644 --- a/walk.go +++ b/walk.go @@ -328,6 +328,117 @@ func (s *PromQLSmith) walkLabelMatchers() []*labels.Matcher { matchers = append(matchers, labels.MustNewMatcher(labels.MatchEqual, labels.MetricName, metricName)) } } + matchers = append(matchers, s.enforceMatchers...) + + return matchers +} + +// walkSelectors is similar to walkLabelMatchers, but used for generating various +// types of matchers more than simple equal matcher. +func (s *PromQLSmith) walkSelectors() []*labels.Matcher { + if len(s.seriesSet) == 0 { + return nil + } + orders := s.rnd.Perm(len(s.labelNames)) + items := randRange((len(s.labelNames)+1)/2, len(s.labelNames)) + matchers := make([]*labels.Matcher, 0, items) + + var ( + value string + repeat bool + ) + for i := 0; i < items; { + res := s.rnd.Intn(4) + name := s.labelNames[orders[i]] + matchType := labels.MatchType(res) + switch matchType { + case labels.MatchEqual: + val := s.rnd.Float64() + if val > 0.9 { + value = "" + } else if val > 0.8 { + value = "not_exist_value" + } else { + idx := s.rnd.Intn(len(s.labelValues[name])) + value = s.labelValues[name][idx] + } + case labels.MatchNotEqual: + switch s.rnd.Intn(3) { + case 0: + value = "" + case 1: + value = "not_exist_value" + default: + idx := s.rnd.Intn(len(s.labelValues[name])) + value = s.labelValues[name][idx] + } + case labels.MatchRegexp: + val := s.rnd.Float64() + if val > 0.95 { + value = "" + } else if val > 0.9 { + value = "not_exist_value" + } else if val > 0.8 { + value = ".*" + } else if val > 0.7 { + value = ".+" + } else if val > 0.5 { + // Prefix + idx := s.rnd.Intn(len(s.labelValues[name])) + value = s.labelValues[name][idx][:len(s.labelValues[name][idx])/2] + ".*" + } else { + valueOrders := s.rnd.Perm(len(s.labelValues[name])) + valueItems := s.rnd.Intn(len(s.labelValues[name])) + var sb strings.Builder + for j := 0; j < valueItems; j++ { + sb.WriteString(s.labelValues[name][valueOrders[j]]) + if j < valueItems-1 { + sb.WriteString("|") + } + } + // Randomly attach a non-existent value. + if s.rnd.Intn(2) == 1 { + sb.WriteString("|not_exist_value") + } + } + case labels.MatchNotRegexp: + val := s.rnd.Float64() + if val > 0.8 { + value = "" + } else if val > 0.6 { + value = "not_exist_value" + } else if val > 0.4 { + // Prefix + idx := s.rnd.Intn(len(s.labelValues[name])) + value = s.labelValues[name][idx][:len(s.labelValues[name][idx])/2] + ".*" + } else { + valueOrders := s.rnd.Perm(len(s.labelValues[name])) + valueItems := s.rnd.Intn(len(s.labelValues[name])) + var sb strings.Builder + for j := 0; j < valueItems; j++ { + sb.WriteString(s.labelValues[name][valueOrders[j]]) + if j < valueItems-1 { + sb.WriteString("|") + } + } + // Randomly attach a non-existent value. + if s.rnd.Intn(2) == 1 { + sb.WriteString("|not_exist_value") + } + } + default: + panic("unsupported label matcher type") + } + matchers = append(matchers, labels.MustNewMatcher(matchType, name, value)) + + if !repeat && s.rnd.Intn(3) == 0 { + repeat = true + } else { + i++ + } + } + matchers = append(matchers, s.enforceMatchers...) + return matchers } @@ -540,3 +651,7 @@ func getOutputSeries(expr parser.Expr) ([]labels.Labels, bool) { } return lbls, stop } + +func randRange(min, max int) int { + return rand.Intn(max-min) + min +}