Make sure LSP server is waited until stopped.

Author: heejaechang

src/client/activation/jedi/languageServerProxy.ts

@@ -23,12 +23,8 @@ import { traceDecoratorError, traceDecoratorVerbose, traceError } from '../../lo
 export class JediLanguageServerProxy implements ILanguageServerProxy {
     public languageClient: LanguageClient | undefined;
 
-    private languageServerTask: Promise<void> | undefined;
-
     private readonly disposables: Disposable[] = [];
 
-    private disposed = false;
-
     private lsVersion: string | undefined;
 
     constructor(
@@ -42,36 +38,9 @@ export class JediLanguageServerProxy implements ILanguageServerProxy {
         };
     }
 
-    @traceDecoratorVerbose('Stopping language server')
+    @traceDecoratorVerbose('Disposing language server')
     public dispose(): void {
-        if (this.languageClient) {
-            // eslint-disable-next-line @typescript-eslint/no-explicit-any
-            const pid: number | undefined = ((this.languageClient as any)._serverProcess as ChildProcess)?.pid;
-            const killServer = () => {
-                if (pid) {
-                    killPid(pid);
-                }
-            };
-
-            // Do not await on this.
-            this.languageClient.stop().then(
-                () => killServer(),
-                (ex) => {
-                    traceError('Stopping language client failed', ex);
-                    killServer();
-                },
-            );
-
-            this.languageClient = undefined;
-            this.languageServerTask = undefined;
-        }
-
-        while (this.disposables.length > 0) {
-            const d = this.disposables.shift()!;
-            d.dispose();
-        }
-
-        this.disposed = true;
+        this.stop().ignoreErrors();
     }
 
     @traceDecoratorError('Failed to start language server')
@@ -87,19 +56,41 @@ export class JediLanguageServerProxy implements ILanguageServerProxy {
         interpreter: PythonEnvironment | undefined,
         options: LanguageClientOptions,
     ): Promise<void> {
-        if (this.languageServerTask) {
-            await this.languageServerTask;
-            return;
-        }
-
         this.lsVersion =
             (options.middleware ? (<LanguageClientMiddleware>options.middleware).serverVersion : undefined) ?? '0.19.3';
 
         this.languageClient = await this.factory.createLanguageClient(resource, interpreter, options);
         this.registerHandlers();
 
-        this.languageServerTask = this.languageClient.start();
-        await this.languageServerTask;
+        await this.languageClient.start();
+    }
+
+    @traceDecoratorVerbose('Stopping language server')
+    public async stop(): Promise<void> {
+        if (this.languageClient) {
+            // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            const pid: number | undefined = ((this.languageClient as any)._serverProcess as ChildProcess)?.pid;
+            const killServer = () => {
+                if (pid) {
+                    killPid(pid);
+                }
+            };
+
+            try {
+                await this.languageClient.stop();
+                killServer();
+            } catch (ex) {
+                traceError('Stopping language client failed', ex);
+                killServer();
+            }
+
+            this.languageClient = undefined;
+        }
+
+        while (this.disposables.length > 0) {
+            const d = this.disposables.shift()!;
+            d.dispose();
+        }
     }
 
     // eslint-disable-next-line class-methods-use-this
@@ -115,11 +106,6 @@ export class JediLanguageServerProxy implements ILanguageServerProxy {
         JediLanguageServerProxy.versionTelemetryProps,
     )
     private registerHandlers() {
-        if (this.disposed) {
-            // Check if it got disposed in the interim.
-            return;
-        }
-
         const progressReporting = new ProgressReporting(this.languageClient!);
         this.disposables.push(progressReporting);
 

src/client/activation/jedi/manager.ts

@@ -59,9 +59,7 @@ export class JediLanguageServerManager implements ILanguageServerManager {
     }
 
     public dispose(): void {
-        if (this.languageProxy) {
-            this.languageProxy.dispose();
-        }
+        this.stopLanguageServer().ignoreErrors();
         JediLanguageServerManager.commandDispose.dispose();
         this.disposables.forEach((d) => d.dispose());
     }
@@ -120,9 +118,7 @@ export class JediLanguageServerManager implements ILanguageServerManager {
     @traceDecoratorError('Failed to restart language server')
     @traceDecoratorVerbose('Restarting language server')
     protected async restartLanguageServer(): Promise<void> {
-        if (this.languageProxy) {
-            this.languageProxy.dispose();
-        }
+        await this.stopLanguageServer();
         await this.startLanguageServer();
     }
 
@@ -147,4 +143,11 @@ export class JediLanguageServerManager implements ILanguageServerManager {
         // Then use this middleware to start a new language client.
         await this.languageServerProxy.start(this.resource, this.interpreter, options);
     }
+
+    @traceDecoratorVerbose('Stopping language server')
+    protected async stopLanguageServer(): Promise<void> {
+        if (this.languageServerProxy) {
+            await this.languageServerProxy.stop();
+        }
+    }
 }

src/client/activation/node/languageServerProxy.ts

@@ -10,7 +10,6 @@ import {
 } from 'vscode-languageclient/node';
 
 import { IExperimentService, IExtensions, IInterpreterPathService, Resource } from '../../common/types';
-import { noop } from '../../common/utils/misc';
 import { IEnvironmentVariablesProvider } from '../../common/variables/types';
 import { PythonEnvironment } from '../../pythonEnvironments/info';
 import { captureTelemetry } from '../../telemetry';
@@ -51,14 +50,10 @@ namespace GetExperimentValue {
 export class NodeLanguageServerProxy implements ILanguageServerProxy {
     public languageClient: LanguageClient | undefined;
 
-    private languageServerTask: Promise<void> | undefined;
-
     private cancellationStrategy: FileBasedCancellationStrategy | undefined;
 
     private readonly disposables: Disposable[] = [];
 
-    private disposed = false;
-
     private lsVersion: string | undefined;
 
     constructor(
@@ -76,24 +71,9 @@ export class NodeLanguageServerProxy implements ILanguageServerProxy {
         };
     }
 
-    @traceDecoratorVerbose('Stopping language server')
+    @traceDecoratorVerbose('Disposing language server')
     public dispose(): void {
-        if (this.languageClient) {
-            // Do not await on this.
-            this.languageClient.stop().then(noop, (ex) => traceError('Stopping language client failed', ex));
-
-            this.languageClient = undefined;
-            this.languageServerTask = undefined;
-        }
-        if (this.cancellationStrategy) {
-            this.cancellationStrategy.dispose();
-            this.cancellationStrategy = undefined;
-        }
-        while (this.disposables.length > 0) {
-            const d = this.disposables.shift()!;
-            d.dispose();
-        }
-        this.disposed = true;
+        this.stop().ignoreErrors();
     }
 
     @traceDecoratorError('Failed to start language server')
@@ -109,11 +89,6 @@ export class NodeLanguageServerProxy implements ILanguageServerProxy {
         interpreter: PythonEnvironment | undefined,
         options: LanguageClientOptions,
     ): Promise<void> {
-        if (this.languageServerTask) {
-            await this.languageServerTask;
-            return;
-        }
-
         const extension = this.extensions.getExtension(PYLANCE_EXTENSION_ID);
         this.lsVersion = extension?.packageJSON.version || '0';
 
@@ -129,8 +104,27 @@ export class NodeLanguageServerProxy implements ILanguageServerProxy {
             }),
         );
 
-        this.languageServerTask = this.languageClient.start();
-        await this.languageServerTask;
+        await this.languageClient.start();
+    }
+
+    @traceDecoratorVerbose('Disposing language server')
+    public async stop(): Promise<void> {
+        if (this.languageClient) {
+            try {
+                await this.languageClient.stop();
+            } catch (ex) {
+                traceError('Stopping language client failed', ex);
+            }
+            this.languageClient = undefined;
+        }
+        if (this.cancellationStrategy) {
+            this.cancellationStrategy.dispose();
+            this.cancellationStrategy = undefined;
+        }
+        while (this.disposables.length > 0) {
+            const d = this.disposables.shift()!;
+            d.dispose();
+        }
     }
 
     // eslint-disable-next-line class-methods-use-this
@@ -146,11 +140,6 @@ export class NodeLanguageServerProxy implements ILanguageServerProxy {
         NodeLanguageServerProxy.versionTelemetryProps,
     )
     private registerHandlers(_resource: Resource) {
-        if (this.disposed) {
-            // Check if it got disposed in the interim.
-            return;
-        }
-
         const progressReporting = new ProgressReporting(this.languageClient!);
         this.disposables.push(progressReporting);
 

src/client/activation/node/manager.ts

@@ -59,9 +59,7 @@ export class NodeLanguageServerManager implements ILanguageServerManager {
     }
 
     public dispose(): void {
-        if (this.languageProxy) {
-            this.languageProxy.dispose();
-        }
+        this.stopLanguageServer().ignoreErrors();
         NodeLanguageServerManager.commandDispose.dispose();
         this.disposables.forEach((d) => d.dispose());
     }
@@ -110,9 +108,7 @@ export class NodeLanguageServerManager implements ILanguageServerManager {
     @traceDecoratorError('Failed to restart language server')
     @traceDecoratorVerbose('Restarting language server')
     protected async restartLanguageServer(): Promise<void> {
-        if (this.languageProxy) {
-            this.languageProxy.dispose();
-        }
+        await this.stopLanguageServer();
         await this.startLanguageServer();
     }
 
@@ -137,4 +133,11 @@ export class NodeLanguageServerManager implements ILanguageServerManager {
         // Then use this middleware to start a new language client.
         await this.languageServerProxy.start(this.resource, this.interpreter, options);
     }
+
+    @traceDecoratorVerbose('Stopping language server')
+    protected async stopLanguageServer(): Promise<void> {
+        if (this.languageServerProxy) {
+            await this.languageServerProxy.stop();
+        }
+    }
 }

src/client/activation/types.ts

@@ -138,6 +138,7 @@ export interface ILanguageServerProxy extends IDisposable {
         interpreter: PythonEnvironment | undefined,
         options: LanguageClientOptions,
     ): Promise<void>;
+    stop(): Promise<void>;
     /**
      * Sends a request to LS so as to load other extensions.
      * This is used as a plugin loader mechanism.

src/client/languageServer/jediLSExtensionManager.ts

@@ -72,9 +72,9 @@ export class JediLSExtensionManager extends LanguageServerCapabilities
         this.serverManager.connect();
     }
 
-    stopLanguageServer(): void {
+    async stopLanguageServer(): Promise<void> {
         this.serverManager.disconnect();
-        this.serverProxy.dispose();
+        await this.serverProxy.stop();
     }
 
     // eslint-disable-next-line class-methods-use-this

src/client/languageServer/noneLSExtensionManager.ts

@@ -46,8 +46,8 @@ export class NoneLSExtensionManager implements ILanguageServer, ILanguageServerE
         return Promise.resolve();
     }
 
-    stopLanguageServer(): void {
-        // Nothing to do here.
+    stopLanguageServer(): Promise<void> {
+        return Promise.resolve();
     }
 
     canStartLanguageServer(): boolean {

src/client/languageServer/pylanceLSExtensionManager.ts

@@ -84,9 +84,9 @@ export class PylanceLSExtensionManager extends LanguageServerCapabilities
         this.serverManager.connect();
     }
 
-    stopLanguageServer(): void {
+    async stopLanguageServer(): Promise<void> {
         this.serverManager.disconnect();
-        this.serverProxy.dispose();
+        await this.serverProxy.stop();
     }
 
     canStartLanguageServer(): boolean {

src/client/languageServer/types.ts

@@ -29,7 +29,7 @@ export interface ILanguageServerCapabilities extends ILanguageServer {
  */
 export interface ILanguageServerExtensionManager extends ILanguageServerCapabilities {
     startLanguageServer(resource: Resource, interpreter?: PythonEnvironment): Promise<void>;
-    stopLanguageServer(): void;
+    stopLanguageServer(): Promise<void>;
     canStartLanguageServer(): boolean;
     languageServerNotAvailable(): Promise<void>;
     dispose(): void;

src/client/languageServer/watcher.ts

@@ -125,7 +125,7 @@ export class LanguageServerWatcher
 
         // Destroy the old language server if it's different.
         if (currentInterpreter && interpreter !== currentInterpreter) {
-            this.stopLanguageServer(lsResource);
+            await this.stopLanguageServer(lsResource);
         }
 
         // If the interpreter is Python 2 and the LS setting is explicitly set to Jedi, turn it off.
@@ -190,12 +190,12 @@ export class LanguageServerWatcher
 
     // Private methods
 
-    private stopLanguageServer(resource?: Resource): void {
+    private async stopLanguageServer(resource?: Resource): Promise<void> {
         const key = this.getWorkspaceKey(resource, this.languageServerType);
         const languageServerExtensionManager = this.workspaceLanguageServers.get(key);
 
         if (languageServerExtensionManager) {
-            languageServerExtensionManager.stopLanguageServer();
+            await languageServerExtensionManager.stopLanguageServer();
             languageServerExtensionManager.dispose();
             this.workspaceLanguageServers.delete(key);
         }
@@ -246,7 +246,7 @@ export class LanguageServerWatcher
         const languageServerType = this.configurationService.getSettings(lsResource).languageServer;
 
         if (languageServerType !== this.languageServerType) {
-            this.stopLanguageServer(resource);
+            await this.stopLanguageServer(resource);
             await this.startLanguageServer(languageServerType, lsResource);
         }
     }
@@ -286,9 +286,9 @@ export class LanguageServerWatcher
         // Since Jedi is the only language server type where we instantiate multiple language servers,
         // Make sure to dispose of them only in that scenario.
         if (event.removed.length && this.languageServerType === LanguageServerType.Jedi) {
-            event.removed.forEach((workspace) => {
-                this.stopLanguageServer(workspace.uri);
-            });
+            for (const workspace of event.removed) {
+                await this.stopLanguageServer(workspace.uri);
+            }
         }
     }
 

src/test/languageServer/watcher.unit.test.ts

@@ -723,7 +723,7 @@ suite('Language server watcher', () => {
             }
 
             const stopLanguageServerStub = sandbox.stub(extensionLSCls.prototype, 'stopLanguageServer');
-            stopLanguageServerStub.returns();
+            stopLanguageServerStub.returns(Promise.resolve());
 
             let onDidChangeWorkspaceFoldersListener: (event: WorkspaceFoldersChangeEvent) => Promise<void> = () =>
                 Promise.resolve();