execution: push limits into telemetry

engine/engine.go

@@ -446,7 +446,7 @@ func (e *Engine) makeQueryOpts(start time.Time, end time.Time, step time.Duratio
 		EnableAnalysis:           e.enableAnalysis,
 		NoStepSubqueryIntervalFn: e.noStepSubqueryIntervalFn,
 		DecodingConcurrency:      e.decodingConcurrency,
-		SampleTracker:            query.NewSampleTracker(e.maxSamplesPerQuery),
+		SampleLimiter:            telemetry.NewSampleLimiter(e.maxSamplesPerQuery, start.UnixMilli(), end.UnixMilli(), telemetry.StepTrackingInterval(step)),
 	}
 
 	if opts == nil {
@@ -459,7 +459,6 @@ func (e *Engine) makeQueryOpts(start time.Time, end time.Time, step time.Duratio
 	if opts.EnablePerStepStats() {
 		res.EnablePerStepStats = opts.EnablePerStepStats()
 	}
-
 	if opts.DecodingConcurrency != 0 {
 		res.DecodingConcurrency = opts.DecodingConcurrency
 	}

engine/explain.go

@@ -8,7 +8,6 @@ import (
 
 	"github.com/thanos-io/promql-engine/execution/model"
 	"github.com/thanos-io/promql-engine/execution/telemetry"
-	"github.com/thanos-io/promql-engine/logicalplan"
 
 	"github.com/prometheus/prometheus/promql"
 )
@@ -64,10 +63,10 @@ func (a *AnalyzeOutputNode) aggregateSamples() {
 			childPeak := child.PeakSamples()
 			a.peakSamples = max(a.peakSamples, childPeak)
 
-			switch a.OperatorTelemetry.LogicalNode().(type) {
-			case *logicalplan.Subquery:
+			switch {
+			case a.OperatorTelemetry.IsSubquery():
 				// Skip aggregating samples for subquery
-			case *logicalplan.StepInvariantExpr:
+			case a.OperatorTelemetry.IsStepInvariant():
 				childSamples := child.TotalSamples()
 				for i := range a.totalSamplesPerStep {
 					a.totalSamples += childSamples

execution/aggregate/count_values.go

@@ -49,7 +49,7 @@ func NewCountValues(next model.VectorOperator, param string, by bool, grouping [
 		by:         by,
 		grouping:   grouping,
 	}
-	return telemetry.NewOperator(telemetry.NewTelemetry(op, opts), op)
+	return telemetry.NewOperator(telemetry.NewTelemetry(op, opts.EnableAnalysis, opts.EnablePerStepStats, opts.Start.UnixMilli(), opts.End.UnixMilli(), opts.Step, opts.SampleLimiter), op)
 }
 
 func (c *countValuesOperator) Explain() []model.VectorOperator {

execution/aggregate/hashaggregate.go

@@ -69,7 +69,7 @@ func NewHashAggregate(
 		params:      make([]float64, opts.StepsBatch),
 	}
 
-	return telemetry.NewOperator(telemetry.NewTelemetry(a, opts), a), nil
+	return telemetry.NewOperator(telemetry.NewTelemetry(a, opts.EnableAnalysis, opts.EnablePerStepStats, opts.Start.UnixMilli(), opts.End.UnixMilli(), opts.Step, opts.SampleLimiter), a), nil
 }
 
 func (a *aggregate) String() string {

execution/aggregate/khashaggregate.go

@@ -80,7 +80,7 @@ func NewKHashAggregate(
 		stepsBatch:  opts.StepsBatch,
 	}
 
-	return telemetry.NewOperator(telemetry.NewTelemetry(op, opts), op), nil
+	return telemetry.NewOperator(telemetry.NewTelemetry(op, opts.EnableAnalysis, opts.EnablePerStepStats, opts.Start.UnixMilli(), opts.End.UnixMilli(), opts.Step, opts.SampleLimiter), op), nil
 }
 
 func (a *kAggregate) Next(ctx context.Context, buf []model.StepVector) (int, error) {

execution/binary/scalar.go

@@ -56,7 +56,7 @@ func NewScalar(
 		stepsBatch: opts.StepsBatch,
 	}
 
-	return telemetry.NewOperator(telemetry.NewTelemetry(op, opts), op), nil
+	return telemetry.NewOperator(telemetry.NewTelemetry(op, opts.EnableAnalysis, opts.EnablePerStepStats, opts.Start.UnixMilli(), opts.End.UnixMilli(), opts.Step, opts.SampleLimiter), op), nil
 }
 
 func (o *scalarOperator) Explain() (next []model.VectorOperator) {

execution/binary/vector.go

@@ -70,7 +70,7 @@ func NewVectorOperator(
 		stepsBatch: opts.StepsBatch,
 	}
 
-	return telemetry.NewOperator(telemetry.NewTelemetry(op, opts), op), nil
+	return telemetry.NewOperator(telemetry.NewTelemetry(op, opts.EnableAnalysis, opts.EnablePerStepStats, opts.Start.UnixMilli(), opts.End.UnixMilli(), opts.Step, opts.SampleLimiter), op), nil
 }
 
 func (o *vectorOperator) String() string {

execution/exchange/coalesce.go

@@ -63,7 +63,7 @@ func NewCoalesce(opts *query.Options, batchSize int64, operators ...model.Vector
 		batchSize:     batchSize,
 	}
 
-	return telemetry.NewOperator(telemetry.NewTelemetry(oper, opts), oper)
+	return telemetry.NewOperator(telemetry.NewTelemetry(oper, opts.EnableAnalysis, opts.EnablePerStepStats, opts.Start.UnixMilli(), opts.End.UnixMilli(), opts.Step, opts.SampleLimiter), oper)
 }
 
 func (c *coalesce) Explain() (next []model.VectorOperator) {

execution/exchange/concurrent.go

@@ -47,7 +47,7 @@ func NewConcurrent(next model.VectorOperator, bufferSize int, opts *query.Option
 		returnChan: make(chan []model.StepVector, bufferSize+2),
 	}
 
-	return telemetry.NewOperator(telemetry.NewTelemetry(oper, opts), oper)
+	return telemetry.NewOperator(telemetry.NewTelemetry(oper, opts.EnableAnalysis, opts.EnablePerStepStats, opts.Start.UnixMilli(), opts.End.UnixMilli(), opts.Step, opts.SampleLimiter), oper)
 }
 
 func (c *concurrencyOperator) Explain() (next []model.VectorOperator) {

execution/exchange/dedup.go

@@ -44,7 +44,7 @@ func NewDedupOperator(next model.VectorOperator, opts *query.Options) model.Vect
 	oper := &dedupOperator{
 		next: next,
 	}
-	return telemetry.NewOperator(telemetry.NewTelemetry(oper, opts), oper)
+	return telemetry.NewOperator(telemetry.NewTelemetry(oper, opts.EnableAnalysis, opts.EnablePerStepStats, opts.Start.UnixMilli(), opts.End.UnixMilli(), opts.Step, opts.SampleLimiter), oper)
 }
 
 func (d *dedupOperator) Next(ctx context.Context, buf []model.StepVector) (int, error) {

execution/exchange/duplicate_label.go

@@ -29,7 +29,7 @@ func NewDuplicateLabelCheck(next model.VectorOperator, opts *query.Options) mode
 	oper := &duplicateLabelCheckOperator{
 		next: next,
 	}
-	return telemetry.NewOperator(telemetry.NewTelemetry(oper, opts), oper)
+	return telemetry.NewOperator(telemetry.NewTelemetry(oper, opts.EnableAnalysis, opts.EnablePerStepStats, opts.Start.UnixMilli(), opts.End.UnixMilli(), opts.Step, opts.SampleLimiter), oper)
 }
 
 func (d *duplicateLabelCheckOperator) Next(ctx context.Context, buf []model.StepVector) (int, error) {

execution/function/absent.go

@@ -31,7 +31,7 @@ func newAbsentOperator(
 		funcExpr: funcExpr,
 		next:     next,
 	}
-	return telemetry.NewOperator(telemetry.NewTelemetry(oper, opts), oper)
+	return telemetry.NewOperator(telemetry.NewTelemetry(oper, opts.EnableAnalysis, opts.EnablePerStepStats, opts.Start.UnixMilli(), opts.End.UnixMilli(), opts.Step, opts.SampleLimiter), oper)
 }
 
 func (o *absentOperator) String() string {

execution/function/histogram.go

@@ -92,7 +92,7 @@ func newHistogramOperator(
 	default:
 		panic("unsupported function passed")
 	}
-	return telemetry.NewOperator(telemetry.NewTelemetry(o, opts), o)
+	return telemetry.NewOperator(telemetry.NewTelemetry(o, opts.EnableAnalysis, opts.EnablePerStepStats, opts.Start.UnixMilli(), opts.End.UnixMilli(), opts.Step, opts.SampleLimiter), o)
 }
 
 func (o *histogramOperator) String() string {

execution/function/operator.go

@@ -75,7 +75,7 @@ func newNoArgsFunctionOperator(funcExpr *logicalplan.FunctionCall, stepsBatch in
 		op.sampleIDs = []uint64{0}
 	}
 
-	return telemetry.NewOperator(telemetry.NewTelemetry(op, opts), op), nil
+	return telemetry.NewOperator(telemetry.NewTelemetry(op, opts.EnableAnalysis, opts.EnablePerStepStats, opts.Start.UnixMilli(), opts.End.UnixMilli(), opts.Step, opts.SampleLimiter), op), nil
 }
 
 // functionOperator returns []model.StepVector after processing input with desired function.
@@ -122,7 +122,7 @@ func newInstantVectorFunctionOperator(funcExpr *logicalplan.FunctionCall, nextOp
 	// Check selector type.
 	switch funcExpr.Args[f.vectorIndex].ReturnType() {
 	case parser.ValueTypeVector, parser.ValueTypeScalar:
-		return telemetry.NewOperator(telemetry.NewTelemetry(f, opts), f), nil
+		return telemetry.NewOperator(telemetry.NewTelemetry(f, opts.EnableAnalysis, opts.EnablePerStepStats, opts.Start.UnixMilli(), opts.End.UnixMilli(), opts.Step, opts.SampleLimiter), f), nil
 	default:
 		return nil, errors.Wrapf(parse.ErrNotImplemented, "got %s:", funcExpr.String())
 	}

execution/function/relabel.go

@@ -35,7 +35,7 @@ func newRelabelOperator(
 		next:     next,
 		funcExpr: funcExpr,
 	}
-	return telemetry.NewOperator(telemetry.NewTelemetry(oper, opts), oper)
+	return telemetry.NewOperator(telemetry.NewTelemetry(oper, opts.EnableAnalysis, opts.EnablePerStepStats, opts.Start.UnixMilli(), opts.End.UnixMilli(), opts.Step, opts.SampleLimiter), oper)
 }
 
 func (o *relabelOperator) String() string {

execution/function/scalar.go

@@ -23,7 +23,7 @@ func newScalarOperator(next model.VectorOperator, opts *query.Options) model.Vec
 		next: next,
 	}
 
-	return telemetry.NewOperator(telemetry.NewTelemetry(oper, opts), oper)
+	return telemetry.NewOperator(telemetry.NewTelemetry(oper, opts.EnableAnalysis, opts.EnablePerStepStats, opts.Start.UnixMilli(), opts.End.UnixMilli(), opts.Step, opts.SampleLimiter), oper)
 }
 
 func (o *scalarOperator) String() string {

execution/function/timestamp.go

@@ -26,7 +26,7 @@ func newTimestampOperator(next model.VectorOperator, opts *query.Options) model.
 	oper := &timestampOperator{
 		next: next,
 	}
-	return telemetry.NewOperator(telemetry.NewTelemetry(oper, opts), oper)
+	return telemetry.NewOperator(telemetry.NewTelemetry(oper, opts.EnableAnalysis, opts.EnablePerStepStats, opts.Start.UnixMilli(), opts.End.UnixMilli(), opts.Step, opts.SampleLimiter), oper)
 }
 
 func (o *timestampOperator) Explain() (next []model.VectorOperator) {

execution/remote/operator.go

@@ -33,17 +33,17 @@ type Execution struct {
 }
 
 func NewExecution(query promql.Query, queryRangeStart, queryRangeEnd time.Time, engineLabels []labels.Labels, opts *query.Options, _ storage.SelectHints) model.VectorOperator {
-	storage := newStorageFromQuery(query, opts, engineLabels)
+	stor := newStorageFromQuery(query, opts, engineLabels)
 	oper := &Execution{
-		storage:         storage,
+		storage:         stor,
 		query:           query,
 		opts:            opts,
 		queryRangeStart: queryRangeStart,
 		queryRangeEnd:   queryRangeEnd,
-		vectorSelector:  promstorage.NewVectorSelector(storage, opts, 0, 0, false, 0, 1),
+		vectorSelector:  promstorage.NewVectorSelector(stor, opts, 0, 0, false, 0, 1),
 	}
 
-	return telemetry.NewOperator(telemetry.NewTelemetry(oper, opts), oper)
+	return telemetry.NewOperator(telemetry.NewTelemetry(oper, opts.EnableAnalysis, opts.EnablePerStepStats, opts.Start.UnixMilli(), opts.End.UnixMilli(), opts.Step, opts.SampleLimiter), oper)
 }
 
 func (e *Execution) Series(ctx context.Context) ([]labels.Labels, error) {

execution/scan/literal_selector.go

@@ -38,7 +38,7 @@ func NewNumberLiteralSelector(opts *query.Options, val float64) model.VectorOper
 		val:         val,
 	}
 
-	return telemetry.NewOperator(telemetry.NewTelemetry(oper, opts), oper)
+	return telemetry.NewOperator(telemetry.NewTelemetry(oper, opts.EnableAnalysis, opts.EnablePerStepStats, opts.Start.UnixMilli(), opts.End.UnixMilli(), opts.Step, opts.SampleLimiter), oper)
 }
 
 func (o *numberLiteralSelector) Explain() (next []model.VectorOperator) {

execution/scan/subquery.go

@@ -20,8 +20,6 @@ import (
 	"github.com/prometheus/prometheus/model/labels"
 )
 
-const sampleLimitCheckPercentage = 0.05
-
 type subqueryOperator struct {
 	next      model.VectorOperator
 	paramOp   model.VectorOperator
@@ -54,9 +52,6 @@ type subqueryOperator struct {
 	paramBuf  []model.StepVector
 	param2Buf []model.StepVector
 	tempBuf   []model.StepVector
-
-	currentTrackedSamples int
-	lastTrackedSamples    int
 }
 
 func NewSubqueryOperator(next, paramOp, paramOp2 model.VectorOperator, opts *query.Options, funcExpr *logicalplan.FunctionCall, subQuery *logicalplan.Subquery) (model.VectorOperator, error) {
@@ -86,7 +81,7 @@ func NewSubqueryOperator(next, paramOp, paramOp2 model.VectorOperator, opts *que
 		params:        make([]float64, opts.StepsBatch),
 		params2:       make([]float64, opts.StepsBatch),
 	}
-	o.telemetry = telemetry.NewSubqueryTelemetry(o, opts)
+	o.telemetry = telemetry.NewSubqueryTelemetry(o, opts.EnableAnalysis, opts.EnablePerStepStats, opts.Start.UnixMilli(), opts.End.UnixMilli(), opts.Step, opts.SampleLimiter)
 	return telemetry.NewOperator(o.telemetry, o), nil
 }
 
@@ -155,9 +150,6 @@ func (o *subqueryOperator) Next(ctx context.Context, buf []model.StepVector) (in
 		for _, b := range o.buffers {
 			b.Reset(mint, maxt+o.subQuery.Offset.Milliseconds())
 		}
-		o.currentTrackedSamples = 0
-		o.lastTrackedSamples = 0
-		checkSampleLimitCounter := 0
 		if len(o.lastVectors) > 0 {
 			for _, v := range o.lastVectors[o.lastCollected+1:] {
 				if v.T > maxt {
@@ -167,7 +159,6 @@ func (o *subqueryOperator) Next(ctx context.Context, buf []model.StepVector) (in
 				o.lastCollected++
 			}
 			if o.lastCollected == len(o.lastVectors)-1 {
-
 				o.lastVectors = nil
 				o.lastCollected = -1
 			}
@@ -192,18 +183,6 @@ func (o *subqueryOperator) Next(ctx context.Context, buf []model.StepVector) (in
 				o.collect(vector, mint)
 			}
 
-			checkSampleLimitCounter++
-			if o.shouldCheckSampleLimit(checkSampleLimitCounter) {
-				if err := o.checkSampleLimit(); err != nil {
-					return 0, err
-				}
-				checkSampleLimitCounter = 0
-			}
-		}
-		if checkSampleLimitCounter > 0 {
-			if err := o.checkSampleLimit(); err != nil {
-				return 0, err
-			}
 		}
 
 		buf[n].Reset(o.currentStep)
@@ -221,7 +200,9 @@ func (o *subqueryOperator) Next(ctx context.Context, buf []model.StepVector) (in
 					buf[n].AppendSampleWithSizeHint(uint64(sampleId), f, hint)
 				}
 			}
-			o.telemetry.IncrementSamplesAtTimestamp(rangeSamples.SampleCount(), buf[n].T)
+			if err := o.telemetry.IncrementSamplesAtTimestamp(rangeSamples.SampleCount(), buf[n].T); err != nil {
+				return 0, err
+			}
 		}
 		n++
 		o.currentStep += o.step
@@ -230,15 +211,6 @@ func (o *subqueryOperator) Next(ctx context.Context, buf []model.StepVector) (in
 	return n, nil
 }
 
-func (o *subqueryOperator) checkSampleLimit() error {
-	delta := o.currentTrackedSamples - o.lastTrackedSamples
-	if delta > 0 {
-		o.opts.SampleTracker.Add(delta)
-	}
-	o.lastTrackedSamples = o.currentTrackedSamples
-	return o.opts.SampleTracker.CheckLimit()
-}
-
 func (o *subqueryOperator) collect(v model.StepVector, mint int64) {
 	if v.T < mint {
 		return
@@ -249,7 +221,6 @@ func (o *subqueryOperator) collect(v model.StepVector, mint int64) {
 			continue
 		}
 		buffer.Push(v.T, ringbuffer.Value{F: s})
-		o.currentTrackedSamples++
 	}
 	for i, s := range v.Histograms {
 		buffer := o.buffers[v.HistogramIDs[i]]
@@ -275,9 +246,7 @@ func (o *subqueryOperator) collect(v model.StepVector, mint int64) {
 			s.CounterResetHint = histogram.UnknownCounterReset
 		}
 		buffer.Push(v.T, ringbuffer.Value{H: s})
-		o.currentTrackedSamples += telemetry.CalculateHistogramSampleCount(s)
 	}
-
 }
 
 func (o *subqueryOperator) Series(ctx context.Context) ([]labels.Labels, error) {
@@ -322,21 +291,3 @@ func (o *subqueryOperator) initSeries(ctx context.Context) error {
 	})
 	return err
 }
-
-func (o *subqueryOperator) shouldCheckSampleLimit(checkSampleLimitCounter int) bool {
-	if len(o.series) == 0 {
-		return checkSampleLimitCounter >= 1
-	}
-
-	limit := o.opts.SampleTracker.Limit()
-	targetSamplesPerCheck := int(float64(limit) * sampleLimitCheckPercentage)
-
-	maxSamplesPerCall := len(o.series) * o.stepsBatch
-	if maxSamplesPerCall == 0 {
-		return checkSampleLimitCounter >= 1
-	}
-
-	interval := max(targetSamplesPerCheck/maxSamplesPerCall, 1)
-
-	return checkSampleLimitCounter >= interval
-}

execution/step_invariant/step_invariant.go

@@ -65,7 +65,7 @@ func NewStepInvariantOperator(
 		u.cacheResult = false
 	}
 
-	return telemetry.NewOperator(telemetry.NewStepInvariantTelemetry(u, opts), u), nil
+	return telemetry.NewOperator(telemetry.NewStepInvariantTelemetry(u, opts.EnableAnalysis, opts.EnablePerStepStats, opts.Start.UnixMilli(), opts.End.UnixMilli(), opts.Step, opts.SampleLimiter), u), nil
 }
 
 func (u *stepInvariantOperator) Series(ctx context.Context) ([]labels.Labels, error) {