Make signature help node building cancellable (#23543)

* Make token building cancellable

* Scope cancellation token, make find all refs and quickinfo cancellable

* Make completion entry details cancellable

* Actually accept public API update

* Add test verifying cancellations within checker for select language service operations

* Document runWithCancellationToken a bit more

* Add post-cancellation verification to test

src/compiler/checker.ts

@@ -325,6 +325,16 @@ namespace ts {
                     return diagnostics;
                 }
             },
+
+            runWithCancellationToken: (token, callback) => {
+                try {
+                    cancellationToken = token;
+                    return callback(checker);
+                }
+                finally {
+                    cancellationToken = undefined;
+                }
+            }
         };
 
         const tupleTypes: GenericType[] = [];
@@ -2988,6 +2998,9 @@ namespace ts {
             }
 
             function typeToTypeNodeHelper(type: Type, context: NodeBuilderContext): TypeNode {
+                if (cancellationToken && cancellationToken.throwIfCancellationRequested) {
+                    cancellationToken.throwIfCancellationRequested();
+                }
                 const inTypeAlias = context.flags & NodeBuilderFlags.InTypeAlias;
                 context.flags &= ~NodeBuilderFlags.InTypeAlias;
 

src/compiler/types.ts

@@ -3000,6 +3000,13 @@ namespace ts {
          * Others are added in computeSuggestionDiagnostics.
          */
         /* @internal */ getSuggestionDiagnostics(file: SourceFile): ReadonlyArray<Diagnostic>;
+
+        /**
+         * Depending on the operation performed, it may be appropriate to throw away the checker
+         * if the cancellation token is triggered. Typically, if it is used for error checking
+         * and the operation is cancelled, then it should be discarded, otherwise it is safe to keep.
+         */
+        runWithCancellationToken<T>(token: CancellationToken, cb: (checker: TypeChecker) => T): T;
     }
 
     /* @internal */

src/harness/unittests/cancellableLanguageServiceOperations.ts

@@ -0,0 +1,95 @@
+/// <reference path="..\harness.ts" />
+
+namespace ts {
+    describe("cancellableLanguageServiceOperations", () => {
+        const file = `
+        function foo(): void;
+        function foo<T>(x: T): T;
+        function foo<T>(x?: T): T | void {}
+        foo(f);
+        `;
+        it("can cancel signature help mid-request", () => {
+            verifyOperationCancelledAfter(file, 4, service => // Two calls are top-level in services, one is the root type, and the second should be for the parameter type
+                service.getSignatureHelpItems("file.ts", file.lastIndexOf("f")),
+                r => assert.exists(r.items[0])
+            );
+        });
+
+        it("can cancel find all references mid-request", () => {
+            verifyOperationCancelledAfter(file, 3, service => // Two calls are top-level in services, one is the root type
+                service.findReferences("file.ts", file.lastIndexOf("o")),
+                r => assert.exists(r[0].definition)
+            );
+        });
+
+        it("can cancel quick info mid-request", () => {
+            verifyOperationCancelledAfter(file, 1, service => // The LS doesn't do any top-level checks on the token for quickinfo, so the first check is within the checker
+                service.getQuickInfoAtPosition("file.ts", file.lastIndexOf("o")),
+                r => assert.exists(r.displayParts)
+            );
+        });
+
+        it("can cancel completion entry details mid-request", () => {
+            const options: FormatCodeSettings = {
+                indentSize: 4,
+                tabSize: 4,
+                newLineCharacter: "\n",
+                convertTabsToSpaces: true,
+                indentStyle: IndentStyle.Smart,
+                insertSpaceAfterConstructor: false,
+                insertSpaceAfterCommaDelimiter: true,
+                insertSpaceAfterSemicolonInForStatements: true,
+                insertSpaceBeforeAndAfterBinaryOperators: true,
+                insertSpaceAfterKeywordsInControlFlowStatements: true,
+                insertSpaceAfterFunctionKeywordForAnonymousFunctions: false,
+                insertSpaceAfterOpeningAndBeforeClosingNonemptyParenthesis: false,
+                insertSpaceAfterOpeningAndBeforeClosingNonemptyBrackets: false,
+                insertSpaceAfterOpeningAndBeforeClosingNonemptyBraces: true,
+                insertSpaceAfterOpeningAndBeforeClosingTemplateStringBraces: false,
+                insertSpaceAfterOpeningAndBeforeClosingJsxExpressionBraces: false,
+                insertSpaceBeforeFunctionParenthesis: false,
+                placeOpenBraceOnNewLineForFunctions: false,
+                placeOpenBraceOnNewLineForControlBlocks: false,
+            };
+            verifyOperationCancelledAfter(file, 1, service => // The LS doesn't do any top-level checks on the token for completion entry details, so the first check is within the checker
+                service.getCompletionEntryDetails("file.ts", file.lastIndexOf("f"), "foo", options, /*content*/ undefined, {}),
+                r => assert.exists(r.displayParts)
+            );
+        });
+    });
+
+    function verifyOperationCancelledAfter<T>(content: string, cancelAfter: number, operation: (service: LanguageService) => T, validator: (arg: T) => void) {
+        let checks = 0;
+        const token: HostCancellationToken = {
+            isCancellationRequested() {
+                checks++;
+                const result = checks >= cancelAfter;
+                if (result) {
+                    checks = -Infinity; // Cancel just once, then disable cancellation, effectively
+                }
+                return result;
+            }
+        };
+        const adapter = new Harness.LanguageService.NativeLanguageServiceAdapter(token);
+        const host = adapter.getHost();
+        host.addScript("file.ts", content, /*isRootFile*/ true);
+        const service = adapter.getLanguageService();
+        assertCancelled(() => operation(service));
+        validator(operation(service));
+    }
+
+    /**
+     * We don't just use `assert.throws` because it doesn't validate instances for thrown objects which do not inherit from `Error`
+     */
+    function assertCancelled(cb: () => void) {
+        let caught: any;
+        try {
+            cb();
+        }
+        catch (e) {
+            caught = e;
+        }
+        assert.exists(caught, "Expected operation to be cancelled, but was not");
+        assert.instanceOf(caught, OperationCanceledException);
+    }
+}

src/services/completions.ts

@@ -534,6 +534,7 @@ namespace ts.Completions {
         formatContext: formatting.FormatContext,
         getCanonicalFileName: GetCanonicalFileName,
         preferences: UserPreferences,
+        cancellationToken: CancellationToken,
     ): CompletionEntryDetails {
         const typeChecker = program.getTypeChecker();
         const compilerOptions = program.getCompilerOptions();
@@ -544,7 +545,7 @@ namespace ts.Completions {
             const stringLiteralCompletions = !contextToken || !isStringLiteralLike(contextToken)
                 ? undefined
                 : getStringLiteralCompletionEntries(sourceFile, contextToken, position, typeChecker, compilerOptions, host);
-            return stringLiteralCompletions && stringLiteralCompletionDetails(name, contextToken, stringLiteralCompletions, sourceFile, typeChecker);
+            return stringLiteralCompletions && stringLiteralCompletionDetails(name, contextToken, stringLiteralCompletions, sourceFile, typeChecker, cancellationToken);
         }
 
         // Compute all the completion symbols again.
@@ -566,7 +567,7 @@ namespace ts.Completions {
             case "symbol": {
                 const { symbol, location, symbolToOriginInfoMap, previousToken } = symbolCompletion;
                 const { codeActions, sourceDisplay } = getCompletionEntryCodeActionsAndSourceDisplay(symbolToOriginInfoMap, symbol, program, typeChecker, host, compilerOptions, sourceFile, previousToken, formatContext, getCanonicalFileName, program.getSourceFiles(), preferences);
-                return createCompletionDetailsForSymbol(symbol, typeChecker, sourceFile, location, codeActions, sourceDisplay);
+                return createCompletionDetailsForSymbol(symbol, typeChecker, sourceFile, location, cancellationToken, codeActions, sourceDisplay);
             }
             case "none":
                 // Didn't find a symbol with this name.  See if we can find a keyword instead.
@@ -574,12 +575,15 @@ namespace ts.Completions {
         }
     }
 
-    function createCompletionDetailsForSymbol(symbol: Symbol, checker: TypeChecker, sourceFile: SourceFile, location: Node, codeActions?: CodeAction[], sourceDisplay?: SymbolDisplayPart[]): CompletionEntryDetails {
-        const { displayParts, documentation, symbolKind, tags } = SymbolDisplay.getSymbolDisplayPartsDocumentationAndSymbolKind(checker, symbol, sourceFile, location, location, SemanticMeaning.All);
+    function createCompletionDetailsForSymbol(symbol: Symbol, checker: TypeChecker, sourceFile: SourceFile, location: Node, cancellationToken: CancellationToken, codeActions?: CodeAction[], sourceDisplay?: SymbolDisplayPart[]): CompletionEntryDetails {
+        const { displayParts, documentation, symbolKind, tags } =
+            checker.runWithCancellationToken(cancellationToken, checker =>
+                SymbolDisplay.getSymbolDisplayPartsDocumentationAndSymbolKind(checker, symbol, sourceFile, location, location, SemanticMeaning.All)
+            );
         return createCompletionDetails(symbol.name, SymbolDisplay.getSymbolModifiers(symbol), symbolKind, displayParts, documentation, tags, codeActions, sourceDisplay);
     }
 
-    function stringLiteralCompletionDetails(name: string, location: Node, completion: StringLiteralCompletion, sourceFile: SourceFile, checker: TypeChecker): CompletionEntryDetails | undefined {
+    function stringLiteralCompletionDetails(name: string, location: Node, completion: StringLiteralCompletion, sourceFile: SourceFile, checker: TypeChecker, cancellationToken: CancellationToken): CompletionEntryDetails | undefined {
         switch (completion.kind) {
             case StringLiteralCompletionKind.Paths: {
                 const match = find(completion.paths, p => p.name === name);
@@ -587,7 +591,7 @@ namespace ts.Completions {
             }
             case StringLiteralCompletionKind.Properties: {
                 const match = find(completion.symbols, s => s.name === name);
-                return match && createCompletionDetailsForSymbol(match, checker, sourceFile, location);
+                return match && createCompletionDetailsForSymbol(match, checker, sourceFile, location, cancellationToken);
             }
             case StringLiteralCompletionKind.Types:
                 return find(completion.types, t => t.value === name) ? createCompletionDetails(name, ScriptElementKindModifier.none, ScriptElementKind.typeElement, [textPart(name)]) : undefined;

src/services/findAllReferences.ts

@@ -45,7 +45,10 @@ namespace ts.FindAllReferences {
         const checker = program.getTypeChecker();
         return !referencedSymbols || !referencedSymbols.length ? undefined : mapDefined<SymbolAndEntries, ReferencedSymbol>(referencedSymbols, ({ definition, references }) =>
             // Only include referenced symbols that have a valid definition.
-            definition && { definition: definitionToReferencedSymbolDefinitionInfo(definition, checker, node), references: references.map(toReferenceEntry) });
+            definition && {
+                definition: checker.runWithCancellationToken(cancellationToken, checker => definitionToReferencedSymbolDefinitionInfo(definition, checker, node)),
+                references: references.map(toReferenceEntry)
+            });
     }
 
     export function getImplementationsAtPosition(program: Program, cancellationToken: CancellationToken, sourceFiles: ReadonlyArray<SourceFile>, sourceFile: SourceFile, position: number): ImplementationLocation[] {

src/services/services.ts

@@ -1424,7 +1424,9 @@ namespace ts {
                 host,
                 formattingOptions && formatting.getFormatContext(formattingOptions),
                 getCanonicalFileName,
-                preferences);
+                preferences,
+                cancellationToken,
+            );
         }
 
         function getCompletionEntrySymbol(fileName: string, position: number, name: string, source?: string): Symbol {
@@ -1465,7 +1467,7 @@ namespace ts {
                             kind: ScriptElementKind.unknown,
                             kindModifiers: ScriptElementKindModifier.none,
                             textSpan: createTextSpanFromNode(node, sourceFile),
-                            displayParts: typeToDisplayParts(typeChecker, type, getContainerNode(node)),
+                            displayParts: typeChecker.runWithCancellationToken(cancellationToken, typeChecker => typeToDisplayParts(typeChecker, type, getContainerNode(node))),
                             documentation: type.symbol ? type.symbol.getDocumentationComment(typeChecker) : undefined,
                             tags: type.symbol ? type.symbol.getJsDocTags() : undefined
                         };
@@ -1474,7 +1476,9 @@ namespace ts {
                 return undefined;
             }
 
-            const { symbolKind, displayParts, documentation, tags } = SymbolDisplay.getSymbolDisplayPartsDocumentationAndSymbolKind(typeChecker, symbol, sourceFile, getContainerNode(node), node);
+            const { symbolKind, displayParts, documentation, tags } = typeChecker.runWithCancellationToken(cancellationToken, typeChecker =>
+                SymbolDisplay.getSymbolDisplayPartsDocumentationAndSymbolKind(typeChecker, symbol, sourceFile, getContainerNode(node), node)
+            );
             return {
                 kind: symbolKind,
                 kindModifiers: SymbolDisplay.getSymbolModifiers(symbol),

src/services/signatureHelp.ts

@@ -41,16 +41,16 @@ namespace ts.SignatureHelp {
             // We didn't have any sig help items produced by the TS compiler.  If this is a JS
             // file, then see if we can figure out anything better.
             if (isSourceFileJavaScript(sourceFile)) {
-                return createJavaScriptSignatureHelpItems(argumentInfo, program);
+                return createJavaScriptSignatureHelpItems(argumentInfo, program, cancellationToken);
             }
 
             return undefined;
         }
 
-        return createSignatureHelpItems(candidates, resolvedSignature, argumentInfo, typeChecker);
+        return typeChecker.runWithCancellationToken(cancellationToken, typeChecker => createSignatureHelpItems(candidates, resolvedSignature, argumentInfo, typeChecker));
     }
 
-    function createJavaScriptSignatureHelpItems(argumentInfo: ArgumentListInfo, program: Program): SignatureHelpItems {
+    function createJavaScriptSignatureHelpItems(argumentInfo: ArgumentListInfo, program: Program, cancellationToken: CancellationToken): SignatureHelpItems {
         if (argumentInfo.invocation.kind !== SyntaxKind.CallExpression) {
             return undefined;
         }
@@ -76,7 +76,7 @@ namespace ts.SignatureHelp {
                         if (type) {
                             const callSignatures = type.getCallSignatures();
                             if (callSignatures && callSignatures.length) {
-                                return createSignatureHelpItems(callSignatures, callSignatures[0], argumentInfo, typeChecker);
+                                return typeChecker.runWithCancellationToken(cancellationToken, typeChecker => createSignatureHelpItems(callSignatures, callSignatures[0], argumentInfo, typeChecker));
                             }
                         }
                     }

tests/baselines/reference/api/tsserverlibrary.d.ts

@@ -1842,6 +1842,12 @@ declare namespace ts {
         getSuggestionForNonexistentModule(node: Identifier, target: Symbol): string | undefined;
         getBaseConstraintOfType(type: Type): Type | undefined;
         getDefaultFromTypeParameter(type: Type): Type | undefined;
+        /**
+         * Depending on the operation performed, it may be appropriate to throw away the checker
+         * if the cancellation token is triggered. Typically, if it is used for error checking
+         * and the operation is cancelled, then it should be discarded, otherwise it is safe to keep.
+         */
+        runWithCancellationToken<T>(token: CancellationToken, cb: (checker: TypeChecker) => T): T;
     }
     enum NodeBuilderFlags {
         None = 0,

tests/baselines/reference/api/typescript.d.ts

@@ -1842,6 +1842,12 @@ declare namespace ts {
         getSuggestionForNonexistentModule(node: Identifier, target: Symbol): string | undefined;
         getBaseConstraintOfType(type: Type): Type | undefined;
         getDefaultFromTypeParameter(type: Type): Type | undefined;
+        /**
+         * Depending on the operation performed, it may be appropriate to throw away the checker
+         * if the cancellation token is triggered. Typically, if it is used for error checking
+         * and the operation is cancelled, then it should be discarded, otherwise it is safe to keep.
+         */
+        runWithCancellationToken<T>(token: CancellationToken, cb: (checker: TypeChecker) => T): T;
     }
     enum NodeBuilderFlags {
         None = 0,