Perform read lock on LSP requests

packages/langium/src/lsp/language-server.ts

@@ -552,12 +552,14 @@ export function addCodeLensHandler(connection: Connection, services: LangiumShar
 
 export function addWorkspaceSymbolHandler(connection: Connection, services: LangiumSharedServices): void {
     const workspaceSymbolProvider = services.lsp.WorkspaceSymbolProvider;
+    const lock = services.workspace.WorkspaceLock;
     if (workspaceSymbolProvider) {
         const documentBuilder = services.workspace.DocumentBuilder;
         connection.onWorkspaceSymbol(async (params, token) => {
             try {
                 await documentBuilder.waitUntil(DocumentState.IndexedContent, token);
-                return await workspaceSymbolProvider.getSymbols(params, token);
+                const result = await lock.read(() => workspaceSymbolProvider.getSymbols(params, token), true);
+                return result;
             } catch (err) {
                 return responseError(err);
             }
@@ -567,7 +569,8 @@ export function addWorkspaceSymbolHandler(connection: Connection, services: Lang
             connection.onWorkspaceSymbolResolve(async (workspaceSymbol, token) => {
                 try {
                     await documentBuilder.waitUntil(DocumentState.IndexedContent, token);
-                    return await resolveWorkspaceSymbol(workspaceSymbol, token);
+                    const result = await lock.read(() => resolveWorkspaceSymbol(workspaceSymbol, token), true);
+                    return result;
                 } catch (err) {
                     return responseError(err);
                 }
@@ -654,6 +657,7 @@ export function createHierarchyRequestHandler<P extends TypeHierarchySupertypesP
     serviceCall: (services: LangiumCoreAndPartialLSPServices, params: P, cancelToken: CancellationToken) => HandlerResult<R, E>,
     sharedServices: LangiumSharedServices,
 ): ServerRequestHandler<P, R, PR, E> {
+    const lock = sharedServices.workspace.WorkspaceLock;
     const serviceRegistry = sharedServices.ServiceRegistry;
     return async (params: P, cancelToken: CancellationToken) => {
         const uri = URI.parse(params.item.uri);
@@ -668,7 +672,10 @@ export function createHierarchyRequestHandler<P extends TypeHierarchySupertypesP
         }
         const language = serviceRegistry.getServices(uri);
         try {
-            return await serviceCall(language, params, cancelToken);
+            const result = await lock.read(async () => {
+                return await serviceCall(language, params, cancelToken);
+            }, true); // Give this priority, since we already waited until the target state
+            return result;
         } catch (err) {
             return responseError<E>(err);
         }
@@ -681,6 +688,7 @@ export function createServerRequestHandler<P extends { textDocument: TextDocumen
     targetState?: DocumentState
 ): ServerRequestHandler<P, R, PR, E> {
     const documents = sharedServices.workspace.LangiumDocuments;
+    const lock = sharedServices.workspace.WorkspaceLock;
     const serviceRegistry = sharedServices.ServiceRegistry;
     return async (params: P, cancelToken: CancellationToken) => {
         const uri = URI.parse(params.textDocument.uri);
@@ -694,9 +702,17 @@ export function createServerRequestHandler<P extends { textDocument: TextDocumen
             return responseError<E>(new Error(errorText));
         }
         const language = serviceRegistry.getServices(uri);
+        const document = documents.getDocument(uri);
+        if (!document) {
+            const errorText = `Could not find document for uri: '${uri}'`;
+            console.debug(errorText);
+            return responseError<E>(new Error(errorText));
+        }
         try {
-            const document = await documents.getOrCreateDocument(uri);
-            return await serviceCall(language, document, params, cancelToken);
+            const result = await lock.read(async () => {
+                return await serviceCall(language, document, params, cancelToken);
+            }, true); // Give this priority, since we already waited until the target state
+            return result;
         } catch (err) {
             return responseError<E>(err);
         }
@@ -709,6 +725,7 @@ export function createRequestHandler<P extends { textDocument: TextDocumentIdent
     targetState?: DocumentState
 ): RequestHandler<P, R | null, E> {
     const documents = sharedServices.workspace.LangiumDocuments;
+    const lock = sharedServices.workspace.WorkspaceLock;
     const serviceRegistry = sharedServices.ServiceRegistry;
     return async (params: P, cancelToken: CancellationToken) => {
         const uri = URI.parse(params.textDocument.uri);
@@ -721,9 +738,17 @@ export function createRequestHandler<P extends { textDocument: TextDocumentIdent
             return null;
         }
         const language = serviceRegistry.getServices(uri);
+        const document = documents.getDocument(uri);
+        if (!document) {
+            const errorText = `Could not find document for uri: '${uri}'`;
+            console.debug(errorText);
+            return responseError<E>(new Error(errorText));
+        }
         try {
-            const document = await documents.getOrCreateDocument(uri);
-            return await serviceCall(language, document, params, cancelToken);
+            const result = await lock.read(async () => {
+                return await serviceCall(language, document, params, cancelToken);
+            }, true); // Give this priority, since we already waited until the target state
+            return result;
         } catch (err) {
             return responseError<E>(err);
         }

packages/langium/src/workspace/workspace-lock.ts

@@ -28,8 +28,13 @@ export interface WorkspaceLock {
      * If a new write action is queued up while a read action is waiting, the write action will receive priority and will be handled before the read action.
      *
      * Note that read actions are not allowed to modify anything in the workspace. Please use {@link write} instead.
+     *
+     * @param action the action to perform.
+     * @param priority
+     * If set to true, the action will receive immediate attention and will be instantly executed. All other queued read/write actions will be postponed.
+     * This should only be used for actions that require a specific document state and have already awaited the necessary write actions.
      */
-    read<T>(action: () => MaybePromise<T>): Promise<T>;
+    read<T>(action: () => MaybePromise<T>, priority?: boolean): Promise<T>;
 
     /**
      * Cancels the last queued write action. All previous write actions already have been cancelled.
@@ -50,7 +55,7 @@ export class DefaultWorkspaceLock implements WorkspaceLock {
     private previousTokenSource = new CancellationTokenSource();
     private writeQueue: LockEntry[] = [];
     private readQueue: LockEntry[] = [];
-    private done = true;
+    private counter = 0;
 
     write(action: (token: CancellationToken) => MaybePromise<void>): Promise<void> {
         this.cancelWrite();
@@ -59,8 +64,25 @@ export class DefaultWorkspaceLock implements WorkspaceLock {
         return this.enqueue(this.writeQueue, action, tokenSource.token);
     }
 
-    read<T>(action: () => MaybePromise<T>): Promise<T> {
-        return this.enqueue(this.readQueue, action);
+    read<T>(action: () => MaybePromise<T>, priority?: boolean): Promise<T> {
+        if (priority) {
+            this.counter++;
+            const deferred = new Deferred<T>();
+            const end = (run: () => void) => {
+                run();
+                // Ensure that in any case the counter is decremented and the next operation is performed
+                this.counter--;
+                this.performNextOperation();
+            };
+            // Instanly run the action and resolve/reject the promise afterwards
+            Promise.resolve(action()).then(
+                result => end(() => deferred.resolve(result)),
+                err => end(() => deferred.reject(err))
+            );
+            return deferred.promise;
+        } else {
+            return this.enqueue(this.readQueue, action);
+        }
     }
 
     private enqueue<T = void>(queue: LockEntry[], action: LockAction<T>, cancellationToken = CancellationToken.None): Promise<T> {
@@ -76,7 +98,7 @@ export class DefaultWorkspaceLock implements WorkspaceLock {
     }
 
     private async performNextOperation(): Promise<void> {
-        if (!this.done) {
+        if (this.counter > 0) {
             return;
         }
         const entries: LockEntry[] = [];
@@ -89,7 +111,7 @@ export class DefaultWorkspaceLock implements WorkspaceLock {
         } else {
             return;
         }
-        this.done = false;
+        this.counter += entries.length;
         await Promise.all(entries.map(async ({ action, deferred, cancellationToken }) => {
             try {
                 // Move the execution of the action to the next event loop tick via `Promise.resolve()`
@@ -104,7 +126,7 @@ export class DefaultWorkspaceLock implements WorkspaceLock {
                 }
             }
         }));
-        this.done = true;
+        this.counter -= entries.length;
         this.performNextOperation();
     }
 

packages/langium/test/workspace/workspace-lock.test.ts

@@ -61,8 +61,8 @@ describe('WorkspaceLock', () => {
         const mutex = new DefaultWorkspaceLock();
         const now = Date.now();
         const magicalNumber = await mutex.read(() => new Promise(resolve => setTimeout(() => resolve(42), 10)));
-        // Confirm that at least 10ms have elapsed
-        expect(Date.now() - now).toBeGreaterThanOrEqual(10);
+        // Confirm that at least a few milliseconds have passed
+        expect(Date.now() - now).toBeGreaterThanOrEqual(5);
         // Confirm the returned value
         expect(magicalNumber).toBe(42);
     });
@@ -112,4 +112,23 @@ describe('WorkspaceLock', () => {
         // and the second action decreases the value again
         expect(counter).toBe(0);
     });
+
+    test('Read actions can receive priority', async () => {
+        let counter = 0;
+        const mutex = new DefaultWorkspaceLock();
+        mutex.write(async () => {
+            await delayNextTick();
+            // Set counter to 1
+            counter = 1;
+        });
+        await mutex.read(() => {
+            // Set counter to 5
+            counter = 5;
+        }, true);
+        // Assert that the read action has received priority
+        expect(counter).toBe(5);
+        await delayNextTick();
+        // Assert that the write action has been successfully finished afterwards
+        expect(counter).toBe(1);
+    });
 });