Fix failure

Signed-off-by: 🌲 Harry 🌊 John 🏔 <[email protected]>

Author: harry671003

engine/engine_test.go

@@ -246,7 +246,7 @@ func TestQueriesAgainstOldEngine(t *testing.T) {
 			query: `predict_linear(http_requests_total{route="/"}[1h:1m] offset 1m, 60)`,
 		},
 		{
-			name: "predict_linear correctness",
+			name: "predict_linear with step invariant",
 			load: `load 30s
 			    http_requests_total{pod="nginx-1"} 1+1.1x10
 			    http_requests_total{pod="nginx-2"} 2+2.3x50`,

ringbuffer/generic.go

@@ -63,6 +63,10 @@ func (r *GenericRingBuffer) ReadIntoLast(f func(*Sample)) {
 	f(&r.items[len(r.items)-1])
 }
 
+func (r *GenericRingBuffer) SetStep(step int64) {
+	r.currentStep = step
+}
+
 // Push adds a new sample to the buffer.
 func (r *GenericRingBuffer) Push(t int64, v Value) {
 	n := len(r.items)

ringbuffer/rate.go

@@ -18,6 +18,7 @@ type Buffer interface {
 	MaxT() int64
 	Push(t int64, v Value)
 	Reset(mint int64, evalt int64)
+	SetStep(int64)
 	Eval(ctx context.Context, _, _ float64, _ *int64) (float64, *histogram.FloatHistogram, bool, error)
 	ReadIntoLast(f func(*Sample))
 }
@@ -96,6 +97,8 @@ func (r *RateBuffer) Len() int { return r.stepRanges[0].numSamples }
 
 func (r *RateBuffer) MaxT() int64 { return r.last.T }
 
+func (r *RateBuffer) SetStep(step int64) { r.step = step }
+
 func (r *RateBuffer) Push(t int64, v Value) {
 	// Detect resets and store the current and previous sample so that
 	// the rate is properly adjusted.

storage/prometheus/matrix_selector.go

@@ -36,6 +36,8 @@ type matrixScanner struct {
 	iterator         chunkenc.Iterator
 	lastSample       ringbuffer.Sample
 	metricAppearedTs *int64
+	timestamp        *int64
+	once             sync.Once
 }
 
 type matrixSelector struct {
@@ -57,6 +59,7 @@ type matrixSelector struct {
 	numSteps      int
 	mint          int64
 	maxt          int64
+	timestamp     *int64
 	step          int64
 	selectRange   int64
 	offset        int64
@@ -85,6 +88,7 @@ func NewMatrixSelector(
 	functionName string,
 	arg float64,
 	arg2 float64,
+	timestamp *int64,
 	opts *query.Options,
 	selectRange, offset time.Duration,
 	batchSize int64,
@@ -108,6 +112,7 @@ func NewMatrixSelector(
 		mint:          opts.Start.UnixMilli(),
 		maxt:          opts.End.UnixMilli(),
 		step:          opts.Step.Milliseconds(),
+		timestamp:     timestamp,
 		isExtFunction: parse.IsExtFunction(functionName),
 
 		selectRange:     selectRange.Milliseconds(),
@@ -250,6 +255,7 @@ func (o *matrixSelector) loadSeries(ctx context.Context) error {
 				iterator:   s.Iterator(nil),
 				lastSample: ringbuffer.Sample{T: math.MinInt64},
 				buffer:     o.newBuffer(ctx),
+				timestamp:  o.timestamp,
 			}
 			o.series[i] = lbls
 		}
@@ -301,6 +307,23 @@ func (o *matrixSelector) String() string {
 	return fmt.Sprintf("[matrixSelector] {%v}[%s] %v mod %v", o.storage.Matchers(), r, o.shard, o.numShards)
 }
 
+func (m *matrixScanner) selectPoints(
+	mint, maxt, evalt int64,
+	fh *histogram.FloatHistogram,
+	isExtFunction bool,
+) error {
+	if m.timestamp != nil {
+		// If timestamp is set, then the points should be selected only once.
+		var err error
+		m.once.Do(func() {
+			err = m.selectPointsInternal(mint, maxt, *m.timestamp, fh, isExtFunction)
+		})
+		m.buffer.SetStep(evalt)
+		return err
+	}
+	return m.selectPointsInternal(mint, maxt, evalt, fh, isExtFunction)
+}
+
 // matrixIterSlice populates a matrix vector covering the requested range for a
 // single time series, with points retrieved from an iterator.
 //
@@ -310,7 +333,7 @@ func (o *matrixSelector) String() string {
 // into the [mint, maxt] range are retained; only points with later timestamps
 // are populated from the iterator.
 // TODO(fpetkovski): Add max samples limit.
-func (m *matrixScanner) selectPoints(
+func (m *matrixScanner) selectPointsInternal(
 	mint, maxt, evalt int64,
 	fh *histogram.FloatHistogram,
 	isExtFunction bool,

storage/prometheus/scanners.go

@@ -134,6 +134,7 @@ func (p Scanners) NewMatrixSelector(
 			call.Func.Name,
 			arg,
 			arg2,
+			logicalNode.VectorSelector.Timestamp,
 			opts,
 			logicalNode.Range,
 			vs.Offset,