Commands to trash all connections to server (#32)

deephaven · Oct 23, 2024 · ad97e50 · ad97e50
1 parent 89f66c6
commit ad97e50
Show file tree

Hide file tree

Showing 6 changed files with 62 additions and 9 deletions.
diff --git a/package.json b/package.json
@@ -725,7 +725,7 @@
       "view/item/context": [
         {
           "command": "vscode-deephaven.connectToServer",
-          "when": "view == vscode-deephaven.serverTree && (viewItem == isManagedServerDisconnected || viewItem == isServerRunningDisconnected || viewItem == isDHEServerRunning)",
+          "when": "view == vscode-deephaven.serverTree && (viewItem == isManagedServerDisconnected || viewItem == isServerRunningDisconnected || viewItem == isDHEServerRunningConnected || viewItem == isDHEServerRunningDisconnected)",
           "group": "inline"
         },
         {
@@ -740,17 +740,17 @@
         },
         {
           "command": "vscode-deephaven.disconnectFromServer",
-          "when": "view == vscode-deephaven.serverConnectionTree && viewItem == isConnection",
+          "when": "(view == vscode-deephaven.serverTree && (viewItem == isServerRunningConnected || viewItem == isDHEServerRunningConnected)) || (view == vscode-deephaven.serverConnectionTree && viewItem == isConnection)",
           "group": "inline@2"
         },
         {
           "command": "vscode-deephaven.generateDHEKeyPair",
-          "when": "view == vscode-deephaven.serverTree && viewItem == isDHEServerRunning",
+          "when": "view == vscode-deephaven.serverTree && (viewItem == isDHEServerRunningConnected || viewItem == isDHEServerRunningDisconnected)",
           "group": "inline"
         },
         {
           "command": "vscode-deephaven.openInBrowser",
-          "when": "view == vscode-deephaven.serverTree && (viewItem == isManagedServerConnected || viewItem == isServerRunningConnected || viewItem == isServerRunningDisconnected || viewItem == isManagedServerDisconnected || viewItem == isDHEServerRunning)",
+          "when": "view == vscode-deephaven.serverTree && (viewItem == isManagedServerConnected || viewItem == isManagedServerDisconnected || viewItem == isServerRunningConnected || viewItem == isServerRunningDisconnected || viewItem == isDHEServerRunningConnected || viewItem == isDHEServerRunningDisconnected)",
           "group": "inline"
         },
         {

diff --git a/src/common/constants.ts b/src/common/constants.ts
@@ -97,7 +97,8 @@ export const PIP_SERVER_STATUS_DIRECTORY = 'pip-server-status';
 
 export const SERVER_TREE_ITEM_CONTEXT = {
   canStartServer: 'canStartServer',
-  isDHEServerRunning: 'isDHEServerRunning',
+  isDHEServerRunningConnected: 'isDHEServerRunningConnected',
+  isDHEServerRunningDisconnected: 'isDHEServerRunningDisconnected',
   isManagedServerConnected: 'isManagedServerConnected',
   isManagedServerConnecting: 'isManagedServerConnecting',
   isManagedServerDisconnected: 'isManagedServerDisconnected',

diff --git a/src/controllers/ExtensionController.ts b/src/controllers/ExtensionController.ts
@@ -579,9 +579,26 @@ export class ExtensionController implements Disposable {
    * Handle disconnecting from a server.
    */
   onDisconnectFromServer = async (
-    connectionState: ConnectionState
+    serverOrConnectionState: ServerState | ConnectionState
   ): Promise<void> => {
-    this._serverManager?.disconnectFromServer(connectionState.serverUrl);
+    // ConnectionState
+    if ('serverUrl' in serverOrConnectionState) {
+      this._serverManager?.disconnectFromServer(
+        serverOrConnectionState.serverUrl
+      );
+      return;
+    }
+
+    // DHC ServerState
+    if (serverOrConnectionState.type === 'DHC') {
+      this._coreCredentialsCache?.delete(serverOrConnectionState.url);
+      this._serverManager?.disconnectFromServer(serverOrConnectionState.url);
+      return;
+    }
+
+    // DHE ServerState
+    this._dheCredentialsCache?.delete(serverOrConnectionState.url);
+    this._serverManager?.disconnectFromDHEServer(serverOrConnectionState.url);
   };
 
   /**

diff --git a/src/services/ServerManager.ts b/src/services/ServerManager.ts
@@ -279,6 +279,37 @@ export class ServerManager implements IServerManager {
     this._onDidUpdate.fire();
   };
 
+  /**
+   * Completely disconnect from a DHE server. This including all workers plus
+   * the primary DHE client connection.
+   * @param dheServerUrl The URL of the DHE server to disconnect from.
+   * @returns Promise that resolves when all connections have been discarded.
+   */
+  disconnectFromDHEServer = async (dheServerUrl: URL): Promise<void> => {
+    const workerUrls = [...this._workerURLToServerURLMap.entries()].filter(
+      ([, url]) => url.toString() === dheServerUrl.toString()
+    );
+
+    for (const [workerUrl] of workerUrls) {
+      await this.disconnectFromServer(workerUrl);
+    }
+
+    this._dheServiceCache.invalidate(dheServerUrl);
+
+    const serverState = this._serverMap.get(dheServerUrl);
+    if (serverState == null) {
+      return;
+    }
+
+    this._serverMap.set(dheServerUrl, {
+      ...serverState,
+      isConnected: false,
+      connectionCount: 0,
+    });
+
+    this._onDidUpdate.fire();
+  };
+
   disconnectFromServer = async (
     serverOrWorkerUrl: URL | WorkerURL
   ): Promise<void> => {

diff --git a/src/types/serviceTypes.d.ts b/src/types/serviceTypes.d.ts
@@ -115,6 +115,7 @@ export interface IServerManager extends Disposable {
     workerConsoleType?: ConsoleType
   ) => Promise<ConnectionState | null>;
   disconnectEditor: (uri: vscode.Uri) => void;
+  disconnectFromDHEServer: (dheServerUrl: URL) => Promise<void>;
   disconnectFromServer: (serverUrl: URL) => Promise<void>;
   loadServerConfig: () => Promise<void>;
 

diff --git a/src/util/treeViewUtils.ts b/src/util/treeViewUtils.ts
@@ -118,7 +118,9 @@ export function getServerContextValue({
 
   if (isRunning) {
     if (isDHE) {
-      return SERVER_TREE_ITEM_CONTEXT.isDHEServerRunning;
+      return isConnected
+        ? SERVER_TREE_ITEM_CONTEXT.isDHEServerRunningConnected
+        : SERVER_TREE_ITEM_CONTEXT.isDHEServerRunningDisconnected;
     }
 
     return isConnected
@@ -252,7 +254,8 @@ export function getServerTreeItem(server: ServerState): vscode.TreeItem {
   const canConnect =
     contextValue === SERVER_TREE_ITEM_CONTEXT.isManagedServerDisconnected ||
     contextValue === SERVER_TREE_ITEM_CONTEXT.isServerRunningDisconnected ||
-    contextValue === SERVER_TREE_ITEM_CONTEXT.isDHEServerRunning;
+    contextValue === SERVER_TREE_ITEM_CONTEXT.isDHEServerRunningConnected ||
+    contextValue === SERVER_TREE_ITEM_CONTEXT.isDHEServerRunningDisconnected;
 
   return {
     label: new URL(urlStr).host,