feat(react-components): add RevealKeepAlive component to cache the vi…

…ewer when unmounting the RevealContainer (#3552) * feat: add RevealKeepAlive component to cache the viewer when unmounting the RevealContainer * fix: restore models based on transform as well to distinguish copies of the same model * chore: remove commented code * fix: broken stories --------- Co-authored-by: cognite-bulldozer[bot] <51074376+cognite-bulldozer[bot]@users.noreply.github.com>
cognitedata · Aug 8, 2023 · b08cbaf · b08cbaf
1 parent 839ebde
commit b08cbaf
Show file tree

Hide file tree

Showing 11 changed files with 258 additions and 35 deletions.
diff --git a/react-components/package.json b/react-components/package.json
@@ -20,6 +20,7 @@
     "@cognite/cogs.js": ">=9",
     "@cognite/reveal": "4.4.0",
     "react": ">=18",
+    "react-dom": ">=18",
     "styled-components": ">=5"
   },
   "devDependencies": {
@@ -39,6 +40,7 @@
     "@tanstack/react-query-devtools": "^4.29.19",
     "@types/lodash": "^4.14.190",
     "@types/react": "18.2.7",
+    "@types/react-dom": "^18.2.7",
     "@types/styled-components": "5.1.26",
     "@types/three": "0.155.0",
     "@typescript-eslint/eslint-plugin": "^5.50.0",

diff --git a/react-components/src/components/CadModelContainer/CadModelContainer.tsx b/react-components/src/components/CadModelContainer/CadModelContainer.tsx
@@ -11,9 +11,10 @@ import {
   DefaultNodeAppearance
 } from '@cognite/reveal';
 import { useReveal } from '../RevealContainer/RevealContext';
-import { type Matrix4 } from 'three';
+import { Matrix4 } from 'three';
 import { useSDK } from '../RevealContainer/SDKProvider';
 import { type CogniteClient } from '@cognite/sdk';
+import { useRevealKeepAlive } from '../RevealKeepAlive/RevealKeepAliveContext';
 
 export type NodeStylingGroup = {
   nodeIds: number[];
@@ -43,6 +44,7 @@ export function CadModelContainer({
   styling,
   onLoad
 }: CogniteCadModelProps): ReactElement {
+  const cachedViewerRef = useRevealKeepAlive();
   const [model, setModel] = useState<CogniteCadModel>();
   const viewer = useReveal();
   const sdk = useSDK();
@@ -51,7 +53,6 @@ export function CadModelContainer({
 
   useEffect(() => {
     addModel(modelId, revisionId, transform, onLoad).catch(console.error);
-    return removeModel;
   }, [modelId, revisionId, geometryFilter]);
 
   useEffect(() => {
@@ -70,6 +71,8 @@ export function CadModelContainer({
     };
   }, [styling, model]);
 
+  useEffect(() => removeModel, [model]);
+
   return <></>;
 
   async function addModel(
@@ -78,18 +81,35 @@ export function CadModelContainer({
     transform?: Matrix4,
     onLoad?: (model: CogniteCadModel) => void
   ): Promise<CogniteCadModel> {
-    const cadModel = await viewer.addCadModel({ modelId, revisionId });
+    const cadModel = await getOrAddModel();
     if (transform !== undefined) {
       cadModel.setModelTransformation(transform);
     }
     setModel(cadModel);
     onLoad?.(cadModel);
 
     return cadModel;
+
+    async function getOrAddModel(): Promise<CogniteCadModel> {
+      const viewerModel = viewer.models.find(
+        (model) =>
+          model.modelId === modelId &&
+          model.revisionId === revisionId &&
+          model.getModelTransformation().equals(transform ?? new Matrix4())
+      );
+      if (viewerModel !== undefined) {
+        return await Promise.resolve(viewerModel as CogniteCadModel);
+      }
+      return await viewer.addCadModel({ modelId, revisionId });
+    }
   }
 
   function removeModel(): void {
     if (model === undefined || !viewer.models.includes(model)) return;
+
+    if (cachedViewerRef !== undefined && !cachedViewerRef.isRevealContainerMountedRef.current)
+      return;
+
     viewer.removeModel(model);
     setModel(undefined);
   }

diff --git a/react-components/src/components/Image360CollectionContainer/Image360CollectionContainer.tsx b/react-components/src/components/Image360CollectionContainer/Image360CollectionContainer.tsx
@@ -4,6 +4,7 @@
 import { type ReactElement, useEffect, useRef } from 'react';
 import { useReveal } from '../RevealContainer/RevealContext';
 import { type Image360Collection } from '@cognite/reveal';
+import { useRevealKeepAlive } from '../RevealKeepAlive/RevealKeepAliveContext';
 
 type Image360CollectionContainerProps = {
   siteId: string;
@@ -14,6 +15,7 @@ export function Image360CollectionContainer({
   siteId,
   onLoad
 }: Image360CollectionContainerProps): ReactElement {
+  const cachedViewerRef = useRevealKeepAlive();
   const modelRef = useRef<Image360Collection>();
   const viewer = useReveal();
 
@@ -25,13 +27,27 @@ export function Image360CollectionContainer({
   return <></>;
 
   async function add360Collection(): Promise<void> {
-    const image360Collection = await viewer.add360ImageSet('events', { site_id: siteId });
+    const image360Collection = await getOrAdd360Collection();
     modelRef.current = image360Collection;
     onLoad?.(image360Collection);
+
+    async function getOrAdd360Collection(): Promise<Image360Collection> {
+      const collections = viewer.get360ImageCollections();
+      const collection = collections.find((collection) => collection.id === siteId);
+      if (collection !== undefined) {
+        return collection;
+      }
+
+      return await viewer.add360ImageSet('events', { site_id: siteId });
+    }
   }
 
   function remove360Collection(): void {
     if (modelRef.current === undefined) return;
+
+    if (cachedViewerRef !== undefined && !cachedViewerRef.isRevealContainerMountedRef.current)
+      return;
+
     viewer.remove360ImageSet(modelRef.current);
     modelRef.current = undefined;
   }

diff --git a/react-components/src/components/PointCloudContainer/PointCloudContainer.tsx b/react-components/src/components/PointCloudContainer/PointCloudContainer.tsx
@@ -10,8 +10,9 @@ import {
 } from '@cognite/reveal';
 
 import { useEffect, type ReactElement, useState } from 'react';
-import { type Matrix4 } from 'three';
+import { Matrix4 } from 'three';
 import { useReveal } from '../RevealContainer/RevealContext';
+import { useRevealKeepAlive } from '../RevealKeepAlive/RevealKeepAliveContext';
 
 export type AnnotationIdStylingGroup = {
   annotationIds: number[];
@@ -36,13 +37,13 @@ export function PointCloudContainer({
   transform,
   onLoad
 }: CognitePointCloudModelProps): ReactElement {
+  const cachedViewerRef = useRevealKeepAlive();
   const [model, setModel] = useState<CognitePointCloudModel>();
   const viewer = useReveal();
   const { modelId, revisionId } = addModelOptions;
 
   useEffect(() => {
     addModel(modelId, revisionId, transform).catch(console.error);
-    return removeModel;
   }, [modelId, revisionId]);
 
   useEffect(() => {
@@ -58,21 +59,39 @@ export function PointCloudContainer({
     return cleanStyling;
   }, [styling, model]);
 
+  useEffect(() => removeModel, [model]);
+
   return <></>;
 
   async function addModel(modelId: number, revisionId: number, transform?: Matrix4): Promise<void> {
-    const pointCloudModel = await viewer.addPointCloudModel({ modelId, revisionId });
+    const pointCloudModel = await getOrAddModel();
 
     if (transform !== undefined) {
       pointCloudModel.setModelTransformation(transform);
     }
     setModel(pointCloudModel);
     onLoad?.(pointCloudModel);
+
+    async function getOrAddModel(): Promise<CognitePointCloudModel> {
+      const viewerModel = viewer.models.find(
+        (model) =>
+          model.modelId === modelId &&
+          model.revisionId === revisionId &&
+          model.getModelTransformation().equals(transform ?? new Matrix4())
+      );
+      if (viewerModel !== undefined) {
+        return await Promise.resolve(viewerModel as CognitePointCloudModel);
+      }
+      return await viewer.addPointCloudModel({ modelId, revisionId });
+    }
   }
 
   function removeModel(): void {
     if (model === undefined || !viewer.models.includes(model)) return;
 
+    if (cachedViewerRef !== undefined && !cachedViewerRef.isRevealContainerMountedRef.current)
+      return;
+
     viewer.removeModel(model);
     setModel(undefined);
   }

diff --git a/react-components/src/components/RevealContainer/RevealContainer.tsx b/react-components/src/components/RevealContainer/RevealContainer.tsx
@@ -3,12 +3,14 @@
  */
 import { type CogniteClient } from '@cognite/sdk';
 import { useEffect, useRef, type ReactNode, useState, type ReactElement } from 'react';
+import { createPortal } from 'react-dom';
 import { Cognite3DViewer, type Cognite3DViewerOptions } from '@cognite/reveal';
 import { RevealContext } from './RevealContext';
 import { type Color } from 'three';
 import { ModelsLoadingStateContext } from '../Reveal3DResources/ModelsLoadingContext';
 import { SDKProvider } from './SDKProvider';
 import { QueryClientProvider, QueryClient } from '@tanstack/react-query';
+import { useRevealKeepAlive } from '../RevealKeepAlive/RevealKeepAliveContext';
 
 type RevealContainerProps = {
   color?: Color;
@@ -34,51 +36,64 @@ export function RevealContainer({
   color,
   viewerOptions
 }: RevealContainerProps): ReactElement {
+  const revealKeepAliveData = useRevealKeepAlive();
   const [viewer, setViewer] = useState<Cognite3DViewer>();
-  const revealDomElementRef = useRef<HTMLDivElement>(null);
+  const wrapperDomElement = useRef<HTMLDivElement | null>(null);
+  const viewerDomElement = useRef<HTMLElement | null>(null);
 
   useEffect(() => {
-    initializeViewer();
-    return disposeViewer;
+    const initializedViewer = getOrInitializeViewer();
+    if (revealKeepAliveData === undefined) {
+      return;
+    }
+    revealKeepAliveData.isRevealContainerMountedRef.current = true;
+    return () => {
+      if (revealKeepAliveData === undefined) {
+        initializedViewer.dispose();
+        return;
+      }
+      revealKeepAliveData.isRevealContainerMountedRef.current = false;
+    };
   }, []);
 
   return (
     <SDKProvider sdk={sdk}>
       <QueryClientProvider client={queryClient}>
-        <div
-          style={{ width: '100%', height: '100%', overflow: 'hidden' }}
-          ref={revealDomElementRef}>
+        <div style={{ width: '100%', height: '100%', overflow: 'hidden' }} ref={wrapperDomElement}>
           {mountChildren()}
         </div>
       </QueryClientProvider>
     </SDKProvider>
   );
 
   function mountChildren(): ReactElement {
-    if (viewer === undefined) return <></>;
+    if (viewer === undefined || viewerDomElement.current === null) return <></>;
     return (
       <>
         <RevealContext.Provider value={viewer}>
-          <ModelsLoadingProvider>{children}</ModelsLoadingProvider>
+          <ModelsLoadingProvider>
+            {createPortal(children, viewerDomElement.current)}
+          </ModelsLoadingProvider>
         </RevealContext.Provider>
       </>
     );
   }
 
-  function initializeViewer(): void {
-    const domElement = revealDomElementRef.current;
+  function getOrInitializeViewer(): Cognite3DViewer {
+    const domElement = wrapperDomElement.current;
     if (domElement === null) {
       throw new Error('Failure in mounting RevealContainer to DOM.');
     }
-    const viewer = new Cognite3DViewer({ ...viewerOptions, sdk, domElement });
+    const viewer =
+      revealKeepAliveData?.viewerRef.current ?? new Cognite3DViewer({ ...viewerOptions, sdk });
+    if (revealKeepAliveData !== undefined) {
+      revealKeepAliveData.viewerRef.current = viewer;
+    }
+    domElement.appendChild(viewer.domElement);
+    viewerDomElement.current = viewer.domElement;
     viewer.setBackgroundColor({ color, alpha: 1 });
     setViewer(viewer);
-  }
-
-  function disposeViewer(): void {
-    if (viewer === undefined) return;
-    viewer.dispose();
-    setViewer(undefined);
+    return viewer;
   }
 }
 

diff --git a/react-components/src/components/RevealKeepAlive/RevealKeepAlive.tsx b/react-components/src/components/RevealKeepAlive/RevealKeepAlive.tsx
@@ -0,0 +1,23 @@
+/*!
+ * Copyright 2023 Cognite AS
+ */
+
+import { type Cognite3DViewer } from '@cognite/reveal';
+import { type ReactNode, type ReactElement, useRef, useEffect } from 'react';
+import { RevealKeepAliveContext } from './RevealKeepAliveContext';
+
+export function RevealKeepAlive({ children }: { children?: ReactNode }): ReactElement {
+  const viewerRef = useRef<Cognite3DViewer>();
+  const isRevealContainerMountedRef = useRef<boolean>(false);
+  useEffect(() => {
+    return () => {
+      viewerRef.current?.dispose();
+      viewerRef.current = undefined;
+    };
+  }, []);
+  return (
+    <RevealKeepAliveContext.Provider value={{ viewerRef, isRevealContainerMountedRef }}>
+      {children}
+    </RevealKeepAliveContext.Provider>
+  );
+}
diff --git a/react-components/src/components/RevealKeepAlive/RevealKeepAliveContext.ts b/react-components/src/components/RevealKeepAlive/RevealKeepAliveContext.ts
@@ -0,0 +1,16 @@
+/*!
+ * Copyright 2023 Cognite AS
+ */
+import { type Cognite3DViewer } from '@cognite/reveal';
+import { type MutableRefObject, createContext, useContext } from 'react';
+
+export type RevealKeepAliveData = {
+  viewerRef: MutableRefObject<Cognite3DViewer | undefined>;
+  isRevealContainerMountedRef: MutableRefObject<boolean>;
+};
+
+export const RevealKeepAliveContext = createContext<RevealKeepAliveData | undefined>(undefined);
+
+export const useRevealKeepAlive = (): RevealKeepAliveData | undefined => {
+  return useContext(RevealKeepAliveContext);
+};
diff --git a/react-components/src/index.ts b/react-components/src/index.ts
@@ -13,6 +13,7 @@ export { Image360HistoricalDetails } from './components/Image360HistoricalDetail
 export { ViewerAnchor } from './components/ViewerAnchor/ViewerAnchor';
 export { CameraController } from './components/CameraController/CameraController';
 export { RevealToolbar } from './components/RevealToolbar/RevealToolbar';
+export { RevealKeepAlive } from './components/RevealKeepAlive/RevealKeepAlive';
 
 // Hooks
 export { useReveal } from './components/RevealContainer/RevealContext';