Wait for a didChange rather than polling

Author: davidwengier

src/lsptoolshost/razor/htmlDocument.ts

@@ -10,7 +10,6 @@ export class HtmlDocument {
     public readonly path: string;
     private content = '';
     private checksum = '';
-    private previousVersion = -1;
 
     public constructor(public readonly uri: vscode.Uri, checksum: string) {
         this.path = getUriPath(uri);
@@ -25,21 +24,8 @@ export class HtmlDocument {
         return this.checksum;
     }
 
-    public async setContent(checksum: string, content: string) {
-        const document = await vscode.workspace.openTextDocument(this.uri);
-        // Capture the version _before_ the change, so we can know for sure if it's been seen
-        this.previousVersion = document.version;
+    public setContent(checksum: string, content: string) {
         this.checksum = checksum;
         this.content = content;
     }
-
-    public async waitForBufferUpdate() {
-        const document = await vscode.workspace.openTextDocument(this.uri);
-
-        // Wait for VS Code to process any previous content change. We don't care about finding
-        // a specific version, just that it's moved on from the previous one.
-        while (document.version === this.previousVersion) {
-            await new Promise((resolve) => setTimeout(resolve, 5));
-        }
-    }
 }

src/lsptoolshost/razor/htmlDocumentManager.ts

@@ -17,6 +17,13 @@ import { UriConverter } from '../utils/uriConverter';
 export class HtmlDocumentManager {
     private readonly htmlDocuments: { [hostDocumentPath: string]: HtmlDocument } = {};
     private readonly contentProvider: HtmlDocumentContentProvider;
+    private readonly pendingUpdates: {
+        [documentPath: string]: {
+            promise: Promise<void>;
+            resolve: () => void;
+            reject: (error: any) => void;
+        };
+    } = {};
 
     private readonly razorDocumentClosedRequest: RequestType<TextDocumentIdentifier, void, Error> = new RequestType(
         'razor/documentClosed'
@@ -35,6 +42,20 @@ export class HtmlDocumentManager {
     }
 
     public register() {
+        const didChangeRegistration = vscode.workspace.onDidChangeTextDocument((e) => {
+            // Check if this document is being monitored for updates
+            if (e.document.uri.scheme === HtmlDocumentContentProvider.scheme) {
+                const documentPath = getUriPath(e.document.uri);
+                const pendingUpdate = this.pendingUpdates[documentPath];
+
+                if (pendingUpdate) {
+                    // Document has been updated, resolve the promise
+                    pendingUpdate.resolve();
+                    delete this.pendingUpdates[documentPath];
+                }
+            }
+        });
+
         const didCloseRegistration = vscode.workspace.onDidCloseTextDocument(async (document) => {
             // We log when a virtual document is closed just in case it helps track down future bugs
             if (document.uri.scheme === HtmlDocumentContentProvider.scheme) {
@@ -63,7 +84,7 @@ export class HtmlDocumentManager {
             this.contentProvider
         );
 
-        return vscode.Disposable.from(didCloseRegistration, providerRegistration);
+        return vscode.Disposable.from(didChangeRegistration, didCloseRegistration, providerRegistration);
     }
 
     public async updateDocumentText(uri: vscode.Uri, checksum: string, text: string) {
@@ -79,7 +100,20 @@ export class HtmlDocumentManager {
 
         this.logger.logTrace(`New content for '${uri}', updating '${document.path}', checksum '${checksum}'.`);
 
-        await document.setContent(checksum, text);
+        // Create a promise for this document update
+        let resolve: () => void;
+        let reject: (error: any) => void;
+        const updatePromise = new Promise<void>((res, rej) => {
+            resolve = res;
+            reject = rej;
+        });
+
+        this.pendingUpdates[document.path] = { promise: updatePromise, resolve: resolve!, reject: reject! };
+
+        // Now update the document and fire the change event so VS Code will inform the Html language client.
+        await vscode.workspace.openTextDocument(document.uri);
+
+        document.setContent(checksum, text);
 
         this.contentProvider.fireDidChange(document.uri);
     }
@@ -90,6 +124,13 @@ export class HtmlDocumentManager {
         if (document) {
             this.logger.logTrace(`Removing '${document.uri}' from the document manager.`);
 
+            // Clean up any pending update promises for this document
+            const pendingUpdate = this.pendingUpdates[document.path];
+            if (pendingUpdate) {
+                pendingUpdate.reject(new Error('Document was closed before update completed'));
+                delete this.pendingUpdates[document.path];
+            }
+
             delete this.htmlDocuments[document.path];
         }
     }
@@ -109,10 +150,35 @@ export class HtmlDocumentManager {
             return undefined;
         }
 
-        // No checksum, just give them the latest document and hope they know what to do with it.
-
         await vscode.workspace.openTextDocument(document.uri);
 
+        if (checksum) {
+            // If checksum is supplied, that means we're getting this document because we're about to call an LSP method
+            // on it. We know that we've got the right document, and we've been told about the right content by the server,
+            // but all we can be sure of at this point is that we've fired the change event for it. The event firing
+            // is async, and the didChange notification that it would generate is a notification, so doesn't necessarily
+            // block. Before we actually make a call to the Html server, we should at least make sure that the document
+            // update has been seen by VS Code. We can't get access to the Html language client specifically to check if it
+            // has seen it, but we can trust that ordering will be preserved at least.
+            const pendingUpdate = this.pendingUpdates[document.path];
+            if (pendingUpdate) {
+                try {
+                    // Wait for the update promise with a 5 second timeout
+                    await Promise.race([
+                        pendingUpdate.promise,
+                        new Promise((_, reject) =>
+                            setTimeout(() => reject(new Error('Document update timeout')), 5000)
+                        ),
+                    ]);
+                } catch (error) {
+                    this.logger.logWarning(`Failed to wait for document update: ${error}`);
+                } finally {
+                    // Clean up the promise reference
+                    delete this.pendingUpdates[document.path];
+                }
+            }
+        }
+
         return document;
     }
 

src/lsptoolshost/razor/razorEndpoints.ts

@@ -252,13 +252,6 @@ export function registerRazorEndpoints(
                 return undefined;
             }
 
-            // We know that we've got the right document, and we've been told about the right content by the server,
-            // but all we can be sure of at this point is that we've fired the change event for it. The event firing
-            // is async, and the didChange notification that it would generate is a notification, so doesn't necessarily
-            // block. Before we actually make a call to the Html server, we should at least make sure that the document
-            // version is not still the same as before we updated the content.
-            await document.waitForBufferUpdate();
-
             return invocation(document, params.request);
         });
     }