perf: cancel diagnostics (#2216)

Another part of #2179 is that TypeScript took a long time to type-check the file. When the file is ts or check-js, it'll be getSemanticDiagnostics, and for non-checked js, it's getSuggestionDiagnostics. There is probably not much we can do about it. But we can instead cancel the check when the file is updated. The cancellation can only be checked if there is an async operation. I added a small delay in the typescript check so that the update notification queue might run in that time frame. And we could return early for the svelte and typescript diagnostic. CSS and HTML are pure sync code and generally not that heavy, so it's probably not worth the overhead.

Author: jasonlyu123

packages/language-server/src/lib/DiagnosticsManager.ts

@@ -1,9 +1,18 @@
-import { _Connection, TextDocumentIdentifier, Diagnostic } from 'vscode-languageserver';
+import {
+    Connection,
+    TextDocumentIdentifier,
+    Diagnostic,
+    CancellationTokenSource,
+    CancellationToken
+} from 'vscode-languageserver';
 import { DocumentManager, Document } from './documents';
 import { debounceThrottle } from '../utils';
 
-export type SendDiagnostics = _Connection['sendDiagnostics'];
-export type GetDiagnostics = (doc: TextDocumentIdentifier) => Thenable<Diagnostic[]>;
+export type SendDiagnostics = Connection['sendDiagnostics'];
+export type GetDiagnostics = (
+    doc: TextDocumentIdentifier,
+    cancellationToken?: CancellationToken
+) => Thenable<Diagnostic[]>;
 
 export class DiagnosticsManager {
     constructor(
@@ -13,6 +22,7 @@ export class DiagnosticsManager {
     ) {}
 
     private pendingUpdates = new Set<Document>();
+    private cancellationTokens = new Map<string, { cancel: () => void }>();
 
     private updateAll() {
         this.docManager.getAllOpenedByClient().forEach((doc) => {
@@ -21,14 +31,43 @@ export class DiagnosticsManager {
         this.pendingUpdates.clear();
     }
 
-    scheduleUpdateAll = debounceThrottle(() => this.updateAll(), 1000);
+    scheduleUpdateAll() {
+        this.cancellationTokens.forEach((token) => token.cancel());
+        this.cancellationTokens.clear();
+        this.pendingUpdates.clear();
+        this.debouncedUpdateAll();
+    }
+
+    private debouncedUpdateAll = debounceThrottle(() => this.updateAll(), 1000);
 
     private async update(document: Document) {
-        const diagnostics = await this.getDiagnostics({ uri: document.getURL() });
+        const uri = document.getURL();
+        this.cancelStarted(uri);
+
+        const tokenSource = new CancellationTokenSource();
+        this.cancellationTokens.set(uri, tokenSource);
+
+        const diagnostics = await this.getDiagnostics(
+            { uri: document.getURL() },
+            tokenSource.token
+        );
         this.sendDiagnostics({
             uri: document.getURL(),
             diagnostics
         });
+
+        tokenSource.dispose();
+
+        if (this.cancellationTokens.get(uri) === tokenSource) {
+            this.cancellationTokens.delete(uri);
+        }
+    }
+
+    cancelStarted(uri: string) {
+        const started = this.cancellationTokens.get(uri);
+        if (started) {
+            started.cancel();
+        }
     }
 
     removeDiagnostics(document: Document) {
@@ -44,6 +83,7 @@ export class DiagnosticsManager {
             return;
         }
 
+        this.cancelStarted(document.getURL());
         this.pendingUpdates.add(document);
         this.scheduleBatchUpdate();
     }
@@ -53,5 +93,5 @@ export class DiagnosticsManager {
             this.update(doc);
         });
         this.pendingUpdates.clear();
-    }, 750);
+    }, 700);
 }

packages/language-server/src/plugins/PluginHost.ts

@@ -76,7 +76,10 @@ export class PluginHost implements LSProvider, OnWatchFileChanges {
         this.deferredRequests = {};
     }
 
-    async getDiagnostics(textDocument: TextDocumentIdentifier): Promise<Diagnostic[]> {
+    async getDiagnostics(
+        textDocument: TextDocumentIdentifier,
+        cancellationToken?: CancellationToken
+    ): Promise<Diagnostic[]> {
         const document = this.getDocument(textDocument.uri);
 
         if (
@@ -96,7 +99,7 @@ export class PluginHost implements LSProvider, OnWatchFileChanges {
         return flatten(
             await this.execute<Diagnostic[]>(
                 'getDiagnostics',
-                [document],
+                [document, cancellationToken],
                 ExecuteMode.Collect,
                 'high'
             )

packages/language-server/src/plugins/svelte/SveltePlugin.ts

@@ -49,15 +49,19 @@ export class SveltePlugin
 
     constructor(private configManager: LSConfigManager) {}
 
-    async getDiagnostics(document: Document): Promise<Diagnostic[]> {
+    async getDiagnostics(
+        document: Document,
+        cancellationToken?: CancellationToken
+    ): Promise<Diagnostic[]> {
         if (!this.featureEnabled('diagnostics') || !this.configManager.getIsTrusted()) {
             return [];
         }
 
         return getDiagnostics(
             document,
             await this.getSvelteDoc(document),
-            this.configManager.getConfig().svelte.compilerWarnings
+            this.configManager.getConfig().svelte.compilerWarnings,
+            cancellationToken
         );
     }
 

packages/language-server/src/plugins/svelte/features/getDiagnostics.ts

@@ -1,6 +1,12 @@
 // @ts-ignore
 import { Warning } from 'svelte/types/compiler/interfaces';
-import { Diagnostic, DiagnosticSeverity, Position, Range } from 'vscode-languageserver';
+import {
+    CancellationToken,
+    Diagnostic,
+    DiagnosticSeverity,
+    Position,
+    Range
+} from 'vscode-languageserver';
 import {
     Document,
     isInTag,
@@ -19,15 +25,20 @@ import { SvelteDocument, TranspileErrorSource } from '../SvelteDocument';
 export async function getDiagnostics(
     document: Document,
     svelteDoc: SvelteDocument,
-    settings: CompilerWarningsSettings
+    settings: CompilerWarningsSettings,
+    cancellationToken?: CancellationToken
 ): Promise<Diagnostic[]> {
     const config = await svelteDoc.config;
     if (config?.loadConfigError) {
         return getConfigLoadErrorDiagnostics(config.loadConfigError);
     }
 
+    if (cancellationToken?.isCancellationRequested) {
+        return [];
+    }
+
     try {
-        return await tryGetDiagnostics(document, svelteDoc, settings);
+        return await tryGetDiagnostics(document, svelteDoc, settings, cancellationToken);
     } catch (error) {
         return getPreprocessErrorDiagnostics(document, error);
     }
@@ -39,12 +50,19 @@ export async function getDiagnostics(
 async function tryGetDiagnostics(
     document: Document,
     svelteDoc: SvelteDocument,
-    settings: CompilerWarningsSettings
+    settings: CompilerWarningsSettings,
+    cancellationToken: CancellationToken | undefined
 ): Promise<Diagnostic[]> {
     const transpiled = await svelteDoc.getTranspiled();
+    if (cancellationToken?.isCancellationRequested) {
+        return [];
+    }
 
     try {
         const res = await svelteDoc.getCompiled();
+        if (cancellationToken?.isCancellationRequested) {
+            return [];
+        }
         return (((res.stats as any)?.warnings || res.warnings || []) as Warning[])
             .filter((warning) => settings[warning.code] !== 'ignore')
             .map((warning) => {
@@ -65,7 +83,7 @@ async function tryGetDiagnostics(
             .map((diag) => adjustMappings(diag, document))
             .filter((diag) => isNoFalsePositive(diag, document));
     } catch (err) {
-        return (await createParserErrorDiagnostic(err, document))
+        return createParserErrorDiagnostic(err, document)
             .map((diag) => mapObjWithRangeToOriginal(transpiled, diag))
             .map((diag) => adjustMappings(diag, document));
     }
@@ -74,7 +92,7 @@ async function tryGetDiagnostics(
 /**
  * Try to infer a nice diagnostic error message from the compilation error.
  */
-async function createParserErrorDiagnostic(error: any, document: Document) {
+function createParserErrorDiagnostic(error: any, document: Document) {
     const start = error.start || { line: 1, column: 0 };
     const end = error.end || start;
     const diagnostic: Diagnostic = {

packages/language-server/src/plugins/typescript/features/DiagnosticsProvider.ts

@@ -80,11 +80,20 @@ export class DiagnosticsProviderImpl implements DiagnosticsProvider {
             ];
         }
 
-        let diagnostics: ts.Diagnostic[] = [
-            ...lang.getSyntacticDiagnostics(tsDoc.filePath),
-            ...lang.getSuggestionDiagnostics(tsDoc.filePath),
-            ...lang.getSemanticDiagnostics(tsDoc.filePath)
-        ];
+        let diagnostics: ts.Diagnostic[] = lang.getSyntacticDiagnostics(tsDoc.filePath);
+        const checkers = [lang.getSuggestionDiagnostics, lang.getSemanticDiagnostics];
+
+        for (const checker of checkers) {
+            if (cancellationToken) {
+                // wait a bit so the event loop can check for cancellation
+                // or let completion go first
+                await new Promise((resolve) => setTimeout(resolve, 10));
+                if (cancellationToken.isCancellationRequested) {
+                    return [];
+                }
+            }
+            diagnostics.push(...checker.call(lang, tsDoc.filePath));
+        }
 
         const additionalStoreDiagnostics: ts.Diagnostic[] = [];
         const notGenerated = isNotGenerated(tsDoc.getFullText());

packages/language-server/src/server.ts

@@ -353,6 +353,7 @@ export function startServer(options?: LSOptions) {
 
     connection.onDidCloseTextDocument((evt) => docManager.closeDocument(evt.textDocument.uri));
     connection.onDidChangeTextDocument((evt) => {
+        diagnosticsManager.cancelStarted(evt.textDocument.uri);
         docManager.updateDocument(evt.textDocument, evt.contentChanges);
         pluginHost.didUpdateDocument();
     });
@@ -468,7 +469,7 @@ export function startServer(options?: LSOptions) {
         refreshCrossFilesSemanticFeatures();
     }
 
-    connection.onDidSaveTextDocument(diagnosticsManager.scheduleUpdateAll);
+    connection.onDidSaveTextDocument(diagnosticsManager.scheduleUpdateAll.bind(diagnosticsManager));
     connection.onNotification('$/onDidChangeTsOrJsFile', async (e: any) => {
         const path = urlToPath(e.uri);
         if (path) {

packages/language-server/test/plugins/PluginHost.test.ts

@@ -52,7 +52,7 @@ describe('PluginHost', () => {
         await pluginHost.getDiagnostics(textDocument);
 
         sinon.assert.calledOnce(plugin.getDiagnostics);
-        sinon.assert.calledWithExactly(plugin.getDiagnostics, document);
+        sinon.assert.calledWithExactly(plugin.getDiagnostics, document, undefined);
     });
 
     it('executes doHover on plugins', async () => {