line chart: introduce onContextLost renderer callback

tensorboard/webapp/widgets/line_chart_v2/lib/chart.ts

@@ -66,7 +66,8 @@ export class ChartImpl implements Chart {
         this.renderer = new ThreeRenderer(
           option.container,
           coordinator,
-          option.devicePixelRatio
+          option.devicePixelRatio,
+          option.callbacks.onContextLost
         );
         break;
       }

tensorboard/webapp/widgets/line_chart_v2/lib/chart_types.ts

@@ -42,6 +42,7 @@ export interface Chart {
 
 export interface ChartCallbacks {
   onDrawEnd(): void;
+  onContextLost(): void;
 }
 
 export interface BaseChartOptions {

tensorboard/webapp/widgets/line_chart_v2/lib/integration_test.ts

@@ -42,6 +42,7 @@ describe('line_chart_v2/lib/integration test', () => {
     dom = document.createElementNS('http://www.w3.org/2000/svg', 'svg');
     callbacks = {
       onDrawEnd: jasmine.createSpy(),
+      onContextLost: jasmine.createSpy(),
     };
     chart = new ChartImpl({
       type: RendererType.SVG,
@@ -460,4 +461,44 @@ describe('line_chart_v2/lib/integration test', () => {
       expect(circle.getAttribute('cy')).toBe('5');
     });
   });
+
+  describe('webgl', () => {
+    it('invokes onContextLost after losing webgl context', async () => {
+      const canvas = document.createElement('canvas');
+      chart = new ChartImpl({
+        type: RendererType.WEBGL,
+        container: canvas,
+        callbacks,
+        domDimension: {width: 200, height: 100},
+        useDarkMode: false,
+        devicePixelRatio: 1,
+      });
+      chart.setXScaleType(ScaleType.LINEAR);
+      chart.setYScaleType(ScaleType.LINEAR);
+
+      expect(callbacks.onContextLost).not.toHaveBeenCalled();
+
+      // For more info about forcing context loss, see
+      // https://developer.mozilla.org/en-US/docs/Web/API/WEBGL_lose_context/loseContext
+      const glExtension = canvas
+        .getContext('webgl2')
+        ?.getExtension('WEBGL_lose_context');
+      if (!glExtension) {
+        console.log(
+          'The browser used for testing does not ' +
+            'support WebGL or extensions needed for testing.'
+        );
+        return;
+      }
+
+      // The `loseContext` triggers the event asynchronously, which.
+      const contextLostPromise = new Promise((resolve) => {
+        canvas.addEventListener('webglcontextlost', resolve);
+      });
+      glExtension.loseContext();
+
+      await contextLostPromise;
+      expect(callbacks.onContextLost).toHaveBeenCalledTimes(1);
+    });
+  });
 });

tensorboard/webapp/widgets/line_chart_v2/lib/renderer/threejs_renderer.ts

@@ -268,12 +268,19 @@ export class ThreeRenderer implements ObjectRenderer<CacheValue> {
   constructor(
     canvas: HTMLCanvasElement | OffscreenCanvas,
     private readonly coordinator: ThreeCoordinator,
-    devicePixelRatio: number
+    devicePixelRatio: number,
+    onContextLost?: EventListener
   ) {
     if (isOffscreenCanvasSupported() && canvas instanceof OffscreenCanvas) {
       // THREE.js require the style object which Offscreen canvas lacks.
       (canvas as any).style = (canvas as any).style || {};
     }
+    // WebGL contexts may be abandoned by the browser if too many contexts are
+    // created on the same page.
+    if (onContextLost) {
+      canvas.addEventListener('webglcontextlost', onContextLost);
+    }
+
     this.renderer = new THREE.WebGLRenderer({
       canvas: canvas as HTMLCanvasElement,
       context: canvas.getContext('webgl2', {

tensorboard/webapp/widgets/line_chart_v2/lib/worker/message_types.ts

@@ -83,10 +83,15 @@ export type HostToGuestMessage =
 
 export enum GuestToMainType {
   ON_REDRAW_END,
+  ON_CONTEXT_LOST,
 }
 
 export interface RedrawEndMessage {
   type: GuestToMainType.ON_REDRAW_END;
 }
 
-export type GuestToMainMessage = RedrawEndMessage;
+export interface ContextLostMessage {
+  type: GuestToMainType.ON_CONTEXT_LOST;
+}
+
+export type GuestToMainMessage = RedrawEndMessage | ContextLostMessage;

tensorboard/webapp/widgets/line_chart_v2/lib/worker/worker_chart.ts

@@ -146,6 +146,10 @@ export class WorkerChart implements Chart {
         this.callbacks.onDrawEnd();
         break;
       }
+      case GuestToMainType.ON_CONTEXT_LOST: {
+        this.callbacks.onContextLost();
+        break;
+      }
     }
   }
 }

tensorboard/webapp/widgets/line_chart_v2/lib/worker/worker_chart_bridge.ts

@@ -43,6 +43,11 @@ function createPortHandler(port: MessagePort, initMessage: InitMessage) {
         type: GuestToMainType.ON_REDRAW_END,
       });
     },
+    onContextLost: () => {
+      port.postMessage({
+        type: GuestToMainType.ON_CONTEXT_LOST,
+      });
+    },
   };
 
   let chartOptions: ChartOptions;

tensorboard/webapp/widgets/line_chart_v2/lib/worker/worker_chart_test.ts

@@ -23,6 +23,7 @@ describe('line_chart_v2/lib/worker_chart test', () => {
   let workerPostMessageSpy: jasmine.Spy;
   let channelTxSpy: jasmine.Spy;
   let onDrawEndSpy: jasmine.Spy;
+  let onContextLostSpy: jasmine.Spy;
   let chart: WorkerChart;
 
   beforeEach(() => {
@@ -41,10 +42,11 @@ describe('line_chart_v2/lib/worker_chart test', () => {
     });
 
     onDrawEndSpy = jasmine.createSpy();
+    onContextLostSpy = jasmine.createSpy();
     chart = new WorkerChart({
       type: RendererType.WEBGL,
       devicePixelRatio: 1,
-      callbacks: {onDrawEnd: onDrawEndSpy},
+      callbacks: {onDrawEnd: onDrawEndSpy, onContextLost: onContextLostSpy},
       container: document.createElement('canvas'),
       domDimension: {width: 100, height: 200},
       useDarkMode: false,

tensorboard/webapp/widgets/line_chart_v2/line_chart_component.ts

@@ -296,7 +296,10 @@ export class LineChartComponent
 
     const rendererType = this.getRendererType();
     // Do not yet need to subscribe to the `onDrawEnd`.
-    const callbacks: ChartCallbacks = {onDrawEnd: () => {}};
+    const callbacks: ChartCallbacks = {
+      onDrawEnd: () => {},
+      onContextLost: () => {},
+    };
     let params: ChartOptions | null = null;
 
     this.readAndUpdateDomDimensions();