From d778591fd3b98f779d5194fd30d07dd813c0ae9b Mon Sep 17 00:00:00 2001
From: Ben Ye <benye@amazon.com>
Date: Sun, 21 Jan 2024 14:19:39 -0800
Subject: [PATCH] add more context cancellation check at evaluation time

Signed-off-by: Ben Ye <benye@amazon.com>
---
 promql/engine.go | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/promql/engine.go b/promql/engine.go
index 6349ddbbe..7165631e0 100644
--- a/promql/engine.go
+++ b/promql/engine.go
@@ -1471,6 +1471,9 @@ func (ev *evaluator) eval(expr parser.Expr) (parser.Value, annotations.Annotatio
 		it := storage.NewBuffer(selRange)
 		var chkIter chunkenc.Iterator
 		for i, s := range selVS.Series {
+			if err := contextDone(ev.ctx, "expression evaluation"); err != nil {
+				ev.error(err)
+			}
 			ev.currentSamples -= len(floats) + totalHPointSize(histograms)
 			if floats != nil {
 				floats = floats[:0]
@@ -1690,6 +1693,9 @@ func (ev *evaluator) eval(expr parser.Expr) (parser.Value, annotations.Annotatio
 		it := storage.NewMemoizedEmptyIterator(durationMilliseconds(ev.lookbackDelta))
 		var chkIter chunkenc.Iterator
 		for i, s := range e.Series {
+			if err := contextDone(ev.ctx, "expression evaluation"); err != nil {
+				ev.error(err)
+			}
 			chkIter = s.Iterator(chkIter)
 			it.Reset(chkIter)
 			ss := Series{