Dispose services / connections when caches are cleared (#22)

src/common/types.d.ts

@@ -4,3 +4,7 @@ export type ConnectionAndSession<TConnection, TSession> = {
   cn: TConnection;
   session: TSession;
 };
+
+export interface Disposable {
+  dispose(): Promise<void>;
+}

src/extension.ts

@@ -44,20 +44,23 @@ export function activate(context: vscode.ExtensionContext) {
     outputChannel
   );
 
-  dhcServiceRegistry.addEventListener('disconnect', () => {
+  dhcServiceRegistry.addEventListener('disconnect', serverUrl => {
+    vscode.window.showInformationMessage(
+      `Disconnected from Deephaven server: ${serverUrl}`
+    );
     clearConnection();
   });
 
   /*
    * Clear connection data
    */
-  function clearConnection() {
+  async function clearConnection() {
     selectedConnectionUrl = null;
     selectedDhService = null;
     const { text, tooltip } = createConnectTextAndTooltip('disconnected');
     connectStatusBarItem.text = text;
     connectStatusBarItem.tooltip = tooltip;
-    dhcServiceRegistry.clearCache();
+    await dhcServiceRegistry.clearCache();
   }
 
   /**
@@ -85,6 +88,7 @@ export function activate(context: vscode.ExtensionContext) {
   const connectStatusBarItem = createConnectStatusBarItem();
 
   context.subscriptions.push(
+    dhcServiceRegistry,
     outputChannel,
     runCodeCmd,
     runSelectionCmd,

src/services/CacheService.ts

@@ -1,9 +1,11 @@
+import { Disposable } from '../common';
+import { isDisposable } from '../util';
 import { EventDispatcher } from './EventDispatcher';
 
-export class CacheService<
-  T,
-  TEventName extends string
-> extends EventDispatcher<TEventName> {
+export class CacheService<T, TEventName extends string>
+  extends EventDispatcher<TEventName>
+  implements Disposable
+{
   constructor(
     label: string,
     loader: (key: string | null) => Promise<T>,
@@ -35,7 +37,25 @@ export class CacheService<
     return this.cachedPromises.get(normalizeKey)!;
   }
 
-  public clearCache(): void {
+  public async clearCache(): Promise<void> {
+    try {
+      const allValues = await Promise.all([...this.cachedPromises.values()]);
+
+      // Since the cache is responsible for creating the promises, it is also
+      // responsible for disposing of them.
+      allValues.forEach(value => {
+        if (isDisposable(value)) {
+          value.dispose();
+        }
+      });
+    } catch (err) {
+      console.error('An error occurred while disposing cached values:', err);
+    }
+
     this.cachedPromises.clear();
   }
+
+  public async dispose(): Promise<void> {
+    await this.clearCache();
+  }
 }

src/services/DhService.ts

@@ -1,7 +1,7 @@
 import * as vscode from 'vscode';
 import type { dh as DhcType } from '../dh/dhc-types';
 import { hasErrorCode } from '../util/typeUtils';
-import { ConnectionAndSession } from '../common';
+import { ConnectionAndSession, Disposable } from '../common';
 import { ExtendedMap, formatTimestamp } from '../util';
 import { EventDispatcher } from './EventDispatcher';
 
@@ -26,10 +26,10 @@ type CommandResultBase = {
   error: string;
 };
 
-export abstract class DhService<
-  TDH,
-  TClient
-> extends EventDispatcher<'disconnect'> {
+export abstract class DhService<TDH, TClient>
+  extends EventDispatcher<'disconnect'>
+  implements Disposable
+{
   constructor(
     serverUrl: string,
     panelRegistry: ExtendedMap<string, vscode.WebviewPanel>,
@@ -82,13 +82,23 @@ export abstract class DhService<
     this.cachedCreateSession = null;
     this.cachedInitApi = null;
     this.client = null;
+
+    if (this.cn != null) {
+      this.cn.close();
+      this.dispatchEvent('disconnect', this.serverUrl);
+    }
     this.cn = null;
+
     this.dh = null;
     this.session = null;
 
     this.subscriptions.forEach(dispose => dispose());
   }
 
+  public async dispose(): Promise<void> {
+    this.clearCaches();
+  }
+
   public get isInitialized(): boolean {
     return this.cachedInitApi != null;
   }
@@ -134,12 +144,6 @@ export abstract class DhService<
         this.subscriptions.push(
           cn.addEventListener('disconnect', () => {
             this.clearCaches();
-
-            vscode.window.showInformationMessage(
-              `Disconnected from Deephaven server: ${this.serverUrl}`
-            );
-
-            this.dispatchEvent('disconnect');
           })
         );
       }

src/services/DhServiceRegistry.ts

@@ -25,8 +25,9 @@ export class DhServiceRegistry<T extends DhcService> extends CacheService<
           outputChannel
         );
 
-        dhService.addEventListener('disconnect', () => {
-          this.dispatchEvent('disconnect');
+        // Propagate service events as registry events.
+        dhService.addEventListener('disconnect', event => {
+          this.dispatchEvent('disconnect', event);
         });
 
         return dhService;

src/services/EventDispatcher.ts

@@ -34,6 +34,8 @@ export class EventDispatcher<TEventName extends string> {
    * @param event The event to dispatch to all listeners
    */
   dispatchEvent = <TEvent>(eventName: TEventName, event?: TEvent): void => {
-    this.listeners.get(eventName)?.forEach(listener => listener(event));
+    this.listeners.get(eventName)?.forEach(listener => {
+      listener(event);
+    });
   };
 }

src/util/index.ts

@@ -1,5 +1,6 @@
 export * from './downloadUtils';
 export * from './ExtendedMap';
+export * from './isDisposable';
 export * from './panelUtils';
 export * from './polyfillUtils';
 export * from './uiUtils';

src/util/isDisposable.ts

@@ -0,0 +1,17 @@
+import { Disposable } from '../common';
+
+/**
+ * Typeguard to determine if given object has a `disposable` method.
+ * @param maybeDisposable The object to check.
+ * @returns `true` if the object has a `dispose` method, `false` otherwise.
+ */
+export function isDisposable(
+  maybeDisposable: unknown
+): maybeDisposable is Disposable {
+  return (
+    maybeDisposable != null &&
+    typeof maybeDisposable === 'object' &&
+    'dispose' in maybeDisposable &&
+    typeof maybeDisposable.dispose === 'function'
+  );
+}