execution,engine: respect configured max samples

engine/engine.go

@@ -30,9 +30,11 @@ import (
 	"github.com/thanos-io/promql-engine/execution"
 	"github.com/thanos-io/promql-engine/execution/model"
 	"github.com/thanos-io/promql-engine/execution/parse"
+	"github.com/thanos-io/promql-engine/execution/tracking"
 	"github.com/thanos-io/promql-engine/execution/warnings"
 	"github.com/thanos-io/promql-engine/logicalplan"
 	"github.com/thanos-io/promql-engine/parser"
+	"github.com/thanos-io/promql-engine/query"
 )
 
 type QueryType int
@@ -47,6 +49,8 @@ const (
 	subsystem    string    = "engine"
 	InstantQuery QueryType = 1
 	RangeQuery   QueryType = 2
+
+	stepsBatch = 10
 )
 
 type Opts struct {
@@ -217,6 +221,7 @@ func New(opts Opts) *compatibilityEngine {
 		timeout:           opts.Timeout,
 		metrics:           metrics,
 		extLookbackDelta:  opts.ExtLookbackDelta,
+		maxSamples:        opts.MaxSamples,
 		enableAnalysis:    opts.EnableAnalysis,
 	}
 }
@@ -234,6 +239,7 @@ type compatibilityEngine struct {
 	metrics           *engineMetrics
 
 	extLookbackDelta time.Duration
+	maxSamples       int
 	enableAnalysis   bool
 }
 
@@ -268,7 +274,20 @@ func (e *compatibilityEngine) NewInstantQuery(ctx context.Context, q storage.Que
 	})
 	lplan = lplan.Optimize(e.logicalOptimizers)
 
-	exec, err := execution.New(ctx, lplan.Expr(), q, ts, ts, 0, opts.LookbackDelta, e.extLookbackDelta, e.enableAnalysis)
+	qopts := &query.Options{
+		Context:          ctx,
+		Start:            ts,
+		End:              ts,
+		Step:             1,
+		LookbackDelta:    e.lookbackDelta,
+		StepsBatch:       stepsBatch,
+		ExtLookbackDelta: e.extLookbackDelta,
+		EnableAnalysis:   e.enableAnalysis,
+	}
+
+	tracker := tracking.NewTracker(e.maxSamples, qopts)
+
+	exec, err := execution.New(ctx, tracker, lplan.Expr(), q, qopts)
 	if e.triggerFallback(err) {
 		e.metrics.queries.WithLabelValues("true").Inc()
 		return e.prom.NewInstantQuery(ctx, q, opts, qs, ts)
@@ -320,7 +339,20 @@ func (e *compatibilityEngine) NewRangeQuery(ctx context.Context, q storage.Query
 	})
 	lplan = lplan.Optimize(e.logicalOptimizers)
 
-	exec, err := execution.New(ctx, lplan.Expr(), q, start, end, step, opts.LookbackDelta, e.extLookbackDelta, e.enableAnalysis)
+	qopts := &query.Options{
+		Context:          ctx,
+		Start:            start,
+		End:              end,
+		Step:             step,
+		LookbackDelta:    e.lookbackDelta,
+		StepsBatch:       stepsBatch,
+		ExtLookbackDelta: e.extLookbackDelta,
+		EnableAnalysis:   e.enableAnalysis,
+	}
+
+	tracker := tracking.NewTracker(e.maxSamples, qopts)
+
+	exec, err := execution.New(ctx, tracker, lplan.Expr(), q, qopts)
 	if e.triggerFallback(err) {
 		e.metrics.queries.WithLabelValues("true").Inc()
 		return e.prom.NewRangeQuery(ctx, q, opts, qs, start, end, step)

engine/engine_test.go

@@ -114,6 +114,55 @@ func TestQueryExplain(t *testing.T) {
 	}
 }
 
+func TestQueryLimits(t *testing.T) {
+	load := `load 1s
+				example{foo="bar"} 1+0x3600
+				example{foo="baz"} 1+0x3600
+  `
+	test, err := promql.NewTest(t, load)
+	testutil.Ok(t, err)
+	defer test.Close()
+	testutil.Ok(t, test.Run())
+
+	ctx := test.Context()
+
+	newEngine := engine.New(engine.Opts{
+		EngineOpts: promql.EngineOpts{
+			MaxSamples: 1000,
+			Timeout:    10 * time.Second,
+		}})
+
+	t.Run("one series too many samples", func(t *testing.T) {
+		query := `sum_over_time(example{foo="bar"}[1h])`
+		q1, err := newEngine.NewInstantQuery(ctx, test.Queryable(), nil, query, time.Unix(1800, 0))
+		testutil.Ok(t, err)
+
+		newResult := q1.Exec(ctx)
+		testutil.NotOk(t, newResult.Err)
+		testutil.Equals(t, "query processing would load too many samples into memory in query execution", newResult.Err.Error())
+	})
+	t.Run("two series too many samples", func(t *testing.T) {
+		query := `sum_over_time(example[1h])`
+		q1, err := newEngine.NewInstantQuery(ctx, test.Queryable(), nil, query, time.Unix(900, 0))
+		testutil.Ok(t, err)
+
+		newResult := q1.Exec(ctx)
+		testutil.NotOk(t, newResult.Err)
+		testutil.Equals(t, "query processing would load too many samples into memory in query execution", newResult.Err.Error())
+	})
+	t.Run("range query should only account for samples at each batch", func(t *testing.T) {
+		query := `sum(example)`
+		start := time.Unix(0, 0)
+		end := start.Add(time.Hour)
+		step := time.Second
+
+		q1, err := newEngine.NewRangeQuery(ctx, test.Queryable(), nil, query, start, end, step)
+		testutil.Ok(t, err)
+		newResult := q1.Exec(ctx)
+		testutil.Ok(t, newResult.Err)
+	})
+}
+
 func assertExecutionTimeNonZero(t *testing.T, got *engine.AnalyzeOutputNode) bool {
 	if got != nil {
 		if got.OperatorTelemetry.ExecutionTimeTaken() <= 0 {
@@ -2787,8 +2836,8 @@ func TestInstantQuery(t *testing.T) {
 		{
 			name: "duplicate label set",
 			load: `load 5m
-  testmetric1{src="a",dst="b"} 0
-  testmetric2{src="a",dst="b"} 1`,
+testmetric1{src="a",dst="b"} 0
+testmetric2{src="a",dst="b"} 1`,
 			query: "changes({__name__=~'testmetric1|testmetric2'}[5m])",
 		},
 		{
@@ -3661,7 +3710,7 @@ func TestInstantQuery(t *testing.T) {
 								if hasNaNs(oldResult) {
 									t.Log("Applying comparison with NaN equality.")
 									equalsWithNaNs(t, oldResult, newResult)
-								} else if oldResult.Err != nil {
+								} else if oldResult.Err != nil && newResult.Err != nil {
 									testutil.Equals(t, oldResult.Err.Error(), newResult.Err.Error())
 								} else {
 									testutil.Equals(t, oldResult, newResult)