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 <[email protected]>

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",

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 {