fix: handle context loss and restore, update opacity, reset shared resources, reuse temp positions map

Author: asuzuki-jumptrading

src/features/Overview/ShredsProgression/WebGl/Chart.tsx

@@ -1,29 +1,51 @@
-import { useLayoutEffect, useRef } from "react";
+import { useCallback, useLayoutEffect, useRef, useState } from "react";
 
 import { useMeasure, useRafLoop } from "react-use";
 import { Box, Flex } from "@radix-ui/themes";
 import type { FlexProps } from "@radix-ui/themes";
 import { useShredsChartScale } from "../useShredsChartScale";
 import { useSetAtom } from "jotai";
-import { minDirtySlotByChartAtom } from "../atoms";
+import { isWebgl2SupportedAtom, minDirtySlotByChartAtom } from "../atoms";
 import type { RendererObj, TsRange } from "./chartUtils";
 import { setUpRenderer, draw } from "./chartUtils";
 import ShredsSlotLabels from "../ShredsSlotLabels";
 import { MChartAxes, xAxisHeight } from "./ChartAxes";
 import { createLabelsState, type LabelsState } from "../utils";
+import { disposeWebglResources } from "../webglUtils";
 
 const REDRAW_INTERVAL_MS = 15;
+/**
+ * How long to wait for the GPU to restore a lost WebGL context before falling back to the canvas chart.
+ * Context loss is usually transient (tab backgrounded, GPU reset, driver hiccup) and the browser restores it within
+ * 1 ~ 2 frames.
+ */
+const CONTEXT_RESTORE_TIMEOUT_MS = 10_000;
 
-type FlexPropsSubset = Pick<FlexProps, "height" | "minHeight" | "flexGrow">;
-
-interface ShredsChartProps {
+interface ShredsChartProps
+  extends Pick<FlexProps, "height" | "minHeight" | "flexGrow"> {
   chartId: string;
 }
-export default function ShredsChart({
+export default function ShredsChart(props: ShredsChartProps) {
+  const [key, setKey] = useState(0);
+  const triggerRemount = useCallback(() => {
+    setKey((prev) => prev + 1);
+  }, [setKey]);
+  return (
+    <ShredsChartInner key={key} {...props} triggerRemount={triggerRemount} />
+  );
+}
+
+interface ShredsChartInnerProps extends ShredsChartProps {
+  triggerRemount: () => void;
+}
+
+export function ShredsChartInner({
   chartId,
+  triggerRemount,
   ...flexProps
-}: ShredsChartProps & FlexPropsSubset) {
+}: ShredsChartInnerProps) {
   const setMinDirtySlotByChart = useSetAtom(minDirtySlotByChartAtom);
+  const setWebgl2Supported = useSetAtom(isWebgl2SupportedAtom);
 
   const prevTimeDiffsRef = useRef<number[]>([]);
   const lastRedrawRef = useRef(0);
@@ -41,9 +63,31 @@ export default function ShredsChart({
     prevLabels: createLabelsState(),
     tempNewLabels: createLabelsState(),
   });
+  /**
+   * timeout to fall back to canvas chart if context is lost and not restored
+   */
+  const contextRestoreTimeoutRef = useRef<ReturnType<typeof setTimeout>>();
 
   const scale = useShredsChartScale();
 
+  const handleContextLost = useCallback(
+    (event: Event) => {
+      // preventDefault so that the browser fires webglcontextrestored
+      event.preventDefault();
+
+      clearTimeout(contextRestoreTimeoutRef.current);
+      contextRestoreTimeoutRef.current = setTimeout(() => {
+        setWebgl2Supported(false);
+      }, CONTEXT_RESTORE_TIMEOUT_MS);
+    },
+    [setWebgl2Supported],
+  );
+
+  const handleContextRestored = useCallback(() => {
+    clearTimeout(contextRestoreTimeoutRef.current);
+    triggerRemount();
+  }, [triggerRemount]);
+
   useLayoutEffect(() => {
     // setup dirty slot tracking
     setMinDirtySlotByChart((prev) => {
@@ -58,8 +102,9 @@ export default function ShredsChart({
         return prev;
       });
 
+      clearTimeout(contextRestoreTimeoutRef.current);
+
       // dispose of WebGL resources
-      // NOTE: keep sharedMaterial and unitQuad to be reused across mounts
       const obj = rendererRef.current;
       if (!obj) return;
 
@@ -69,13 +114,30 @@ export default function ShredsChart({
       for (const slotMesh of obj.availableMeshes) {
         slotMesh.mesh.geometry.dispose();
       }
+      // dispose this chart's own unitQuad / sharedMaterial
+      disposeWebglResources(obj.resources);
+
+      obj.renderer.domElement.removeEventListener(
+        "webglcontextlost",
+        handleContextLost,
+      );
+      obj.renderer.domElement.removeEventListener(
+        "webglcontextrestored",
+        handleContextRestored,
+      );
+
       obj.renderer.dispose();
       // force release of GL context on repeated mount / unmounts
       obj.renderer.forceContextLoss();
       obj.renderer.domElement.remove();
       rendererRef.current = undefined;
     };
-  }, [chartId, setMinDirtySlotByChart]);
+  }, [
+    chartId,
+    handleContextLost,
+    handleContextRestored,
+    setMinDirtySlotByChart,
+  ]);
 
   // handle chart resize
   useLayoutEffect(() => {
@@ -95,6 +157,10 @@ export default function ShredsChart({
   }, [scale, width, height, chartId]);
 
   useRafLoop(function drawShredsLoop(time: number) {
+    // Don't draw while waiting for context restore.
+    // but keep canvas mounted to listen for restore event.
+    if (contextRestoreTimeoutRef.current != null) return;
+
     if (
       lastRedrawRef.current == null ||
       time - lastRedrawRef.current >= REDRAW_INTERVAL_MS
@@ -105,7 +171,10 @@ export default function ShredsChart({
         if (!rendererObj) return;
 
         rendererRef.current = rendererObj;
-        containerRef.current?.replaceChildren(rendererObj.renderer.domElement);
+        const canvas = rendererObj.renderer.domElement;
+        canvas.addEventListener("webglcontextlost", handleContextLost);
+        canvas.addEventListener("webglcontextrestored", handleContextRestored);
+        containerRef.current?.replaceChildren(canvas);
       } else {
         draw(
           chartId,

src/features/Overview/ShredsProgression/WebGl/chartUtils.ts

@@ -16,13 +16,14 @@ import {
 } from "../atoms";
 import { shredEventDescPriorities } from "../const";
 import { updateLabels } from "../shredsProgressionPlugin";
-import type { SlotMesh } from "../webglUtils";
+import type { SlotMesh, WebglResources } from "../webglUtils";
 import {
   createSlotMesh,
   updateSlotMeshCounts,
   ensureCapacity,
   addRectangleToMesh,
   convertToWebGlColor,
+  createWebglResources,
 } from "../webglUtils";
 import {
   shredPublishedColor,
@@ -47,13 +48,20 @@ const store = getDefaultStore();
 const SKIPPED_SLOT_DOT_DURATION_MS = 10;
 const msPerDay = 24 * 60 * 60 * 1000;
 
+const tempEventPositions = new Map<
+  Exclude<ShredEvent, ShredEvent.slot_complete>,
+  { x: number; w: number }
+>();
+
 export type RendererObj = {
   renderer: THREE.WebGLRenderer;
   camera: THREE.OrthographicCamera;
   scene: THREE.Scene;
   meshes: Map<number, SlotMesh>;
   availableMeshes: SlotMesh[];
   worldTsRange: TsRange;
+  // resources shared by this renderer's slot meshes
+  resources: WebglResources;
 };
 
 const MAX_PIXEL_RATIO = 2;
@@ -96,6 +104,7 @@ export function setUpRenderer(canvasWidth: number, canvasHeight: number) {
 
     const meshes = new Map<number, SlotMesh>();
     const availableMeshes: SlotMesh[] = [];
+    const resources = createWebglResources();
     renderer.render(scene, camera);
 
     return {
@@ -105,6 +114,7 @@ export function setUpRenderer(canvasWidth: number, canvasHeight: number) {
       meshes,
       availableMeshes,
       worldTsRange,
+      resources,
     };
   } catch {
     // context creation can still fail despite the probe (e.g. too many live
@@ -186,10 +196,6 @@ export function draw(
 
   let anythingDrawn = false;
   const minDirtySlot = store.get(minDirtySlotByChartAtom).get(chartId);
-  const tempEventPositions = new Map<
-    Exclude<ShredEvent, ShredEvent.slot_complete>,
-    { x: number; w: number }
-  >();
 
   for (const slotNumber of orderedSlotNumbers) {
     const slot = liveShreds.slots.get(slotNumber);
@@ -199,7 +205,7 @@ export function draw(
     const isNewMesh = !slotMesh;
     if (!slotMesh) {
       const lastMesh = rendererObj.availableMeshes.pop();
-      slotMesh = lastMesh ?? createSlotMesh();
+      slotMesh = lastMesh ?? createSlotMesh(rendererObj.resources);
       rendererObj.meshes.set(slotNumber, slotMesh);
       rendererObj.scene.add(slotMesh.mesh);
     }
@@ -216,6 +222,7 @@ export function draw(
       const shred = slot.shreds[shredIdx];
       if (!shred) continue;
 
+      tempEventPositions.clear();
       const rectanglesAdded = addEventsForRow({
         tempEventPositions,
         slotMesh,
@@ -325,8 +332,6 @@ function addEventsForRow({
   y,
   visibleTsRange,
 }: AddEventsForRowArgs) {
-  tempEventPositions.clear();
-
   let endTs: number =
     slotCompletionTsDelta == null
       ? // event goes to max x

src/features/Overview/ShredsProgression/webglUtils.ts

@@ -1,6 +1,8 @@
 import * as THREE from "three";
 import { SHRED_EVENT_TYPES_COUNT } from "../../../api/entities";
 
+const SHREDS_OPACITY = 0.8;
+
 export type SlotMesh = {
   mesh: THREE.Mesh;
   rectArray: Float32Array;
@@ -39,26 +41,53 @@ void main() {
 }
 `;
 
-const unitQuad = new THREE.BufferGeometry();
-unitQuad.setAttribute(
-  "position",
-  new THREE.BufferAttribute(
-    new Float32Array([-0.5, -0.5, 0.5, -0.5, -0.5, 0.5, 0.5, 0.5]),
-    2,
-  ),
-);
-unitQuad.setIndex([0, 1, 2, 1, 3, 2]);
-unitQuad.boundingSphere = new THREE.Sphere();
-
-const sharedMaterial = new THREE.RawShaderMaterial({
-  vertexShader,
-  fragmentShader,
-  side: THREE.FrontSide,
-  transparent: true,
-  uniforms: {
-    uOpacity: { value: 1 },
-  },
-});
+/**
+ * Resources shared by all slot meshes of a single chart / renderer.
+ * Compiled shaders / uploaded buffers are bound to a specific GL
+ * context, so a fresh renderer (e.g. after a context loss) needs its own copy.
+ */
+export type WebglResources = {
+  unitQuad: THREE.BufferGeometry;
+  sharedMaterial: THREE.RawShaderMaterial;
+};
+
+function createUnitQuad() {
+  const geometry = new THREE.BufferGeometry();
+  geometry.setAttribute(
+    "position",
+    new THREE.BufferAttribute(
+      new Float32Array([-0.5, -0.5, 0.5, -0.5, -0.5, 0.5, 0.5, 0.5]),
+      2,
+    ),
+  );
+  geometry.setIndex([0, 1, 2, 1, 3, 2]);
+  geometry.boundingSphere = new THREE.Sphere();
+  return geometry;
+}
+
+function createSharedMaterial() {
+  return new THREE.RawShaderMaterial({
+    vertexShader,
+    fragmentShader,
+    side: THREE.FrontSide,
+    transparent: true,
+    uniforms: {
+      uOpacity: { value: SHREDS_OPACITY },
+    },
+  });
+}
+
+export function createWebglResources(): WebglResources {
+  return {
+    unitQuad: createUnitQuad(),
+    sharedMaterial: createSharedMaterial(),
+  };
+}
+
+export function disposeWebglResources(resources: WebglResources) {
+  resources.unitQuad.dispose();
+  resources.sharedMaterial.dispose();
+}
 
 // 700 shreds, all events except completion could have a rectangle
 const INITIAL_CAPACITY = 700 * (SHRED_EVENT_TYPES_COUNT - 1);
@@ -67,7 +96,7 @@ const INITIAL_CAPACITY = 700 * (SHRED_EVENT_TYPES_COUNT - 1);
  * Create one mesh per slot. That way, we can easily add / delete slots without updating
  * unchanged slots.
  */
-export function createSlotMesh(): SlotMesh {
+export function createSlotMesh(resources: WebglResources): SlotMesh {
   const rectArray = new Float32Array(INITIAL_CAPACITY * 4);
   const colorArray = new Float32Array(INITIAL_CAPACITY * 3);
 
@@ -78,14 +107,17 @@ export function createSlotMesh(): SlotMesh {
   colorAttr.setUsage(THREE.DynamicDrawUsage);
 
   const geometry = new THREE.InstancedBufferGeometry();
-  geometry.index = unitQuad.index;
-  geometry.setAttribute("position", unitQuad.getAttribute("position"));
+  geometry.index = resources.unitQuad.index;
+  geometry.setAttribute(
+    "position",
+    resources.unitQuad.getAttribute("position"),
+  );
   geometry.setAttribute("instanceRect", rectAttr);
   geometry.setAttribute("instanceColor", colorAttr);
   geometry.instanceCount = 0;
   geometry.boundingSphere = new THREE.Sphere();
 
-  const mesh = new THREE.Mesh(geometry, sharedMaterial);
+  const mesh = new THREE.Mesh(geometry, resources.sharedMaterial);
   mesh.frustumCulled = false;
 
   return {