Update max_samples tracking to selector operators

Signed-off-by: Paurush Garg <[email protected]>

Author: PaurushGarg

engine/engine.go

@@ -102,9 +102,6 @@ type QueryOpts struct {
 
 	// LogicalOptimizers can be used to override the LogicalOptimizers engine setting.
 	LogicalOptimizers []logicalplan.Optimizer
-
-	// MaxSamples can be used to override the MaxSamples engine setting.
-	MaxSamples int
 }
 
 func (opts QueryOpts) LookbackDelta() time.Duration { return opts.LookbackDeltaParam }
@@ -453,12 +450,11 @@ func (e *Engine) makeQueryOpts(start time.Time, end time.Time, step time.Duratio
 		EnableAnalysis:           e.enableAnalysis,
 		NoStepSubqueryIntervalFn: e.noStepSubqueryIntervalFn,
 		DecodingConcurrency:      e.decodingConcurrency,
-		MaxSamples:               e.maxSamples,
 	}
 
 	// Initialize SampleTracker if MaxSamples is set
-	if res.MaxSamples > 0 {
-		res.SampleTracker = query.NewSampleTracker(res.MaxSamples)
+	if e.maxSamples > 0 {
+		res.SampleTracker = query.NewSampleTracker(e.maxSamples)
 	}
 
 	if opts == nil {
@@ -476,12 +472,6 @@ func (e *Engine) makeQueryOpts(start time.Time, end time.Time, step time.Duratio
 		res.DecodingConcurrency = opts.DecodingConcurrency
 	}
 
-	// Allow per-query override of MaxSamples
-	if opts.MaxSamples > 0 {
-		res.MaxSamples = opts.MaxSamples
-		res.SampleTracker = query.NewSampleTracker(res.MaxSamples)
-	}
-
 	return res
 }
 
@@ -757,10 +747,5 @@ func recoverEngine(logger *slog.Logger, plan logicalplan.Plan, errp *error) {
 
 		logger.Error("runtime panic in engine", "expr", plan.Root().String(), "err", e, "stacktrace", string(buf))
 		*errp = errors.Wrap(err, "unexpected error")
-	case error:
-		// Handle regular errors (like maxSamples exceeded)
-		*errp = err
-	default:
-		*errp = errors.Newf("unexpected panic: %v", e)
 	}
 }

execution/exchange/concurrent.go

@@ -119,15 +119,6 @@ func (c *concurrencyOperator) Next(ctx context.Context, buf []model.StepVector)
 
 func (c *concurrencyOperator) pull(ctx context.Context) {
 	defer close(c.buffer)
-	defer func() {
-		if r := recover(); r != nil {
-			if err, ok := r.(error); ok {
-				c.buffer <- maybeStepVector{err: err}
-			} else {
-				c.buffer <- maybeStepVector{err: fmt.Errorf("%v", r)}
-			}
-		}
-	}()
 
 	for {
 		select {

execution/model/vector.go

@@ -8,7 +8,6 @@ import (
 
 	"github.com/prometheus/prometheus/model/histogram"
 	"github.com/prometheus/prometheus/model/labels"
-	"github.com/thanos-io/promql-engine/query"
 )
 
 type Series struct {
@@ -25,12 +24,6 @@ type StepVector struct {
 
 	HistogramIDs []uint64
 	Histograms   []*histogram.FloatHistogram
-
-	tracker *query.SampleTracker
-}
-
-func (s *StepVector) SetTracker(tracker *query.SampleTracker) {
-	s.tracker = tracker
 }
 
 // Reset resets the StepVector to the given timestamp while preserving slice capacity.
@@ -69,13 +62,6 @@ func (s *StepVector) AppendSampleWithSizeHint(id uint64, val float64, hint int)
 	}
 	s.SampleIDs = append(s.SampleIDs, id)
 	s.Samples = append(s.Samples, val)
-
-	// Track sample added
-	if s.tracker != nil {
-		if err := s.tracker.Add(1); err != nil {
-			panic(err)
-		}
-	}
 }
 
 func (s *StepVector) AppendSamples(ids []uint64, vals []float64) {
@@ -84,23 +70,11 @@ func (s *StepVector) AppendSamples(ids []uint64, vals []float64) {
 	}
 	s.SampleIDs = append(s.SampleIDs, ids...)
 	s.Samples = append(s.Samples, vals...)
-
-	// Track samples added
-	if s.tracker != nil {
-		if err := s.tracker.Add(len(vals)); err != nil {
-			panic(err)
-		}
-	}
 }
 
 func (s *StepVector) RemoveSample(index int) {
 	s.Samples = slices.Delete(s.Samples, index, index+1)
 	s.SampleIDs = slices.Delete(s.SampleIDs, index, index+1)
-
-	// Track sample removed
-	if s.tracker != nil {
-		s.tracker.Remove(1)
-	}
 }
 
 func (s *StepVector) AppendHistogram(histogramID uint64, h *histogram.FloatHistogram) {
@@ -122,15 +96,6 @@ func (s *StepVector) AppendHistogramWithSizeHint(histogramID uint64, h *histogra
 	}
 	s.HistogramIDs = append(s.HistogramIDs, histogramID)
 	s.Histograms = append(s.Histograms, h)
-
-	// Track histogram added (histograms count as multiple samples based on bucket count)
-	if s.tracker != nil && h != nil {
-		// Count histogram buckets as samples
-		count := len(h.PositiveBuckets) + len(h.NegativeBuckets) + 2 // +2 for count and sum
-		if err := s.tracker.Add(count); err != nil {
-			panic(err)
-		}
-	}
 }
 
 func (s *StepVector) AppendHistograms(histogramIDs []uint64, hs []*histogram.FloatHistogram) {
@@ -139,33 +104,9 @@ func (s *StepVector) AppendHistograms(histogramIDs []uint64, hs []*histogram.Flo
 	}
 	s.HistogramIDs = append(s.HistogramIDs, histogramIDs...)
 	s.Histograms = append(s.Histograms, hs...)
-
-	// Track histograms added
-	if s.tracker != nil {
-		totalCount := 0
-		for _, h := range hs {
-			if h != nil {
-				totalCount += len(h.PositiveBuckets) + len(h.NegativeBuckets) + 2
-			}
-		}
-		if totalCount > 0 {
-			if err := s.tracker.Add(totalCount); err != nil {
-				panic(err)
-			}
-		}
-	}
 }
 
 func (s *StepVector) RemoveHistogram(index int) {
-	// Track histogram being removed
-	if s.tracker != nil && index < len(s.Histograms) {
-		h := s.Histograms[index]
-		if h != nil {
-			count := len(h.PositiveBuckets) + len(h.NegativeBuckets) + 2
-			s.tracker.Remove(count)
-		}
-	}
-
 	s.Histograms = slices.Delete(s.Histograms, index, index+1)
 	s.HistogramIDs = slices.Delete(s.HistogramIDs, index, index+1)
 }

execution/scan/subquery.go

@@ -52,6 +52,8 @@ type subqueryOperator struct {
 	paramBuf  []model.StepVector
 	param2Buf []model.StepVector
 	tempBuf   []model.StepVector
+
+	lastTrackedSamples int
 }
 
 func NewSubqueryOperator(next, paramOp, paramOp2 model.VectorOperator, opts *query.Options, funcExpr *logicalplan.FunctionCall, subQuery *logicalplan.Subquery) (model.VectorOperator, error) {
@@ -207,6 +209,22 @@ func (o *subqueryOperator) Next(ctx context.Context, buf []model.StepVector) (in
 		o.currentStep += o.step
 	}
 
+	if o.opts.SampleTracker != nil {
+		totalSamplesInBatch := 0
+		for _, b := range o.buffers {
+			totalSamplesInBatch += b.SampleCount()
+		}
+		if totalSamplesInBatch > o.lastTrackedSamples {
+			o.opts.SampleTracker.Add(totalSamplesInBatch - o.lastTrackedSamples)
+		} else if totalSamplesInBatch < o.lastTrackedSamples {
+			o.opts.SampleTracker.Remove(o.lastTrackedSamples - totalSamplesInBatch)
+		}
+		o.lastTrackedSamples = totalSamplesInBatch
+		if err := o.opts.SampleTracker.CheckLimit(); err != nil {
+			return 0, err
+		}
+	}
+
 	return n, nil
 }
 
@@ -277,7 +295,7 @@ func (o *subqueryOperator) initSeries(ctx context.Context) error {
 		o.series = make([]labels.Labels, len(series))
 		o.buffers = make([]*ringbuffer.GenericRingBuffer, len(series))
 		for i := range o.buffers {
-			o.buffers[i] = ringbuffer.New(ctx, 8, o.subQuery.Range.Milliseconds(), o.subQuery.Offset.Milliseconds(), o.call, o.opts)
+			o.buffers[i] = ringbuffer.New(ctx, 8, o.subQuery.Range.Milliseconds(), o.subQuery.Offset.Milliseconds(), o.call)
 		}
 		var b labels.ScratchBuilder
 		for i, s := range series {

query/options.go

@@ -18,8 +18,7 @@ type Options struct {
 	NoStepSubqueryIntervalFn func(time.Duration) time.Duration
 	EnableAnalysis           bool
 	DecodingConcurrency      int
-	MaxSamples               int            // Maximum samples allowed in query execution
-	SampleTracker            *SampleTracker // Tracks current samples in memory (internal)
+	SampleTracker            *SampleTracker // Tracks current samples in memory
 }
 
 // TotalSteps returns the total number of steps in the query, regardless of batching.

query/sample_tracker.go

@@ -4,52 +4,49 @@
 package query
 
 import (
+	"fmt"
 	"sync/atomic"
-
-	"github.com/efficientgo/core/errors"
 )
 
-// SampleTracker tracks the number of samples currently in memory during query execution.
-// It enforces a maximum sample limit to prevent out-of-memory errors.
 type SampleTracker struct {
-	current atomic.Int64 // Current samples in memory
-	limit   int64        // Maximum samples allowed (0 = no limit)
+	current atomic.Int64
+	limit   int64
 }
 
-// NewSampleTracker creates a new sample tracker with the given limit.
 func NewSampleTracker(maxSamples int) *SampleTracker {
 	return &SampleTracker{
 		limit: int64(maxSamples),
 	}
 }
 
-// Add increments the current sample count and checks against the limit.
-// Returns an error if adding these samples would exceed the limit.
-func (st *SampleTracker) Add(count int) error {
-	if count <= 0 {
-		return nil
-	}
-
-	newCurrent := st.current.Add(int64(count))
-
-	// Check limit
-	if st.limit > 0 && newCurrent > st.limit {
-		return errors.Newf("query exceeded maximum samples limit: current=%d, limit=%d", newCurrent, st.limit)
-	}
-
-	return nil
+func (st *SampleTracker) Add(count int) {
+	st.current.Add(int64(count))
 }
 
-// Remove decrements the current sample count.
-// This should be called when samples are freed from memory.
 func (st *SampleTracker) Remove(count int) {
-	if count <= 0 {
-		return
-	}
 	st.current.Add(-int64(count))
 }
 
-// Current returns the current number of samples in memory.
 func (st *SampleTracker) Current() int64 {
 	return st.current.Load()
 }
+
+func (st *SampleTracker) CheckLimit() error {
+	if st.limit <= 0 {
+		return nil
+	}
+	current := st.current.Load()
+	if current > st.limit {
+		return ErrMaxSamplesExceeded{Current: current, Limit: st.limit}
+	}
+	return nil
+}
+
+type ErrMaxSamplesExceeded struct {
+	Current int64
+	Limit   int64
+}
+
+func (e ErrMaxSamplesExceeded) Error() string {
+	return fmt.Sprintf("query exceeded maximum samples limit: current=%d, limit=%d", e.Current, e.Limit)
+}

ringbuffer/generic.go

@@ -9,7 +9,6 @@ import (
 
 	"github.com/thanos-io/promql-engine/execution/telemetry"
 	"github.com/thanos-io/promql-engine/warnings"
-	"github.com/thanos-io/promql-engine/query"
 
 	"github.com/prometheus/prometheus/model/histogram"
 )
@@ -47,27 +46,20 @@ type GenericRingBuffer struct {
 	selectRange int64
 	extLookback int64
 	call        FunctionCall
-
-	tracker *query.SampleTracker
 }
 
-func New(ctx context.Context, size int, selectRange, offset int64, call FunctionCall, opts *query.Options) *GenericRingBuffer {
-	return NewWithExtLookback(ctx, size, selectRange, offset, 0, call, opts)
+func New(ctx context.Context, size int, selectRange, offset int64, call FunctionCall) *GenericRingBuffer {
+	return NewWithExtLookback(ctx, size, selectRange, offset, 0, call)
 }
 
-func NewWithExtLookback(ctx context.Context, size int, selectRange, offset, extLookback int64, call FunctionCall, opts *query.Options) *GenericRingBuffer {
-	var tracker *query.SampleTracker
-	if opts != nil {
-		tracker = opts.SampleTracker
-	}
+func NewWithExtLookback(ctx context.Context, size int, selectRange, offset, extLookback int64, call FunctionCall) *GenericRingBuffer {
 	return &GenericRingBuffer{
 		ctx:         ctx,
 		items:       make([]Sample, 0, size),
 		selectRange: selectRange,
 		offset:      offset,
 		extLookback: extLookback,
 		call:        call,
-		tracker:     tracker,
 	}
 }
 
@@ -118,35 +110,11 @@ func (r *GenericRingBuffer) Push(t int64, v Value) {
 	} else {
 		r.items[n].V.H = nil
 	}
-
-	// Track samples added (count histogram buckets like Prometheus)
-	if r.tracker != nil {
-		samplesAdded := 1
-		if v.H != nil {
-			// Count histogram size in 16-byte units (same as Prometheus)
-			samplesAdded = (v.H.Size() + 8) / 16
-		}
-		if err := r.tracker.Add(samplesAdded); err != nil {
-			panic(err)
-		}
-	}
 }
 
 func (r *GenericRingBuffer) Reset(mint int64, evalt int64) {
 	r.currentStep = evalt
 	if r.extLookback == 0 && (len(r.items) == 0 || r.items[len(r.items)-1].T < mint) {
-		// Track samples removed (count histogram sizes like Prometheus)
-		if r.tracker != nil && len(r.items) > 0 {
-			samplesRemoved := 0
-			for _, item := range r.items {
-				if item.V.H != nil {
-					samplesRemoved += (item.V.H.Size() + 8) / 16
-				} else {
-					samplesRemoved++
-				}
-			}
-			r.tracker.Remove(samplesRemoved)
-		}
 		r.items = r.items[:0]
 		return
 	}
@@ -157,19 +125,6 @@ func (r *GenericRingBuffer) Reset(mint int64, evalt int64) {
 		drop--
 	}
 
-	// Track samples removed (count histogram sizes like Prometheus)
-	if r.tracker != nil && drop > 0 {
-		samplesRemoved := 0
-		for i := 0; i < drop; i++ {
-			if r.items[i].V.H != nil {
-				samplesRemoved += (r.items[i].V.H.Size() + 8) / 16
-			} else {
-				samplesRemoved++
-			}
-		}
-		r.tracker.Remove(samplesRemoved)
-	}
-
 	keep := len(r.items) - drop
 	r.tail = resize(r.tail, drop)
 	copy(r.tail, r.items[:drop])