fix: Disposal Logic + Error toast

src/controllers/ConnectionController.ts

@@ -13,7 +13,7 @@
   createConnectionQuickPick,
   createConnectionQuickPickOptions,
   createConnectStatusBarItem,
   getConnectionsForConsoleType,
   getConsoleType,
   getEditorForUri,
   isSupportedLanguageId,
@@ -153,6 +153,7 @@
       // disconnect from it.
       if (err instanceof UnsupportedConsoleTypeError && newConnectionUrl) {
         this._serverManager.disconnectFromServer(newConnectionUrl);
+        this._toaster.error(err.message);
       }
 
       throw err;

src/controllers/ExtensionController.ts

@@ -117,11 +117,14 @@ export class ExtensionController implements Disposable {
   readonly _config: IConfigService;
 
   private _connectionController: ConnectionController | null = null;
-  private _coreClientCache: URLMap<CoreAuthenticatedClient> | null = null;
+  private _coreClientCache: URLMap<
+    CoreAuthenticatedClient & Disposable
+  > | null = null;
   private _coreClientFactory: ICoreClientFactory | null = null;
   private _coreJsApiCache: IAsyncCacheService<URL, typeof DhcType> | null =
     null;
-  private _dheClientCache: URLMap<DheAuthenticatedClient> | null = null;
+  private _dheClientCache: URLMap<DheAuthenticatedClient & Disposable> | null =
+    null;
   private _dheClientFactory: IDheClientFactory | null = null;
   private _dheServiceCache: IAsyncCacheService<URL, IDheService> | null = null;
   private _panelController: PanelController | null = null;
@@ -161,6 +164,7 @@ export class ExtensionController implements Disposable {
     const codelensProvider = new RunCommandCodeLensProvider();
 
     this._context.subscriptions.push(
+      codelensProvider,
       vscode.languages.registerCodeLensProvider('groovy', codelensProvider),
       vscode.languages.registerCodeLensProvider('python', codelensProvider)
     );
@@ -329,22 +333,43 @@ export class ExtensionController implements Disposable {
 
     this._coreClientFactory = async (
       url: URL
-    ): Promise<CoreUnauthenticatedClient> => {
+    ): Promise<CoreUnauthenticatedClient & Disposable> => {
       assertDefined(this._coreJsApiCache, 'coreJsApiCache');
       const dhc = await this._coreJsApiCache.get(url);
-      return new dhc.CoreClient(url.toString()) as CoreUnauthenticatedClient;
+
+      const client = new dhc.CoreClient(
+        url.toString()
+      ) as CoreUnauthenticatedClient;
+
+      // Attach a dispose method so that client caches can dispose of the client
+      return Object.assign(client, {
+        dispose: async () => {
+          client.disconnect();
+        },
+      });
     };
 
     this._dheClientFactory = async (
       url: URL
-    ): Promise<DheUnauthenticatedClient> => {
+    ): Promise<DheUnauthenticatedClient & Disposable> => {
       assertDefined(this._dheJsApiCache, 'dheJsApiCache');
       const dhe = await this._dheJsApiCache.get(url);
-      return createDheClient(dhe, getWsUrl(url));
+
+      const client = await createDheClient(dhe, getWsUrl(url));
+
+      // Attach a dispose method so that client caches can dispose of the client
+      return Object.assign(client, {
+        dispose: async () => {
+          client.disconnect();
+        },
+      });
     };
 
-    this._coreClientCache = new URLMap();
-    this._dheClientCache = new URLMap();
+    this._coreClientCache = new URLMap<CoreAuthenticatedClient & Disposable>();
+    this._context.subscriptions.push(this._coreClientCache);
+
+    this._dheClientCache = new URLMap<DheAuthenticatedClient & Disposable>();
+    this._context.subscriptions.push(this._dheClientCache);
 
     this._panelService = new PanelService();
     this._context.subscriptions.push(this._panelService);
@@ -382,6 +407,7 @@ export class ExtensionController implements Disposable {
 
     this._serverManager.onDidDisconnect(
       serverUrl => {
+        this._panelService?.clearServerData(serverUrl);
         this._outputChannel?.appendLine(
           `Disconnected from server: '${serverUrl}'.`
         );
@@ -529,7 +555,10 @@ export class ExtensionController implements Disposable {
     this._context.subscriptions.push(
       this._serverTreeView,
       this._serverConnectionTreeView,
-      this._serverConnectionPanelTreeView
+      this._serverConnectionPanelTreeView,
+      this._serverTreeProvider,
+      this._serverConnectionTreeProvider,
+      this._serverConnectionPanelTreeProvider
     );
   };
 

src/providers/RunCommandCodeLensProvider.ts

@@ -1,11 +1,12 @@
 import * as vscode from 'vscode';
 import { ICON_ID } from '../common';
+import type { Disposable } from '../types';
 
 /**
  * Provides inline editor code lenses for running Deephaven code.
  */
 export class RunCommandCodeLensProvider
-  implements vscode.CodeLensProvider<vscode.CodeLens>
+  implements vscode.CodeLensProvider<vscode.CodeLens>, Disposable
 {
   constructor() {
     vscode.workspace.onDidChangeConfiguration(() => {
@@ -22,6 +23,10 @@ export class RunCommandCodeLensProvider
   readonly onDidChangeCodeLenses: vscode.Event<void> =
     this._onDidChangeCodeLenses.event;
 
+  dispose = async (): Promise<void> => {
+    this._onDidChangeCodeLenses.dispose();
+  };
+
   provideCodeLenses(
     document: vscode.TextDocument,
     _token: vscode.CancellationToken

src/providers/ServerConnectionPanelTreeProvider.ts

@@ -11,6 +11,7 @@ import {
   getPanelVariableTreeItem,
   sortByStringProp,
 } from '../util';
+import { getFirstSupportedConsoleType } from '../services';
 
 export class ServerConnectionPanelTreeProvider extends TreeDataProviderBase<ServerConnectionPanelNode> {
   constructor(serverManager: IServerManager, panelService: IPanelService) {
@@ -31,7 +32,10 @@ export class ServerConnectionPanelTreeProvider extends TreeDataProviderBase<Serv
       return getPanelVariableTreeItem(connectionOrVariable);
     }
 
-    return getPanelConnectionTreeItem(connectionOrVariable);
+    return getPanelConnectionTreeItem(
+      connectionOrVariable,
+      getFirstSupportedConsoleType
+    );
   };
 
   getChildren = (

src/providers/TreeDataProviderBase.ts

@@ -1,11 +1,11 @@
 import * as vscode from 'vscode';
-import type { IServerManager } from '../types';
+import type { Disposable, IServerManager } from '../types';
 
 /**
  * Base class for tree view data providers.
  */
 export abstract class TreeDataProviderBase<T>
-  implements vscode.TreeDataProvider<T>
+  implements vscode.TreeDataProvider<T>, Disposable
 {
   constructor(serverManager: IServerManager) {
     this.serverManager = serverManager;
@@ -27,6 +27,10 @@ export abstract class TreeDataProviderBase<T>
 
   abstract getChildren(element?: T | undefined): vscode.ProviderResult<T[]>;
 
+  dispose = async (): Promise<void> => {
+    this._onDidChangeTreeData.dispose();
+  };
+
   refresh(): void {
     this._onDidChangeTreeData.fire();
   }

src/services/DhcService.ts

@@ -178,7 +178,7 @@ export class DhcService implements IDhcService {
       try {
         const { cn, session } = await this.initSessionPromise;
 
-        cn.subscribeToFieldUpdates(changes => {
+        const fieldUpdateSubscription = cn.subscribeToFieldUpdates(changes => {
           this.panelService.updateVariables(
             this.serverUrl,
             changes as VariableChanges
@@ -198,15 +198,15 @@ export class DhcService implements IDhcService {
             panelVariablesToUpdate
           );
         });
+        this.subscriptions.push(fieldUpdateSubscription);
 
         // TODO: Use constant 'disconnect' event name
-        this.subscriptions.push(
-          cn.addEventListener('disconnect', () => {
-            this.clearCaches();
-          })
-        );
+        const disconnectSubscription = cn.addEventListener('disconnect', () => {
+          this.clearCaches();
+        });
+        this.subscriptions.push(disconnectSubscription);
 
-        session.onLogMessage(logItem => {
+        const logMessageSubscription = session.onLogMessage(logItem => {
           // TODO: Should this pull log level from config somewhere?
           if (logItem.logLevel !== 'INFO') {
             const date = new Date(logItem.micros / 1000);
@@ -217,6 +217,7 @@ export class DhcService implements IDhcService {
             );
           }
         });
+        this.subscriptions.push(logMessageSubscription);
       } catch (err) {}
     }
 

src/services/DheService.ts

@@ -256,8 +256,10 @@ export class DheService implements IDheService {
     const querySerials = [...this._querySerialSet];
 
     this._querySerialSet.clear();
-    this._workerInfoMap.clear();
 
-    await this._disposeQueries(querySerials);
+    await Promise.all([
+      this._workerInfoMap.dispose(),
+      this._disposeQueries(querySerials),
+    ]);
   };
 }

src/services/PanelService.ts

@@ -22,8 +22,25 @@ export class PanelService implements IPanelService, Disposable {
   private readonly _cnPanelMap: URLMap<VariablePanelMap>;
   private readonly _cnVariableMap: URLMap<VariableMap>;
 
+  /**
+   * Clear panel data for the given connection url.
+   * @param url The connection url.
+   */
+  clearServerData = (url: URL): void => {
+    this._cnPanelMap.delete(url);
+    this._cnVariableMap.delete(url);
+  };
+
+  /**
+   * Cleanup resources.
+   */
   dispose = async (): Promise<void> => {
-    this._cnPanelMap.clear();
+    this._onDidUpdate.dispose();
+
+    await Promise.all([
+      this._cnPanelMap.dispose(),
+      this._cnVariableMap.dispose(),
+    ]);
   };
 
   /**

src/services/PollingService.spec.ts

@@ -0,0 +1,88 @@
+import { beforeEach, describe, it, expect, vi, afterAll } from 'vitest';
+import { pollUntilTrue } from './PollingService';
+
+// See __mocks__/vscode.ts for the mock implementation
+vi.mock('vscode');
+
+beforeEach(() => {
+  vi.clearAllMocks();
+  vi.useFakeTimers();
+});
+
+afterAll(() => {
+  vi.useRealTimers();
+});
+
+const resolved = vi.fn().mockName('resolved');
+const rejected = vi.fn().mockName('rejected');
+
+describe('pollUntilTrue', () => {
+  const poll = vi.fn<() => Promise<boolean>>().mockName('poll');
+
+  it('should resolve when poll function returns true', async () => {
+    const intervalMs = 100;
+    poll.mockResolvedValue(false);
+
+    const { promise } = pollUntilTrue(poll, intervalMs);
+    promise.then(resolved);
+
+    // Initial polling call that is scheduled via setTimeout(..., 0)
+    await vi.advanceTimersToNextTimerAsync();
+    expect(poll).toHaveBeenCalledTimes(1);
+    expect(resolved).not.toHaveBeenCalled();
+
+    // 2nd poll (after first intervalMs)
+    await vi.advanceTimersByTimeAsync(intervalMs);
+    expect(poll).toHaveBeenCalledTimes(2);
+    expect(resolved).not.toHaveBeenCalled();
+
+    poll.mockResolvedValue(true);
+
+    // 3rd poll
+    await vi.advanceTimersByTimeAsync(intervalMs);
+    expect(poll).toHaveBeenCalledTimes(3);
+    expect(resolved).toHaveBeenCalledWith(true);
+
+    // Advance intervalMs. No more polling expected since resolved
+    await vi.advanceTimersByTimeAsync(intervalMs);
+    expect(poll).toHaveBeenCalledTimes(3);
+    expect(resolved).toHaveBeenCalledOnce();
+  });
+
+  it('should cancel polling if timeout exceeded', async () => {
+    const intervalMs = 100;
+    const timeoutMs = 1000;
+
+    poll.mockResolvedValue(false);
+
+    const { promise } = pollUntilTrue(poll, intervalMs, timeoutMs);
+    promise.then(resolved).catch(rejected);
+
+    expect(resolved).not.toHaveBeenCalled();
+    expect(rejected).not.toHaveBeenCalled();
+
+    await vi.advanceTimersByTimeAsync(timeoutMs);
+
+    expect(resolved).not.toHaveBeenCalled();
+    expect(rejected).toHaveBeenCalledWith(new Error('Polling cancelled'));
+  });
+
+  it('should cancel polling if cancel explicitly called', async () => {
+    const intervalMs = 100;
+
+    poll.mockResolvedValue(false);
+
+    const { promise, cancel } = pollUntilTrue(poll, intervalMs);
+    promise.then(resolved).catch(rejected);
+
+    expect(resolved).not.toHaveBeenCalled();
+    expect(rejected).not.toHaveBeenCalled();
+
+    cancel();
+
+    await vi.advanceTimersToNextTimerAsync();
+
+    expect(resolved).not.toHaveBeenCalled();
+    expect(rejected).toHaveBeenCalledWith(new Error('Polling cancelled'));
+  });
+});