From b4e2909bda0393d47ce1b8b2872a6053e9c4a20b Mon Sep 17 00:00:00 2001
From: Stephan Renatus <stephan@styra.com>
Date: Fri, 6 Dec 2024 15:21:15 +0100
Subject: [PATCH] topdown: move context.Context cancellation check (#7210)

The topdown Cancel machinery is there because it's cheap to check. ctx.Err() is
not.

This change moves the "Is the context the cause for cancellation?" check into
the branch were evaluation has been aborted through the topdown.Cancel call.

When evaluation has already been cancelled, an expensive check no longer
matters much -- when it's still ongoing, it'll affect the overall performance.

Signed-off-by: Stephan Renatus <stephan@styra.com>
---
 topdown/eval.go      | 15 +++++++--------
 topdown/eval_test.go |  6 +++++-
 2 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/topdown/eval.go b/topdown/eval.go
index e716cee0d6..ade0f0d3a0 100644
--- a/topdown/eval.go
+++ b/topdown/eval.go
@@ -340,15 +340,14 @@ func (e *eval) evalExpr(iter evalIterator) error {
 		return &earlyExitError{prev: err, e: e}
 	}
 
-	if e.ctx != nil && e.ctx.Err() != nil {
-		return &Error{
-			Code:    CancelErr,
-			Message: e.ctx.Err().Error(),
-			err:     e.ctx.Err(),
-		}
-	}
-
 	if e.cancel != nil && e.cancel.Cancelled() {
+		if e.ctx != nil && e.ctx.Err() != nil {
+			return &Error{
+				Code:    CancelErr,
+				Message: e.ctx.Err().Error(),
+				err:     e.ctx.Err(),
+			}
+		}
 		return &Error{
 			Code:    CancelErr,
 			Message: "caller cancelled query execution",
diff --git a/topdown/eval_test.go b/topdown/eval_test.go
index f135c48fc0..88d72dddb6 100644
--- a/topdown/eval_test.go
+++ b/topdown/eval_test.go
@@ -1625,12 +1625,16 @@ func TestContextErrorHandling(t *testing.T) {
 			txn := storage.NewTransactionOrDie(ctx, store)
 			defer store.Abort(ctx, txn)
 
+			c := NewCancel()
 			query := NewQuery(ast.MustParseBody("")).
 				WithCompiler(compiler).
 				WithStore(store).
-				WithTransaction(txn)
+				WithTransaction(txn).
+				WithCancel(c)
 
 			testCtx := tc.before()
+			c.Cancel()
+
 			qrs, err := query.Run(testCtx)
 
 			if err == nil {