Merged PR posit-dev/positron-python#108: Safely shutdown LSP on language runtime shutdown or restart

petetronic · wesm · commit 9e90676b7255 · 2024-03-28T15:04:00.000-05:00
Merge pull request #108 from posit-dev/feature/631 Safely shutdown LSP on language runtime shutdown or restart -------------------- Commit message for posit-dev/positron-python@34eb739: Safely shutdown LSP on language runtime shutdown or restart Also refactor LSP client related code into the recently added pythonLanguageRuntime.ts file. Authored-by: Pete Farland <pete.farland@posit.co> Signed-off-by: Pete Farland <pete.farland@posit.co>
diff --git a/extensions/positron-python/src/client/activation/jedi/positronLanguageRuntimes.ts b/extensions/positron-python/src/client/activation/jedi/positronLanguageRuntimes.ts
@@ -1,29 +1,26 @@
-/* eslint-disable @typescript-eslint/no-explicit-any */
 /* eslint-disable class-methods-use-this */
 /* eslint-disable global-require */
 /*---------------------------------------------------------------------------------------------
  *  Copyright (C) 2023 Posit Software, PBC. All rights reserved.
  *--------------------------------------------------------------------------------------------*/
 
-import { Socket } from 'net';
 import * as path from 'path';
 import * as fs from 'fs';
 // eslint-disable-next-line import/no-unresolved
 import * as positron from 'positron';
 import * as vscode from 'vscode';
 import * as crypto from 'crypto';
-import { Disposable, DocumentFilter, LanguageClient, LanguageClientOptions, ServerOptions, StreamInfo } from 'vscode-languageclient/node';
+import { Disposable, DocumentFilter, LanguageClientOptions } from 'vscode-languageclient/node';
 
 import { compare } from 'semver';
 import { EXTENSION_ROOT_DIR, PYTHON_LANGUAGE } from '../../common/constants';
 import { IConfigurationService, IInstaller, InstallerResponse, Product, Resource } from '../../common/types';
 import { InstallOptions } from '../../common/installer/types';
 import { IInterpreterService } from '../../interpreter/contracts';
 import { IServiceContainer } from '../../ioc/types';
-import { traceError, traceVerbose } from '../../logging';
+import { traceVerbose } from '../../logging';
 import { PythonEnvironment } from '../../pythonEnvironments/info';
 import { PythonVersion } from '../../pythonEnvironments/info/pythonVersion';
-import { ProgressReporting } from '../progress';
 import { ILanguageServerProxy } from '../types';
 import { PythonLanguageRuntime } from './pythonLanguageRuntime';
 import { JupyterAdapterApi } from '../../jupyter-adapter.d'
@@ -42,8 +39,6 @@ export class PositronJediLanguageServerProxy implements ILanguageServerProxy {
 
     private readonly disposables: Disposable[] = [];
 
-    private readonly languageClients: LanguageClient[] = [];
-
     private extensionVersion: string | undefined;
 
     private readonly installer: IInstaller;
@@ -119,16 +114,6 @@ export class PositronJediLanguageServerProxy implements ILanguageServerProxy {
             const r = this.disposables.shift()!;
             r.dispose();
         }
-
-        // Dispose of any language clients
-        for (const client of this.languageClients) {
-            try {
-                await client.stop();
-                await client.dispose();
-            } catch (ex) {
-                traceError('Stopping language client failed', ex);
-            }
-        }
     }
 
     public dispose(): void {
@@ -242,17 +227,7 @@ export class PositronJediLanguageServerProxy implements ILanguageServerProxy {
 
         // Create an adapter for the kernel as our language runtime
         const runtime = new PythonLanguageRuntime(
-            kernelSpec, metadata, jupyterAdapterApi);
-
-        runtime.onDidChangeRuntimeState((state) => {
-            if (state === positron.RuntimeState.Ready) {
-                jupyterAdapterApi.findAvailablePort([], 25).then((port: number) => {
-                    runtime.emitJupyterLog(`Starting Positron LSP server on port ${port}`);
-                    runtime.startPositronLsp(`127.0.0.1:${port}`);
-                    this.startClient(options, port).ignoreErrors();
-                });
-            }
-        });
+            kernelSpec, metadata, jupyterAdapterApi, options);
 
         // Register our language runtime provider
         return positron.runtime.registerLanguageRuntime(runtime);
@@ -295,77 +270,7 @@ export class PositronJediLanguageServerProxy implements ILanguageServerProxy {
         return `${info.major}.${info.minor}.${info.patch}`;
     }
 
-    /**
-     * Start the language client
-     */
-    private async startClient(clientOptions: LanguageClientOptions, port: number): Promise<void> {
-
-        // Configure language client to connect to LSP via TCP on start
-        const serverOptions: ServerOptions = async () => this.getServerOptions(port);
-        const client = new LanguageClient(PYTHON_LANGUAGE, 'Positron Python Jedi', serverOptions, clientOptions);
-        this.registerHandlers(client);
-        await client.start();
-        this.languageClients.push(client);
-    }
-
-    /**
-     * An async function used by the LanguageClient to establish a connection to the LSP on start.
-     * Several attempts to connect are made given recently spawned servers may not be ready immediately
-     * for client connections.
-     * @param port the LSP port
-     */
-    private async getServerOptions(port: number): Promise<StreamInfo> {
-
-        const delay = (ms: number) => new Promise(resolve => setTimeout(resolve, ms));
-        const maxAttempts = 20;
-        const baseDelay = 50;
-        const multiplier = 1.5;
-
-        for (let attempt = 0; attempt < maxAttempts; attempt += 1) {
-            // Retry up to five times then start to back-off
-            const interval = attempt < 6 ? baseDelay : baseDelay * multiplier * attempt;
-            if (attempt > 0) {
-                await delay(interval);
-            }
-
-            try {
-                // Try to connect to LSP port
-                const socket: Socket = await this.tryToConnect(port);
-                return { reader: socket, writer: socket };
-            } catch (error: any) {
-                if (error?.code === 'ECONNREFUSED') {
-                    traceVerbose(`Error '${error.message}' on connection attempt '${attempt}' to Jedi LSP on port '${port}', will retry`);
-                } else {
-                    throw error;
-                }
-            }
-        }
-
-        throw new Error(`Failed to create TCP connection to Jedi LSP on port ${port} after multiple attempts`);
-    }
-
-    /**
-     * Attempts to establish a TCP socket connection to the given port
-     * @param port the server port to connect to
-     */
-    private async tryToConnect(port: number): Promise<Socket> {
-        return new Promise((resolve, reject) => {
-            const socket = new Socket();
-            socket.on('ready', () => {
-                resolve(socket);
-            });
-            socket.on('error', (error) => {
-                reject(error);
-            });
-            socket.connect(port);
-        });
-    }
-
-    private registerHandlers(client: LanguageClient) {
-        const progressReporting = new ProgressReporting(client);
-        this.disposables.push(progressReporting);
-    }
-
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
     private withActiveExtention(ext: vscode.Extension<any>, callback: () => void) {
         if (ext.isActive) {
             callback();
diff --git a/extensions/positron-python/src/client/activation/jedi/pythonLanguageRuntime.ts b/extensions/positron-python/src/client/activation/jedi/pythonLanguageRuntime.ts
@@ -1,21 +1,31 @@
+/* eslint-disable class-methods-use-this */
 /*---------------------------------------------------------------------------------------------
  *  Copyright (C) 2023 Posit Software, PBC. All rights reserved.
  *--------------------------------------------------------------------------------------------*/
 
+import { Socket } from 'net';
+
 // eslint-disable-next-line import/no-unresolved
 import * as positron from 'positron';
 import * as vscode from 'vscode';
+import { Disposable, LanguageClient, LanguageClientOptions, ServerOptions, StreamInfo } from 'vscode-languageclient/node';
+import { PYTHON_LANGUAGE } from '../../common/constants';
 import { JupyterAdapterApi, JupyterKernelSpec, JupyterLanguageRuntime } from '../../jupyter-adapter.d';
+import { ProgressReporting } from '../progress';
 
 /**
  * A Positron Python language runtime that wraps a Jupyter kernel runtime.
  * Fulfills most Language Runtime API calls by delegating to the wrapped kernel.
  */
-export class PythonLanguageRuntime implements JupyterLanguageRuntime {
+export class PythonLanguageRuntime implements JupyterLanguageRuntime, Disposable {
 
     /** The Jupyter kernel-based implementation of the Language Runtime API */
     private _kernel: JupyterLanguageRuntime;
 
+    private _lsp: LanguageClient | undefined;
+
+    private readonly disposables: Disposable[] = [];
+
     /**
      * Create a new PythonLanguageRuntime object to wrap a Jupyter kernel.
      *
@@ -26,12 +36,17 @@ export class PythonLanguageRuntime implements JupyterLanguageRuntime {
     constructor(
         readonly kernelSpec: JupyterKernelSpec,
         readonly metadata: positron.LanguageRuntimeMetadata,
-        readonly adapterApi: JupyterAdapterApi) {
+        readonly adapterApi: JupyterAdapterApi,
+        readonly languageClientOptions: LanguageClientOptions) {
 
         this._kernel = adapterApi.adaptKernel(kernelSpec, metadata);
 
         this.onDidChangeRuntimeState = this._kernel.onDidChangeRuntimeState;
         this.onDidReceiveRuntimeMessage = this._kernel.onDidReceiveRuntimeMessage;
+
+        this.onDidChangeRuntimeState((state) => {
+            this.onStateChange(state);
+        });
     }
 
     startPositronLsp(clientAddress: string): void {
@@ -85,10 +100,115 @@ export class PythonLanguageRuntime implements JupyterLanguageRuntime {
     }
 
     async restart(): Promise<void> {
+
+        // Stop the LSP (it will restart after the kernel restarts)
+        if (this._lsp) {
+            await this._lsp.stop();
+        }
+
+        // Restart the kernel
         return this._kernel.restart();
     }
 
     async shutdown(): Promise<void> {
+
+        // Stop the LSP
+        if (this._lsp) {
+            await this._lsp.stop();
+        }
+
+        // Shutdown the kernel
         return this._kernel.shutdown();
     }
+
+    dispose(): void {
+        while (this.disposables.length > 0) {
+            const d = this.disposables.shift()!;
+            d.dispose();
+        }
+    }
+
+    private onStateChange(state: positron.RuntimeState): void {
+        if (state === positron.RuntimeState.Ready) {
+            this.adapterApi.findAvailablePort([], 25).then((port: number) => {
+                this.emitJupyterLog(`Starting Positron LSP server on port ${port}`);
+                this.startPositronLsp(`127.0.0.1:${port}`);
+                this.startLSPClient(port).ignoreErrors();
+            });
+        } else if (state === positron.RuntimeState.Exiting ||
+            state === positron.RuntimeState.Exited) {
+            if (this._lsp) {
+                this._lsp.stop().ignoreErrors();
+            }
+        }
+    }
+
+    // LSP Client Helpers
+
+    private async startLSPClient(port: number): Promise<void> {
+
+        const serverOptions: ServerOptions = async () => this.getLSPServerOptions(port);
+        const client = new LanguageClient(PYTHON_LANGUAGE, 'Positron Python Jedi', serverOptions, this.languageClientOptions);
+        this.disposables.push(client);
+
+        const progressReporting = new ProgressReporting(client);
+        this.disposables.push(progressReporting);
+
+        await client.start();
+        this._lsp = client;
+    }
+
+    /**
+     * An async function used by the LanguageClient to establish a connection to the LSP on start.
+     * Several attempts to connect are made given recently spawned servers may not be ready immediately
+     * for client connections.
+     * @param port the LSP port
+     */
+    private async getLSPServerOptions(port: number): Promise<StreamInfo> {
+
+        const delay = (ms: number) => new Promise(resolve => setTimeout(resolve, ms));
+        const maxAttempts = 20;
+        const baseDelay = 50;
+        const multiplier = 1.5;
+
+        for (let attempt = 0; attempt < maxAttempts; attempt += 1) {
+            // Retry up to five times then start to back-off
+            const interval = attempt < 6 ? baseDelay : baseDelay * multiplier * attempt;
+            if (attempt > 0) {
+                await delay(interval);
+            }
+
+            try {
+                // Try to connect to LSP port
+                const socket: Socket = await this.tryToConnect(port);
+                return { reader: socket, writer: socket };
+                // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            } catch (error: any) {
+                if (error?.code === 'ECONNREFUSED') {
+                    this.emitJupyterLog(`Error '${error.message}' on connection attempt '${attempt}' to Jedi LSP on port '${port}', will retry`);
+                } else {
+                    throw error;
+                }
+            }
+        }
+
+        throw new Error(`Failed to create TCP connection to Jedi LSP on port ${port} after multiple attempts`);
+    }
+
+    /**
+     * Attempts to establish a TCP socket connection to the given port
+     * @param port the server port to connect to
+     */
+    private async tryToConnect(port: number): Promise<Socket> {
+        return new Promise((resolve, reject) => {
+            const socket = new Socket();
+            socket.on('ready', () => {
+                resolve(socket);
+            });
+            socket.on('error', (error) => {
+                reject(error);
+            });
+            socket.connect(port);
+        });
+    }
 }
