Skip to content

Commit

Permalink
Add cancellation token
Browse files Browse the repository at this point in the history
  • Loading branch information
@msujew
msujew committed Jan 8, 2024
1 parent 607ad4f commit 1eadb4a
Show file tree
Hide file tree
Showing 4 changed files with 127 additions and 18 deletions.
34 changes: 24 additions & 10 deletions packages/langium/src/lsp/language-server.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -487,7 +487,7 @@ export function addWorkspaceSymbolHandler(connection: Connection, services: Lang
const documentBuilder = services.workspace.DocumentBuilder;
connection.onWorkspaceSymbol(async (params, token) => {
try {
await documentBuilder.waitUntil(DocumentState.IndexedContent);
await documentBuilder.waitUntil(DocumentState.IndexedContent, token);
return await workspaceSymbolProvider.getSymbols(params, token);
} catch (err) {
return responseError(err);
Expand All @@ -497,7 +497,7 @@ export function addWorkspaceSymbolHandler(connection: Connection, services: Lang
if (resolveWorkspaceSymbol) {
connection.onWorkspaceSymbolResolve(async (workspaceSymbol, token) => {
try {
await documentBuilder.waitUntil(DocumentState.IndexedContent);
await documentBuilder.waitUntil(DocumentState.IndexedContent, token);
return await resolveWorkspaceSymbol(workspaceSymbol, token);
} catch (err) {
return responseError(err);
Expand Down Expand Up @@ -586,9 +586,11 @@ export function createHierarchyRequestHandler<P extends TypeHierarchySupertypesP
sharedServices: LangiumSharedServices,
): ServerRequestHandler<P, R, PR, E> {
const serviceRegistry = sharedServices.ServiceRegistry;
const documentBuilder = sharedServices.workspace.DocumentBuilder;
return async (params: P, cancelToken: CancellationToken) => {
await documentBuilder.waitUntil(DocumentState.IndexedReferences);
const cancellationError = await waitUntilPhase<E>(sharedServices, cancelToken, DocumentState.IndexedReferences);
if (cancellationError) {
return cancellationError;
}
const uri = URI.parse(params.item.uri);
const language = serviceRegistry.getServices(uri);
if (!language) {
Expand All @@ -610,11 +612,11 @@ export function createServerRequestHandler<P extends { textDocument: TextDocumen
targetState?: DocumentState
): ServerRequestHandler<P, R, PR, E> {
const documents = sharedServices.workspace.LangiumDocuments;
const documentBuilder = sharedServices.workspace.DocumentBuilder;
const serviceRegistry = sharedServices.ServiceRegistry;
return async (params: P, cancelToken: CancellationToken) => {
if (targetState !== undefined) {
await documentBuilder.waitUntil(targetState);
const cancellationError = await waitUntilPhase<E>(sharedServices, cancelToken, targetState);
if (cancellationError) {
return cancellationError;
}
const uri = URI.parse(params.textDocument.uri);
const language = serviceRegistry.getServices(uri);
Expand All @@ -638,11 +640,11 @@ export function createRequestHandler<P extends { textDocument: TextDocumentIdent
targetState?: DocumentState
): RequestHandler<P, R | null, E> {
const documents = sharedServices.workspace.LangiumDocuments;
const documentBuilder = sharedServices.workspace.DocumentBuilder;
const serviceRegistry = sharedServices.ServiceRegistry;
return async (params: P, cancelToken: CancellationToken) => {
if (targetState !== undefined) {
await documentBuilder.waitUntil(targetState);
const cancellationError = await waitUntilPhase<E>(sharedServices, cancelToken, targetState);
if (cancellationError) {
return cancellationError;
}
const uri = URI.parse(params.textDocument.uri);
const language = serviceRegistry.getServices(uri);
Expand All @@ -662,6 +664,18 @@ export function createRequestHandler<P extends { textDocument: TextDocumentIdent
};
}

async function waitUntilPhase<E>(services: LangiumSharedServices, cancelToken: CancellationToken, targetState?: DocumentState): Promise<ResponseError<E> | undefined> {
if (targetState !== undefined) {
const documentBuilder = services.workspace.DocumentBuilder;
try {
await documentBuilder.waitUntil(targetState, cancelToken);
} catch (err) {
return responseError(err);
}
}
return undefined;
}

function responseError<E = void>(err: unknown): ResponseError<E> {
if (isOperationCancelled(err)) {
return new ResponseError(LSPErrorCodes.RequestCancelled, 'The request has been cancelled.');
Expand Down
24 changes: 18 additions & 6 deletions packages/langium/src/workspace/document-builder.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ import type { IndexManager } from '../workspace/index-manager.js';
import type { LangiumDocument, LangiumDocuments, LangiumDocumentFactory } from './documents.js';
import { CancellationToken, Disposable } from 'vscode-languageserver';
import { MultiMap } from '../utils/collections.js';
import { interruptAndCheck } from '../utils/promise-util.js';
import { OperationCancelled, interruptAndCheck } from '../utils/promise-util.js';
import { stream } from '../utils/stream.js';
import { ValidationCategory } from '../validation/validation-registry.js';
import { DocumentState } from './documents.js';
Expand Down Expand Up @@ -82,8 +82,12 @@ export interface DocumentBuilder {

/**
* Wait until the workspace has reached the specified state for all documents.
*
* @param state The desired state. The promise won't resolve until all documents have reached this state
* @param cancelToken Optionally allows to cancel the wait operation, disposing any listeners in the process
* @throws `OperationCancelled` if cancellation has been requested before the state has been reached
*/
waitUntil(state: DocumentState): Promise<void>;
waitUntil(state: DocumentState, cancelToken?: CancellationToken): Promise<void>;
}

export type DocumentUpdateListener = (changed: URI[], deleted: URI[]) => void | Promise<void>
Expand Down Expand Up @@ -306,15 +310,23 @@ export class DefaultDocumentBuilder implements DocumentBuilder {
});
}

waitUntil(state: DocumentState): Promise<void> {
waitUntil(state: DocumentState, cancelToken = CancellationToken.None): Promise<void> {
if (this.currentState >= state) {
return Promise.resolve();
} else if (cancelToken.isCancellationRequested) {
return Promise.reject(OperationCancelled);
}
return new Promise(resolve => {
const disposable = this.onBuildPhase(state, () => {
disposable.dispose();
return new Promise((resolve, reject) => {
const buildDisposable = this.onBuildPhase(state, () => {
buildDisposable.dispose();
cancelDisposable.dispose();
resolve();
});
const cancelDisposable = cancelToken.onCancellationRequested(() => {
buildDisposable.dispose();
cancelDisposable.dispose();
reject(OperationCancelled);
});
});
}

Expand Down
85 changes: 84 additions & 1 deletion packages/langium/test/workspace/document-builder.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,9 +8,10 @@ import type { AstNode, Reference, ValidationChecks } from 'langium';
import { describe, expect, test } from 'vitest';
import { CancellationTokenSource } from 'vscode-languageserver';
import { TextDocument } from 'vscode-languageserver-textdocument';
import { isOperationCancelled, DocumentState, EmptyFileSystem, URI } from 'langium';
import { isOperationCancelled, DocumentState, EmptyFileSystem, URI, delayNextTick } from 'langium';
import { createLangiumGrammarServices, createServicesForGrammar } from 'langium/grammar';
import { setTextDocument } from 'langium/test';
import { fail } from 'assert';

describe('DefaultDocumentBuilder', () => {
const grammarServices = createLangiumGrammarServices(EmptyFileSystem).grammar;
Expand Down Expand Up @@ -238,6 +239,88 @@ describe('DefaultDocumentBuilder', () => {
]);
});

test('waits until a specific workspace stage has been reached', async () => {
const services = await createServices();
const documentFactory = services.shared.workspace.LangiumDocumentFactory;
const documents = services.shared.workspace.LangiumDocuments;
const builder = services.shared.workspace.DocumentBuilder;
const document = documentFactory.fromString('', URI.parse('file:///test1.txt'));
documents.addDocument(document);

const actual: string[] = [];
const expected: string[] = [];
function wait(state: DocumentState): void {
expected.push('B' + state);
expected.push('W' + state);
builder.onBuildPhase(state, async () => {
actual.push('B' + state);
await delayNextTick();
});
builder.waitUntil(state).then(() => actual.push('W' + state));
}
for (let i = 2; i <= 6; i++) {
wait(i);
}
await builder.build([document], { validation: true });
expect(actual).toEqual(expected);
});

test('`waitUntil` will correctly wait even though the build process has been cancelled', async () => {
const services = await createServices();
const documentFactory = services.shared.workspace.LangiumDocumentFactory;
const documents = services.shared.workspace.LangiumDocuments;
const builder = services.shared.workspace.DocumentBuilder;
const document = documentFactory.fromString('', URI.parse('file:///test1.txt'));
documents.addDocument(document);

const actual: string[] = [];
const cancelTokenSource = new CancellationTokenSource();
function wait(state: DocumentState): void {
builder.onBuildPhase(state, async () => {
actual.push('B' + state);
await delayNextTick();
});
}
for (let i = 2; i <= 6; i++) {
wait(i);
}
builder.waitUntil(DocumentState.ComputedScopes).then(() => cancelTokenSource.cancel());
builder.waitUntil(DocumentState.IndexedReferences).then(() => {
actual.push('W' + DocumentState.IndexedReferences);
});
// Build twice but interrupt the first build after the computing scope phase
try {
await builder.build([document], { validation: true }, cancelTokenSource.token);
} catch {
// build has been cancelled, ignore
}
document.state = DocumentState.Parsed;
await builder.build([document], { validation: true });
// The B2 and B3 phases are duplicated because the first build has been cancelled
// W5 still appears as expected after B5
expect(actual).toEqual(['B2', 'B3', 'B2', 'B3', 'B4', 'B5', 'W5', 'B6']);
});

test('`waitUntil` can be cancelled before it gets triggered', async () => {
const services = await createServices();
const documentFactory = services.shared.workspace.LangiumDocumentFactory;
const documents = services.shared.workspace.LangiumDocuments;
const builder = services.shared.workspace.DocumentBuilder;
const document = documentFactory.fromString('', URI.parse('file:///test1.txt'));
documents.addDocument(document);

const cancelTokenSource = new CancellationTokenSource();
builder.waitUntil(DocumentState.IndexedReferences, cancelTokenSource.token).then(() => {
fail('This should have been cancelled');
}).catch(err => {
expect(isOperationCancelled(err)).toBeTruthy();
});
builder.onBuildPhase(DocumentState.ComputedScopes, () => {
cancelTokenSource.cancel();
});
await builder.build([document], { validation: true });
});

});

type TestAstType = {
Expand Down
2 changes: 1 addition & 1 deletion vite.config.mts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ export default defineConfig({
interopDefault: true
},
include: ['**/test/**/*.test.ts'],
exclude: ['**/node_modules/**', '**/dist/**', '**/generated/**', '**/templates/**'],
exclude: ['**/node_modules/**', '**/dist/**', '**/generated/**', '**/templates/**', '**/examples/hello*/**'],
watchExclude: [ '**/examples/hello*/**' /* populated by the yeoman generator test */],
}
});

0 comments on commit 1eadb4a

Please sign in to comment.