fix: DH-20908: Handle table disconnects in Plotly Express

plugins/plotly-express/src/js/src/PlotlyExpressChartModel.test.ts

@@ -12,9 +12,13 @@ import {
 } from './PlotlyExpressChartUtils';
 
 const SMALL_TABLE = TestUtils.createMockProxy<DhType.Table>({
+  addEventListener: jest.fn(() => jest.fn()),
   columns: [{ name: 'x' }, { name: 'y' }] as DhType.Column[],
   size: 500,
-  subscribe: () => TestUtils.createMockProxy(),
+  subscribe: () =>
+    TestUtils.createMockProxy<DhType.TableSubscription>({
+      addEventListener: jest.fn(() => jest.fn()),
+    }),
 });
 
 const LARGE_TABLE = TestUtils.createMockProxy<DhType.Table>({
@@ -113,6 +117,7 @@ function createMockWidget(
       close: jest.fn(),
     })),
     addEventListener: jest.fn(),
+    close: jest.fn(),
     sendMessage: jest.fn(),
   } satisfies Partial<DhType.Widget> as unknown as DhType.Widget;
 }
@@ -187,7 +192,7 @@ const checkEventTypes = (mockSubscribe: jest.Mock, eventTypes: string[]) => {
 };
 
 beforeEach(() => {
-  jest.resetAllMocks();
+  jest.clearAllMocks();
 });
 
 describe('PlotlyExpressChartModel', () => {
@@ -614,4 +619,50 @@ describe('PlotlyExpressChartModel', () => {
       new CustomEvent(ChartModel.EVENT_LAYOUT_UPDATED)
     );
   });
+
+  it('should emit disconnect event when table disconnects', async () => {
+    const mockWidget = createMockWidget([SMALL_TABLE], 'scatter');
+    const chartModel = new PlotlyExpressChartModel(
+      mockDh,
+      mockWidget,
+      jest.fn()
+    );
+
+    const mockSubscribe = jest.fn();
+    await chartModel.subscribe(mockSubscribe);
+    await new Promise(process.nextTick);
+
+    const table0: DhType.Table = await mockWidget.exportedObjects[0].fetch();
+
+    const handler = TestUtils.findLastCall(
+      TestUtils.asMock(table0.addEventListener),
+      () => true
+    );
+
+    handler?.[1](new CustomEvent(mockDh.Table.EVENT_DISCONNECT));
+
+    expect(mockSubscribe).toHaveBeenCalledTimes(1);
+    expect(mockSubscribe).toHaveBeenLastCalledWith(
+      new CustomEvent(ChartModel.EVENT_DISCONNECT)
+    );
+  });
+
+  it('should cleanup subscriptions', async () => {
+    const mockWidget = createMockWidget([SMALL_TABLE], 'scatter');
+    const chartModel = new PlotlyExpressChartModel(
+      mockDh,
+      mockWidget,
+      jest.fn()
+    );
+
+    const mockSubscribe = jest.fn();
+    await chartModel.subscribe(mockSubscribe);
+    await new Promise(process.nextTick);
+
+    expect(chartModel.subscriptionCleanupMap.size).toBe(1);
+
+    chartModel.unsubscribe(jest.fn());
+
+    expect(chartModel.subscriptionCleanupMap.size).toBe(0);
+  });
 });

plugins/plotly-express/src/js/src/PlotlyExpressChartModel.ts

@@ -112,7 +112,7 @@ export class PlotlyExpressChartModel extends ChartModel {
   /**
    * Map of table index to cleanup function for the subscription.
    */
-  subscriptionCleanupMap: Map<number, () => void> = new Map();
+  subscriptionCleanupMap: Map<number, Set<() => void>> = new Map();
 
   /**
    * Map of table index to map of column names to array of paths where the data should be replaced.
@@ -189,6 +189,21 @@ export class PlotlyExpressChartModel extends ChartModel {
    */
   requiredColumns: Set<string> = new Set();
 
+  cleanupSubscriptions(id: number): void {
+    this.subscriptionCleanupMap.get(id)?.forEach(cleanup => {
+      cleanup();
+    });
+
+    try {
+      this.tableSubscriptionMap.get(id)?.close();
+    } catch {
+      // ignore
+    }
+
+    this.subscriptionCleanupMap.delete(id);
+    this.tableSubscriptionMap.delete(id);
+  }
+
   override getData(): Partial<Data>[] {
     const hydratedData = [...this.plotlyData];
 
@@ -688,11 +703,11 @@ export class PlotlyExpressChartModel extends ChartModel {
     log.debug('Updating downsampled table', downsampleInfo);
 
     this.fireDownsampleStart(null);
-    this.subscriptionCleanupMap.get(id)?.();
-    this.tableSubscriptionMap.get(id)?.close();
-    this.subscriptionCleanupMap.delete(id);
-    this.tableSubscriptionMap.delete(id);
+
+    this.cleanupSubscriptions(id);
+
     this.tableReferenceMap.delete(id);
+
     this.addTable(id, oldDownsampleInfo.originalTable);
   }
 
@@ -807,24 +822,35 @@ export class PlotlyExpressChartModel extends ChartModel {
       const columns = table.columns.filter(({ name }) => columnNames.has(name));
       const subscription = table.subscribe(columns);
       this.tableSubscriptionMap.set(id, subscription);
-      this.subscriptionCleanupMap.set(
-        id,
+
+      if (!this.subscriptionCleanupMap.has(id)) {
+        this.subscriptionCleanupMap.set(id, new Set());
+      }
+      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+      const cleanupSet = this.subscriptionCleanupMap.get(id)!;
+
+      cleanupSet.add(
         subscription.addEventListener<DhType.SubscriptionTableData>(
           this.dh.Table.EVENT_UPDATED,
           e => this.handleFigureUpdated(e, id)
         )
       );
+
+      cleanupSet.add(
+        table.addEventListener<DhType.Table>(
+          this.dh.Table.EVENT_DISCONNECT,
+          e => this.fireDisconnect()
+        )
+      );
     }
   }
 
   removeTable(id: number): void {
-    this.subscriptionCleanupMap.get(id)?.();
-    this.tableSubscriptionMap.get(id)?.close();
+    this.cleanupSubscriptions(id);
 
     this.tableReferenceMap.delete(id);
+
     this.downsampleMap.delete(id);
-    this.subscriptionCleanupMap.delete(id);
-    this.tableSubscriptionMap.delete(id);
     this.chartDataMap.delete(id);
     this.tableDataMap.delete(id);
     this.tableColumnReplacementMap.delete(id);