fix: Pip servers handle config changes gracefully

src/common/constants.ts

@@ -1,5 +1,5 @@
 import * as path from 'node:path';
-import type { ConsoleType, Port, VariableType } from '../types';
+import type { ConsoleType, VariableType } from '../types';
 
 export const EXTENSION_ID = 'vscode-deephaven' as const;
 
@@ -12,10 +12,6 @@ export const CONFIG_KEY = {
 export const DEFAULT_CONSOLE_TYPE = 'python' as const;
 // export const DHFS_SCHEME = 'dhfs';
 
-export const DEFAULT_PIP_PORT_RANGE: ReadonlySet<Port> = new Set([
-  10001, 10002, 10003, 10004,
-] as Port[]);
-
 export const PYTHON_ENV_WAIT = 1500 as const;
 
 export const PIP_SERVER_STATUS_CHECK_INTERVAL = 3000;

src/controllers/ExtensionController.ts

@@ -2,7 +2,6 @@ import * as vscode from 'vscode';
 import {
   CONNECT_TO_SERVER_CMD,
   CREATE_NEW_TEXT_DOC_CMD,
-  DEFAULT_PIP_PORT_RANGE,
   DISCONNECT_EDITOR_CMD,
   DISCONNECT_FROM_SERVER_CMD,
   DOWNLOAD_LOGS_CMD,
@@ -179,7 +178,6 @@ export class ExtensionController implements Disposable {
     this._pipServerController = new PipServerController(
       this._context,
       this._serverManager,
-      DEFAULT_PIP_PORT_RANGE,
       this._outputChannel,
       this._toaster
     );

src/controllers/PipServerController.ts

@@ -22,13 +22,11 @@ export class PipServerController implements Disposable {
   constructor(
     context: vscode.ExtensionContext,
     serverManager: IServerManager,
-    portRange: Iterable<Port>,
     outputChannel: vscode.OutputChannel,
     toastService: IToastService
   ) {
     this._context = context;
     this._pollers = new Map();
-    this._portRange = new Set(portRange);
     this._serverUrlTerminalMap = new Map();
     this._serverManager = serverManager;
     this._outputChannel = outputChannel;
@@ -51,17 +49,19 @@ export class PipServerController implements Disposable {
       undefined,
       this._context.subscriptions
     );
+
+    this._serverManager.onDidLoadConfig(this.onDidLoadConfig);
   }
 
   private readonly _context: vscode.ExtensionContext;
   private readonly _outputChannel: vscode.OutputChannel;
   private readonly _pollers: Map<Port, { cancel: () => void }>;
-  private readonly _portRange: ReadonlySet<Port>;
   private readonly _serverUrlTerminalMap: Map<Port, vscode.Terminal>;
   private readonly _serverManager: IServerManager;
   private readonly _toaster: IToastService;
   private _checkPipInstallPromise: Promise<boolean> | null = null;
   private _isPipServerInstalled = false;
+  private _reservedPorts: ReadonlySet<Port> = new Set();
 
   /**
    * Log and show an error message to the user.
@@ -127,10 +127,50 @@ export class PipServerController implements Disposable {
     return this._checkPipInstallPromise;
   };
 
-  getAvailablePorts = (): Port[] => {
-    return [...this._portRange]
-      .filter(port => !this._serverUrlTerminalMap.has(port))
-      .sort();
+  /**
+   * Gets the next available port for starting a pip server.
+   * @returns A port number or `null` if no ports are available.
+   */
+  getNextAvailablePort = (): Port | null => {
+    for (let i = 10000; i < 10050; ++i) {
+      if (
+        !this._serverUrlTerminalMap.has(i as Port) &&
+        !this._reservedPorts.has(i as Port)
+      ) {
+        return i as Port;
+      }
+    }
+
+    return null;
+  };
+
+  /**
+   * Whenever server config loads, reserve any ports that are explicitly
+   * configured.
+   */
+  onDidLoadConfig = (): void => {
+    const servers = this._serverManager.getServers();
+
+    const reservedPorts = new Set<Port>();
+    const toDispose = new Set<Port>();
+
+    for (const server of servers) {
+      const port = parsePort(server.url.port);
+
+      reservedPorts.add(port);
+
+      // If an existing pip managed server port has become explicitly configured,
+      // mark it for disposal.
+      if (!server.isManaged && this._serverUrlTerminalMap.has(port)) {
+        toDispose.add(port);
+      }
+    }
+
+    this._reservedPorts = reservedPorts;
+
+    if (toDispose.size > 0) {
+      this.disposeServers(toDispose);
+    }
   };
 
   pollUntilServerStarts = async (port: Port): Promise<void> => {
@@ -163,7 +203,7 @@ export class PipServerController implements Disposable {
   };
 
   startServer = async (): Promise<void> => {
-    const [port] = this.getAvailablePorts();
+    const port = this.getNextAvailablePort();
 
     if (port == null) {
       this._logAndShowError('No available ports');
@@ -254,7 +294,7 @@ export class PipServerController implements Disposable {
     }
 
     this._serverManager.canStartServer =
-      this._isPipServerInstalled && this.getAvailablePorts().length > 0;
+      this._isPipServerInstalled && this.getNextAvailablePort() != null;
 
     if (!this._isPipServerInstalled) {
       this._serverManager.syncManagedServers([]);
@@ -269,10 +309,24 @@ export class PipServerController implements Disposable {
   disposeServers = async (ports: Iterable<Port>): Promise<void> => {
     for (const port of ports) {
       const terminal = this._serverUrlTerminalMap.get(port);
+      this._serverUrlTerminalMap.delete(port);
+
+      if (terminal != null && terminal.exitStatus == null) {
+        // One time subscription to update server status after terminal is closed
+        const oneTime = vscode.window.onDidCloseTerminal(t => {
+          if (t === terminal) {
+            oneTime.dispose();
+            this._serverManager.updateStatus();
+          }
+        });
 
-      if (terminal != null) {
-        terminal.dispose();
-        this._serverUrlTerminalMap.delete(port);
+        // Send ctrl+c to stop pip server, then exit the terminal. This allows
+        // `onDidCloseTerminal` to fire once the server is actually stopped vs
+        // `terminal.dispose()` which will fire `onDidCloseTerminal` immediately
+        // before the server process has actually finished exiting.
+        const ctrlC = String.fromCharCode(3);
+        terminal.sendText(ctrlC);
+        terminal.sendText('exit');
       }
 
       this._pollers.get(port)?.cancel();

src/services/SerializedKeyMap.ts

@@ -67,7 +67,7 @@ export abstract class SerializedKeyMap<TKey, TValue> {
     }
   }
 
-  *values(): IterableIterator<TValue> {
-    yield* this._map.values();
+  values(): IteratorObject<TValue> {
+    return this._map.values();
   }
 }

src/services/ServerManager.ts

@@ -48,6 +48,9 @@ export class ServerManager implements IServerManager {
   private readonly _onDidDisconnect = new vscode.EventEmitter<URL>();
   readonly onDidDisconnect = this._onDidDisconnect.event;
 
+  private readonly _onDidLoadConfig = new vscode.EventEmitter<void>();
+  readonly onDidLoadConfig = this._onDidLoadConfig.event;
+
   private readonly _onDidServerStatusChange =
     new vscode.EventEmitter<ServerState>();
   readonly onDidServerStatusChange = this._onDidServerStatusChange.event;
@@ -63,21 +66,31 @@ export class ServerManager implements IServerManager {
   canStartServer: boolean;
 
   loadServerConfig = async (): Promise<void> => {
-    const initialDhcServerState = getInitialServerStates(
+    // We want to keep any existing managed servers that aren't overridden by
+    // the latest config so we don't lose the PSKs that were generated when
+    // the servers were created.
+    const managedServersStates = this._serverMap
+      .values()
+      .filter(v => v.isManaged);
+
+    const configuredDhcServerState = getInitialServerStates(
       'DHC',
       this._configService.getCoreServers()
     );
 
-    const initialDheServerState = getInitialServerStates(
+    const configuredDheServerState = getInitialServerStates(
       'DHE',
       this._configService.getEnterpriseServers()
     );
 
     this._serverMap = new URLMap(
-      [...initialDhcServerState, ...initialDheServerState].map(state => [
-        state.url,
-        state,
-      ])
+      [
+        // Managed (pip) servers are first so that they can be overridden by the
+        // configured servers if necessary
+        ...managedServersStates,
+        ...configuredDhcServerState,
+        ...configuredDheServerState,
+      ].map(state => [state.url, state])
     );
 
     // If server config changes in a way that removes servers, disconnect any
@@ -89,6 +102,8 @@ export class ServerManager implements IServerManager {
     }
 
     await this.updateStatus();
+
+    this._onDidLoadConfig.fire();
   };
 
   connectToServer = async (serverUrl: URL): Promise<IDhService | null> => {

src/types/serviceTypes.d.ts

@@ -120,6 +120,7 @@ export interface IServerManager extends Disposable {
 
   onDidConnect: vscode.Event<URL>;
   onDidDisconnect: vscode.Event<URL>;
+  onDidLoadConfig: vscode.Event<void>;
   onDidRegisterEditor: vscode.Event<vscode.Uri>;
   onDidServerStatusChange: vscode.Event<ServerState>;
   onDidUpdate: vscode.Event<void>;