logicalplan: distributed optimizer fuzzing

We should only have remote executions that maintain the partition. Add a
fuzzing test to check that invariant.

Signed-off-by: Michael Hoffmann <mhoffmann@cloudflare.com>

Author: MichaHoffmann

logicalplan/distribute.go

@@ -102,10 +102,18 @@ func (r RemoteExecution) ReturnType() parser.ValueType { return r.Query.ReturnTy
 
 // Deduplicate is a logical plan which deduplicates samples from multiple RemoteExecutions.
 type Deduplicate struct {
-	LeafNode
 	Expressions RemoteExecutions
 }
 
+func (r Deduplicate) Children() []*Node {
+	children := make([]*Node, len(r.Expressions))
+	for i := range r.Expressions {
+		var n Node = r.Expressions[i]
+		children[i] = &n
+	}
+	return children
+}
+
 func (r Deduplicate) Clone() Node {
 	clone := r
 	clone.Expressions = make(RemoteExecutions, len(r.Expressions))

logicalplan/distribute_test.go

@@ -5,13 +5,15 @@ package logicalplan
 
 import (
 	"math"
+	"math/rand"
 	"regexp"
 	"testing"
 	"time"
 
 	"github.com/thanos-io/promql-engine/api"
 	"github.com/thanos-io/promql-engine/query"
 
+	"github.com/cortexproject/promqlsmith"
 	"github.com/efficientgo/core/testutil"
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/prometheus/prometheus/promql/parser"
@@ -1083,3 +1085,78 @@ func newEngineMock(mint, maxt int64, labelSets []labels.Labels) *engineMock {
 func newEngineMockWithExplicitPartition(mint, maxt int64, labelSets, partitionLabelSets []labels.Labels) *engineMock {
 	return &engineMock{minT: mint, maxT: maxt, labelSets: labelSets, partitionLabelSets: partitionLabelSets}
 }
+
+func FuzzDistributedExecutionPreservesPartitionLabels(f *testing.F) {
+	f.Add(int64(0))
+	f.Fuzz(func(t *testing.T, seed int64) {
+		rnd := rand.New(rand.NewSource(seed))
+
+		engines := []api.RemoteEngine{
+			newEngineMock(math.MinInt64, math.MaxInt64, []labels.Labels{labels.FromStrings("region", "east")}),
+			newEngineMock(math.MinInt64, math.MaxInt64, []labels.Labels{labels.FromStrings("region", "west")}),
+		}
+		optimizers := []Optimizer{
+			DistributedExecutionOptimizer{Endpoints: api.NewStaticEndpoints(engines)},
+		}
+		engineLabels := map[string]struct{}{"region": {}}
+
+		lbls := []labels.Labels{
+			labels.FromStrings("__name__", "http_requests_total", "pod", "nginx-1", "region", "east"),
+			labels.FromStrings("__name__", "http_requests_total", "pod", "nginx-2", "region", "west"),
+		}
+
+		psOpts := []promqlsmith.Option{
+			promqlsmith.WithEnableOffset(false),
+			promqlsmith.WithEnableAtModifier(false),
+			promqlsmith.WithEnabledAggrs([]parser.ItemType{
+				parser.SUM, parser.MIN, parser.MAX, parser.AVG, parser.GROUP,
+				parser.COUNT, parser.QUANTILE, parser.STDDEV, parser.STDVAR,
+				parser.COUNT_VALUES, parser.TOPK, parser.BOTTOMK,
+			}),
+			promqlsmith.WithEnableVectorMatching(true),
+		}
+		ps := promqlsmith.New(rnd, lbls, psOpts...)
+
+		opts := &query.Options{
+			Start: time.Unix(0, 0),
+			End:   time.Unix(3600, 0),
+			Step:  time.Minute,
+		}
+		for range testRuns {
+			expr := ps.WalkRangeQuery()
+			exprStr := expr.Pretty(0)
+
+			parsed, err := parser.ParseExpr(exprStr)
+			if err != nil {
+				continue
+			}
+
+			plan, err := NewFromAST(parsed, opts, PlanOptions{})
+			if err != nil {
+				continue
+			}
+
+			optimizedPlan, _ := plan.Optimize(optimizers)
+			root := optimizedPlan.Root()
+
+			Traverse(&root, func(node *Node) {
+				remote, ok := (*node).(RemoteExecution)
+				if !ok {
+					return
+				}
+				if isAbsent(&remote.Query) {
+					return
+				}
+				if !preservesPartitionLabels(remote.Query, engineLabels) {
+					t.Errorf(
+						"remote query does not preserve partition labels\n"+
+							"  original:     %s\n"+
+							"  optimized:    %s\n"+
+							"  remote query: %s",
+						exprStr, root.String(), remote.Query.String(),
+					)
+				}
+			})
+		}
+	})
+}