Add more opportunities for cancellation for validate()

.vscode/settings.json

@@ -10,8 +10,8 @@
     },
 
     // VS Code Test Runner configuration
-    "testing.automaticallyOpenTestResults": "openOnTestStart",
     "testing.automaticallyOpenPeekView": "failureInVisibleDocument",
+    "testing.automaticallyOpenTestResults": "openOnTestStart",
     "testing.defaultGutterClickAction": "run",
     "testing.followRunningTest": true,
 

.vscode/tasks.json

@@ -11,6 +11,9 @@
             "group": {
                 "kind": "build",
                 "isDefault": true
+            },
+            "presentation": {
+                "clear": true
             }
         }
     ]

eslint.config.js

@@ -35,6 +35,7 @@ module.exports = [
                 __dirname: "readonly",
                 module: "readonly",
                 setTimeout: "readonly",
+                console: "readonly",
             },
         },
         plugins: {

src/powerquery-language-services/inspection/invokeExpression/invokeExpression.ts

@@ -65,8 +65,6 @@ async function inspectInvokeExpression(
         correlationId,
     );
 
-    settings.cancellationToken?.throwIfCancelled();
-
     const invokeExpressionXorNode: TXorNode | undefined = NodeIdMapUtils.xor(nodeIdMapCollection, invokeExpressionId);
 
     if (invokeExpressionXorNode === undefined) {
@@ -93,6 +91,8 @@ async function inspectInvokeExpression(
         await tryType(settings, nodeIdMapCollection, previousNode.node.id, typeCache),
     );
 
+    settings.cancellationToken?.throwIfCancelled();
+
     let invokeExpressionArgs: InvokeExpressionArguments | undefined;
 
     if (TypeUtils.isDefinedFunction(functionType)) {
@@ -213,6 +213,7 @@ async function getArgumentTypes(
     for (const xorNode of argXorNodes) {
         // eslint-disable-next-line no-await-in-loop
         const triedArgType: TriedType = await tryType(settings, nodeIdMapCollection, xorNode.node.id, typeCache);
+        settings.cancellationToken?.throwIfCancelled();
 
         if (ResultUtils.isError(triedArgType)) {
             throw triedArgType;

src/powerquery-language-services/promiseUtils.ts

@@ -0,0 +1,45 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+import { ICancellationToken } from "@microsoft/powerquery-parser";
+import { setImmediate } from "timers";
+
+/**
+ * Sequential processing with cancellation support - often better than Promise.all
+ * for cancellation scenarios because we can stop between each operation.
+ * Also yields control before each operation to allow cancellation tokens to take effect.
+ */
+export async function processSequentiallyWithCancellation<T, R>(
+    items: T[],
+    processor: (item: T) => Promise<R>,
+    cancellationToken?: ICancellationToken,
+): Promise<R[]> {
+    const results: R[] = [];
+
+    for (const item of items) {
+        // Yield control to allow async cancellation tokens to take effect
+        // eslint-disable-next-line no-await-in-loop
+        await yieldForCancellation(cancellationToken);
+
+        // eslint-disable-next-line no-await-in-loop
+        const result: R = await processor(item);
+        results.push(result);
+    }
+
+    return results;
+}
+
+export async function yieldForCancellation(cancellationToken?: ICancellationToken): Promise<void> {
+    if (cancellationToken) {
+        // First yield to microtasks (handles synchronous cancellation)
+        await Promise.resolve();
+        cancellationToken.throwIfCancelled();
+
+        // Additional yield for setImmediate-based cancellation tokens
+        // This ensures we yield to the timer queue where setImmediate callbacks execute
+        await new Promise<void>((resolve: () => void) => setImmediate(resolve));
+        cancellationToken.throwIfCancelled();
+    }
+
+    return Promise.resolve();
+}

src/powerquery-language-services/validate/validate.ts

@@ -1,13 +1,15 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
+import { CommonError, Result, ResultUtils } from "@microsoft/powerquery-parser";
 import { NodeIdMap, ParseError, ParseState } from "@microsoft/powerquery-parser/lib/powerquery-parser/parser";
 import { Diagnostic } from "vscode-languageserver-types";
 import { TextDocument } from "vscode-languageserver-textdocument";
 import { Trace } from "@microsoft/powerquery-parser/lib/powerquery-parser/common/trace";
 
+import * as PromiseUtils from "../promiseUtils";
+
 import { Analysis, AnalysisSettings, AnalysisUtils } from "../analysis";
-import { CommonError, Result, ResultUtils } from "@microsoft/powerquery-parser";
 import { TypeCache } from "../inspection";
 import { validateDuplicateIdentifiers } from "./validateDuplicateIdentifiers";
 import { validateFunctionExpression } from "./validateFunctionExpression";
@@ -35,7 +37,10 @@ export function validate(
             initialCorrelationId: trace.id,
         };
 
+        validationSettings.cancellationToken?.throwIfCancelled();
+
         const analysis: Analysis = AnalysisUtils.analysis(textDocument, analysisSettings);
+
         const parseState: ParseState | undefined = ResultUtils.assertOk(await analysis.getParseState());
         const parseError: ParseError.ParseError | undefined = ResultUtils.assertOk(await analysis.getParseError());
 
@@ -58,49 +63,56 @@ export function validate(
             return result;
         }
 
-        let functionExpressionDiagnostics: Diagnostic[];
-        let invokeExpressionDiagnostics: Diagnostic[];
-        let unknownIdentifiersDiagnostics: Diagnostic[];
-
         const nodeIdMapCollection: NodeIdMap.Collection = parseState.contextState.nodeIdMapCollection;
         const typeCache: TypeCache = analysis.getTypeCache();
 
-        if (validationSettings.checkInvokeExpressions && nodeIdMapCollection) {
-            functionExpressionDiagnostics = validateFunctionExpression(validationSettings, nodeIdMapCollection);
+        // Define validation operations to run sequentially
+        const validationOperations: (() => Promise<Diagnostic[]>)[] = [
+            // Parse validation (if there are parse errors)
+            async (): Promise<Diagnostic[]> => await validateParse(parseError, updatedSettings),
+        ];
+
+        // Add conditional validations based on settings
+        if (validationSettings.checkForDuplicateIdentifiers && nodeIdMapCollection) {
+            validationOperations.push(
+                async (): Promise<Diagnostic[]> =>
+                    await validateDuplicateIdentifiers(
+                        textDocument,
+                        nodeIdMapCollection,
+                        updatedSettings,
+                        validationSettings.cancellationToken,
+                    ),
+            );
+        }
 
-            invokeExpressionDiagnostics = await validateInvokeExpression(
-                validationSettings,
-                nodeIdMapCollection,
-                typeCache,
+        if (validationSettings.checkInvokeExpressions && nodeIdMapCollection) {
+            validationOperations.push(
+                async (): Promise<Diagnostic[]> =>
+                    await validateFunctionExpression(validationSettings, nodeIdMapCollection),
+                async (): Promise<Diagnostic[]> =>
+                    await validateInvokeExpression(validationSettings, nodeIdMapCollection, typeCache),
             );
-        } else {
-            functionExpressionDiagnostics = [];
-            invokeExpressionDiagnostics = [];
         }
 
         if (validationSettings.checkUnknownIdentifiers && nodeIdMapCollection) {
-            unknownIdentifiersDiagnostics = await validateUnknownIdentifiers(
-                validationSettings,
-                nodeIdMapCollection,
-                typeCache,
+            validationOperations.push(
+                async (): Promise<Diagnostic[]> =>
+                    await validateUnknownIdentifiers(validationSettings, nodeIdMapCollection, typeCache),
             );
-        } else {
-            unknownIdentifiersDiagnostics = [];
         }
 
+        // Execute all validation operations sequentially with cancellation support
+        const allDiagnostics: Diagnostic[][] = await PromiseUtils.processSequentiallyWithCancellation(
+            validationOperations,
+            (operation: () => Promise<Diagnostic[]>) => operation(),
+            validationSettings.cancellationToken,
+        );
+
+        // Flatten all diagnostics into a single array
+        const flattenedDiagnostics: Diagnostic[] = allDiagnostics.flat();
+
         const result: ValidateOk = {
-            diagnostics: [
-                ...validateDuplicateIdentifiers(
-                    textDocument,
-                    nodeIdMapCollection,
-                    updatedSettings,
-                    validationSettings.cancellationToken,
-                ),
-                ...(await validateParse(parseError, updatedSettings)),
-                ...functionExpressionDiagnostics,
-                ...invokeExpressionDiagnostics,
-                ...unknownIdentifiersDiagnostics,
-            ],
+            diagnostics: flattenedDiagnostics,
             hasSyntaxError: parseError !== undefined,
         };