Dbaeumer/better-cougar-lime (#935)

* Fix race condition

* Handle errors correctly which show up because a pending response got rejected becase the connection got disposed

* Fix test

* Only allow send* calls when client is running or not started yet.

Author: dbaeumer

client-node-tests/src/integration.test.ts

@@ -1566,7 +1566,7 @@ suite('Server tests', () => {
 
 		await assert.rejects(async () => {
 			await client.stop();
-		}, /Connection got disposed/);
+		}, /Pending response rejected since connection got disposed/);
 		assert.strictEqual(client.needsStart(), true);
 		assert.strictEqual(client.needsStop(), false);
 

client/src/common/client.ts

@@ -34,7 +34,7 @@ import {
 	DocumentOnTypeFormattingRequest, RenameRequest, DocumentSymbolRequest, DocumentLinkRequest, DocumentColorRequest, DeclarationRequest, FoldingRangeRequest,
 	ImplementationRequest, SelectionRangeRequest, TypeDefinitionRequest, CallHierarchyPrepareRequest, SemanticTokensRegistrationType, LinkedEditingRangeRequest,
 	TypeHierarchyPrepareRequest, InlineValueRequest, InlayHintRequest, WorkspaceSymbolRequest, TextDocumentRegistrationOptions, FileOperationRegistrationOptions,
-	ConnectionOptions, PositionEncodingKind, DocumentDiagnosticRequest, NotebookDocumentSyncRegistrationType, NotebookDocumentSyncRegistrationOptions
+	ConnectionOptions, PositionEncodingKind, DocumentDiagnosticRequest, NotebookDocumentSyncRegistrationType, NotebookDocumentSyncRegistrationOptions, ErrorCodes
 } from 'vscode-languageserver-protocol';
 
 import * as c2p from './codeConverter';
@@ -431,6 +431,13 @@ export abstract class BaseLanguageClient implements FeatureClient<Middleware, La
 	private _clientOptions: ResolvedClientOptions;
 
 	private _state: ClientState;
+	private _onStart: Promise<void> | undefined;
+	private _onStop: Promise<void> | undefined;
+	private _connection: Connection | undefined;
+	private _idleInterval: Disposable | undefined;
+	private readonly _ignoredRegistrations: Set<string>;
+	// private _idleStart: number | undefined;
+
 	private readonly _notificationHandlers: Map<string, GenericNotificationHandler>;
 	private readonly _notificationDisposables: Map<string, Disposable>;
 	private readonly _pendingNotificationHandlers: Map<string, GenericNotificationHandler>;
@@ -441,12 +448,6 @@ export abstract class BaseLanguageClient implements FeatureClient<Middleware, La
 	private readonly _pendingProgressHandlers: Map<string | number,  { type: ProgressType<any>; handler: NotificationHandler<any> }>;
 	private readonly _progressDisposables: Map<string | number,  Disposable>;
 
-	private _onStart: Promise<void> | undefined;
-	private _onStop: Promise<void> | undefined;
-	private _connection: Connection | undefined;
-	private _idleInterval: Disposable | undefined;
-	// private _idleStart: number | undefined;
-
 	private _initializeResult: InitializeResult | undefined;
 	private _outputChannel: OutputChannel | undefined;
 	private _disposeOutputChannel: boolean;
@@ -510,6 +511,8 @@ export abstract class BaseLanguageClient implements FeatureClient<Middleware, La
 		this._clientOptions.synchronize = this._clientOptions.synchronize || {};
 
 		this._state = ClientState.Initial;
+		this._ignoredRegistrations = new Set();
+
 		this._notificationHandlers = new Map();
 		this._pendingNotificationHandlers = new Map();
 		this._notificationDisposables = new Map();
@@ -642,6 +645,9 @@ export abstract class BaseLanguageClient implements FeatureClient<Middleware, La
 	public sendRequest<R>(method: string, token?: CancellationToken): Promise<R>;
 	public sendRequest<R>(method: string, param: any, token?: CancellationToken): Promise<R>;
 	public async sendRequest<R>(type: string | MessageSignature, ...params: any[]): Promise<R> {
+		if (this.$state === ClientState.StartFailed || this.$state === ClientState.Stopping || this.$state === ClientState.Stopped) {
+			return Promise.reject(new ResponseError(ErrorCodes.ConnectionInactive, `Client is not running`));
+		}
 		try {
 			// Ensure we have a connection before we force the document sync.
 			const connection = await this.$start();
@@ -702,6 +708,9 @@ export abstract class BaseLanguageClient implements FeatureClient<Middleware, La
 	public sendNotification(method: string): Promise<void>;
 	public sendNotification(method: string, params: any): Promise<void>;
 	public async sendNotification<P>(type: string | MessageSignature, params?: P): Promise<void> {
+		if (this.$state === ClientState.StartFailed || this.$state === ClientState.Stopping || this.$state === ClientState.Stopped) {
+			return Promise.reject(new ResponseError(ErrorCodes.ConnectionInactive, `Client is not running`));
+		}
 		try {
 			// Ensure we have a connection before we force the document sync.
 			const connection = await this.$start();
@@ -925,12 +934,15 @@ export abstract class BaseLanguageClient implements FeatureClient<Middleware, La
 		if (this._onStart !== undefined) {
 			return this._onStart;
 		}
+		const [promise, resolve, reject] = this.createOnStartPromise();
+		this._onStart = promise;
+
+		// We are currently stopping the language client. Await the stop
+		// before continuing.
 		if (this._onStop !== undefined) {
 			await this._onStop;
 			this._onStop = undefined;
 		}
-		const [promise, resolve, reject] = this.createOnStartPromise();
-		this._onStart = promise;
 		// If we restart then the diagnostics collection is reused.
 		if (this._diagnostics === undefined) {
 			this._diagnostics = this._clientOptions.diagnosticCollectionName
@@ -1206,12 +1218,12 @@ export abstract class BaseLanguageClient implements FeatureClient<Middleware, La
 		return undefined;
 	}
 
-	public async stop(timeout: number = 2000): Promise<void> {
+	public stop(timeout: number = 2000): Promise<void> {
 		// Wait 2 seconds on stop
 		return this.shutdown('stop', timeout);
 	}
 
-	public async suspend(): Promise<void> {
+	public suspend(): Promise<void> {
 		// Wait 5 seconds on suspend.
 		return this.shutdown('suspend', 5000);
 	}
@@ -1259,6 +1271,7 @@ export abstract class BaseLanguageClient implements FeatureClient<Middleware, La
 			this._onStart = undefined;
 			this._onStop = undefined;
 			this._connection = undefined;
+			this._ignoredRegistrations.clear();
 		});
 	}
 
@@ -1699,6 +1712,15 @@ export abstract class BaseLanguageClient implements FeatureClient<Middleware, La
 	}
 
 	private async handleRegistrationRequest(params: RegistrationParams): Promise<void> {
+		// We will not receive a registration call before a client is running
+		// from a server. However if we stop or shutdown we might which might
+		// try to restart the server. So ignore registrations if we are not running
+		if (!this.isRunning()) {
+			for (const registration of params.registrations) {
+				this._ignoredRegistrations.add(registration.id);
+			}
+			return;
+		}
 		interface WithDocumentSelector {
 			documentSelector: DocumentSelector | undefined;
 		}
@@ -1723,6 +1745,9 @@ export abstract class BaseLanguageClient implements FeatureClient<Middleware, La
 
 	private async handleUnregistrationRequest(params: UnregistrationParams): Promise<void> {
 		for (let unregistration of params.unregisterations) {
+			if (this._ignoredRegistrations.has(unregistration.id)) {
+				continue;
+			}
 			const feature = this._dynamicFeatures.get(unregistration.method);
 			if (!feature) {
 				return Promise.reject(new Error(`No feature implementation for ${unregistration.method} found. Unregistration failed.`));
@@ -1771,6 +1796,11 @@ export abstract class BaseLanguageClient implements FeatureClient<Middleware, La
 	public handleFailedRequest<T>(type: MessageSignature, token: CancellationToken | undefined, error: any, defaultValue: T, showNotification: boolean = true): T {
 		// If we get a request cancel or a content modified don't log anything.
 		if (error instanceof ResponseError) {
+			// The connection got disposed while we were waiting for a response.
+			// Simply return the default value. Is the best we can do.
+			if (error.code === ErrorCodes.PendingResponseRejected || error.code === ErrorCodes.ConnectionInactive) {
+				return defaultValue;
+			}
 			if (error.code === LSPErrorCodes.RequestCancelled || error.code === LSPErrorCodes.ServerCancelled) {
 				if (token !== undefined && token.isCancellationRequested) {
 					return defaultValue;

client/src/common/diagnostic.ts

@@ -361,6 +361,9 @@ class DiagnosticRequestor implements Disposable {
 	}
 
 	public pull(document: TextDocument | Uri, cb?: () => void): void {
+		if (this.isDisposed) {
+			return;
+		}
 		const uri = document instanceof Uri ? document : document.uri;
 		this.pullAsync(document).then(() => {
 			if (cb) {
@@ -372,6 +375,9 @@ class DiagnosticRequestor implements Disposable {
 	}
 
 	private async pullAsync(document: TextDocument | Uri, version?: number | undefined): Promise<void> {
+		if (this.isDisposed) {
+			return;
+		}
 		const isUri = document instanceof Uri;
 		const uri = isUri ? document : document.uri;
 		const key = uri.toString();
@@ -456,6 +462,9 @@ class DiagnosticRequestor implements Disposable {
 	}
 
 	public pullWorkspace(): void {
+		if (this.isDisposed) {
+			return;
+		}
 		this.pullWorkspaceAsync().then(() => {
 			this.workspaceTimeout = RAL().timer.setTimeout(() => {
 				this.pullWorkspace();
@@ -474,7 +483,7 @@ class DiagnosticRequestor implements Disposable {
 	}
 
 	private async pullWorkspaceAsync(): Promise<void> {
-		if (!this.provider.provideWorkspaceDiagnostics) {
+		if (!this.provider.provideWorkspaceDiagnostics || this.isDisposed) {
 			return;
 		}
 		if (this.workspaceCancellation !== undefined) {
@@ -515,6 +524,9 @@ class DiagnosticRequestor implements Disposable {
 						textDocument: { uri: this.client.code2ProtocolConverter.asUri(document instanceof Uri ? document : document.uri) },
 						previousResultId: previousResultId
 					};
+					if (this.isDisposed === true || !this.client.isRunning()) {
+						return { kind: vsdiag.DocumentDiagnosticReportKind.full, items: [] };
+					}
 					return this.client.sendRequest(DocumentDiagnosticRequest.type, params, token).then(async (result) => {
 						if (result === undefined || result === null || this.isDisposed || token.isCancellationRequested) {
 							return { kind: vsdiag.DocumentDiagnosticReportKind.full, items: [] };
@@ -585,6 +597,9 @@ class DiagnosticRequestor implements Disposable {
 						previousResultIds: convertPreviousResultIds(resultIds),
 						partialResultToken: partialResultToken
 					};
+					if (this.isDisposed === true || !this.client.isRunning()) {
+						return { items: [] };
+					}
 					return this.client.sendRequest(WorkspaceDiagnosticRequest.type, params, token).then(async (result): Promise<vsdiag.WorkspaceDiagnosticReport> => {
 						if (token.isCancellationRequested) {
 							return { items: [] };
@@ -640,13 +655,22 @@ class BackgroundScheduler implements Disposable {
 	private endDocument: TextDocument | Uri | undefined;
 	private readonly documents: LinkedMap<string, TextDocument | Uri>;
 	private intervalHandle: Disposable | undefined;
+	// The problem is that there could be outstanding diagnostic requests
+	// when we shutdown which when we receive the result will trigger a
+	// reschedule. So we remember if the background scheduler got disposed
+	// and ignore those re-schedules
+	private isDisposed: boolean;
 
 	public constructor(diagnosticRequestor: DiagnosticRequestor) {
 		this.diagnosticRequestor = diagnosticRequestor;
 		this.documents = new LinkedMap();
+		this.isDisposed = false;
 	}
 
 	public add(document: TextDocument | Uri): void {
+		if (this.isDisposed === true) {
+			return;
+		}
 		const key = document instanceof Uri ? document.toString() : document.uri.toString();
 		if (this.documents.has(key)) {
 			return;
@@ -672,6 +696,9 @@ class BackgroundScheduler implements Disposable {
 	}
 
 	public trigger(): void {
+		if (this.isDisposed === true) {
+			return;
+		}
 		// We have a round running. So simply make sure we run up to the
 		// last document
 		if (this.intervalHandle !== undefined) {
@@ -693,6 +720,7 @@ class BackgroundScheduler implements Disposable {
 	}
 
 	public dispose(): void {
+		this.isDisposed = true;
 		this.stop();
 		this.documents.clear();
 	}
@@ -866,11 +894,10 @@ class DiagnosticFeatureProviderImpl implements DiagnosticProviderShape {
 
 export class DiagnosticFeature extends TextDocumentLanguageFeature<DiagnosticOptions, DiagnosticRegistrationOptions, DiagnosticProviderShape, DiagnosticProviderMiddleware, $DiagnosticPullOptions> {
 
-	private readonly tabs: Tabs;
+	private tabs: Tabs | undefined;
 
 	constructor(client: FeatureClient<DiagnosticProviderMiddleware, $DiagnosticPullOptions>) {
 		super(client, DocumentDiagnosticRequest.type);
-		this.tabs = new Tabs();
 	}
 
 	public fillClientCapabilities(capabilities: ClientCapabilities): void {
@@ -899,11 +926,17 @@ export class DiagnosticFeature extends TextDocumentLanguageFeature<DiagnosticOpt
 	}
 
 	public dispose(): void {
-		this.tabs.dispose();
+		if (this.tabs !== undefined) {
+			this.tabs.dispose();
+			this.tabs = undefined;
+		}
 		super.dispose();
 	}
 
 	protected registerLanguageProvider(options: DiagnosticRegistrationOptions): [Disposable, DiagnosticProviderShape] {
+		if (this.tabs === undefined) {
+			this.tabs = new Tabs();
+		}
 		const provider = new DiagnosticFeatureProviderImpl(this._client, this.tabs, options);
 		return [provider.disposable, provider];
 	}

client/src/common/fileOperations.ts

@@ -91,7 +91,6 @@ abstract class FileOperationFeature<I, E extends Event<I>> implements DynamicFea
 		assign(value, this._clientCapability, true);
 	}
 
-
 	public initialize(capabilities: proto.ServerCapabilities): void {
 		const options = capabilities.workspace?.fileOperations;
 		const capability = options !== undefined ? access(options, this._serverCapability) : undefined;

jsonrpc/src/common/connection.ts

@@ -1392,7 +1392,7 @@ export function createMessageConnection(messageReader: MessageReader, messageWri
 			}
 			state = ConnectionState.Disposed;
 			disposeEmitter.fire(undefined);
-			const error = new Error('Connection got disposed.');
+			const error = new ResponseError(ErrorCodes.PendingResponseRejected, 'Pending response rejected since connection got disposed');
 			for (const promise of responsePromises.values()) {
 				promise.reject(error);
 			}

jsonrpc/src/common/messages.ts

@@ -57,9 +57,27 @@ export namespace ErrorCodes {
 	/** @deprecated use  jsonrpcReservedErrorRangeStart */
 	export const serverErrorStart: -32099 = /* jsonrpcReservedErrorRangeStart */ -32099;
 
+	/**
+	 * An error occurred when write a message to the transport layer.
+	 */
 	export const MessageWriteError: -32099 = -32099;
+
+	/**
+	 * An error occurred when reading a message from the transport layer.
+	 */
 	export const MessageReadError: -32098 = -32098;
 
+	/**
+	 * The connection got disposed or lost and all pending responses got
+	 * rejected.
+	 */
+	export const PendingResponseRejected: -32097 = -32097;
+
+	/**
+	 * The connection is inactive and a use of it failed.
+	 */
+	export const ConnectionInactive: -32096 = -32096;
+
 	/**
 	 * Error code indicating that a server received a notification or
 	 * request before the server has received the `initialize` request.