fix: mount component correctly (#671)

src/Globe/Globe.stories.tsx

@@ -1,8 +1,14 @@
 import { action } from "@storybook/addon-actions";
 import { Meta, StoryObj } from "@storybook/react";
-import { CesiumTerrainProvider, EllipsoidTerrainProvider, IonResource } from "cesium";
-import { StrictMode } from "react";
+import {
+  CesiumTerrainProvider,
+  EllipsoidTerrainProvider,
+  IonResource,
+  Viewer as CesiumViewer,
+} from "cesium";
+import { StrictMode, useEffect, useRef } from "react";
 
+import { CesiumComponentRef } from "../core";
 import Viewer from "../Viewer";
 
 import Globe from "./Globe";
@@ -49,3 +55,95 @@ export const Promise: Story = {
     </StrictMode>
   ),
 };
+
+export const PromiseStrict: Story = {
+  args: { enableLighting: true, isTerrainEnabled: false } as any,
+  render: ({ isTerrainEnabled, ...args }: any) => (
+    <StrictMode>
+      <Viewer full>
+        <Globe
+          {...args}
+          terrainProvider={
+            isTerrainEnabled
+              ? CesiumTerrainProvider.fromUrl(IonResource.fromAssetId(1), {
+                  requestVertexNormals: true,
+                  requestWaterMask: false,
+                })
+              : new EllipsoidTerrainProvider()
+          }
+          onImageryLayersUpdate={action("onImageryLayersUpdate")}
+          onTerrainProviderChange={action("onTerrainProviderChange")}
+        />
+      </Viewer>
+    </StrictMode>
+  ),
+};
+
+const DynamicComp = ({ isTerrainEnabled, ...args }: any) => {
+  const viewerRef = useRef<CesiumComponentRef<CesiumViewer> | null>(null);
+  useEffect(() => {
+    const run = async () => {
+      if (!viewerRef.current?.cesiumElement) {
+        return;
+      }
+
+      viewerRef.current.cesiumElement.terrainProvider = isTerrainEnabled
+        ? await CesiumTerrainProvider.fromUrl(IonResource.fromAssetId(1), {
+            requestVertexNormals: true,
+            requestWaterMask: false,
+          })
+        : new EllipsoidTerrainProvider();
+    };
+    run();
+  }, [viewerRef, isTerrainEnabled]);
+
+  return (
+    <Viewer full ref={viewerRef}>
+      <Globe
+        {...args}
+        onImageryLayersUpdate={action("onImageryLayersUpdate")}
+        onTerrainProviderChange={action("onTerrainProviderChange")}
+      />
+    </Viewer>
+  );
+};
+export const Dynamic: Story = {
+  args: { enableLighting: true, isTerrainEnabled: false } as any,
+  render: ({ ...args }: any) => <DynamicComp {...args} />,
+};
+
+const DynamicStrictComp = ({ isTerrainEnabled, ...args }: any) => {
+  const viewerRef = useRef<CesiumComponentRef<CesiumViewer> | null>(null);
+  useEffect(() => {
+    const run = () => {
+      // TODO: Remove this setTimeout in  strict mode.
+      setTimeout(async () => {
+        if (!viewerRef.current?.cesiumElement) {
+          return;
+        }
+
+        viewerRef.current.cesiumElement.terrainProvider = isTerrainEnabled
+          ? await CesiumTerrainProvider.fromUrl(IonResource.fromAssetId(1), {
+              requestVertexNormals: true,
+              requestWaterMask: false,
+            })
+          : new EllipsoidTerrainProvider();
+      }, 1);
+    };
+    run();
+  }, [viewerRef, isTerrainEnabled]);
+
+  return (
+    <Viewer full ref={viewerRef}>
+      <Globe
+        {...args}
+        onImageryLayersUpdate={action("onImageryLayersUpdate")}
+        onTerrainProviderChange={action("onTerrainProviderChange")}
+      />
+    </Viewer>
+  );
+};
+export const DynamicStrict: Story = {
+  args: { enableLighting: true, isTerrainEnabled: false } as any,
+  render: ({ ...args }: any) => <DynamicStrictComp {...args} />,
+};

src/core/hooks.ts

@@ -135,7 +135,7 @@ export const useCesiumComponent = <Element, Props extends RootComponentInternalP
         em.setEvents(useRootEvent ? null : element.current, props);
       }
 
-      if (update) {
+      if (update && mountedRef.current) {
         const maybePromise = update(element.current, props, prevProps.current, ctx);
         if (isPromise(maybePromise)) {
           await maybePromise;
@@ -155,6 +155,7 @@ export const useCesiumComponent = <Element, Props extends RootComponentInternalP
           );
         }
 
+        beforeUnmount();
         await unmount();
         mountReadyRef.current = mount();
       }
@@ -218,16 +219,24 @@ export const useCesiumComponent = <Element, Props extends RootComponentInternalP
       em.setEvents(useRootEvent ? null : element.current, initialProps.current);
     }
 
-    setMounted(true);
+    if (!unmountReadyRef.current) {
+      setMounted(true);
+    }
   }, []); // eslint-disable-line react-hooks/exhaustive-deps
 
+  const beforeUnmount = useCallback(() => {
+    setMounted(false);
+    mountedRef.current = false;
+  }, []);
+
   const unmount = useCallback(async () => {
     // Wait one tick to resolve Cesium's unmount process.
     await new Promise(r => queueMicrotask(() => r(undefined)));
 
     // Wait mount before unmount
     if (mountReadyRef.current) {
       await mountReadyRef.current;
+      mountReadyRef.current = undefined;
     }
 
     // Destroy cesium element
@@ -255,25 +264,25 @@ export const useCesiumComponent = <Element, Props extends RootComponentInternalP
     provided.current = undefined;
     stateRef.current = undefined;
     element.current = undefined;
-    mountedRef.current = false;
   }, []); // eslint-disable-line react-hooks/exhaustive-deps
 
   // To prevent re-execution by hot loader, execute only once
   useLayoutEffect(() => {
     const run = async () => {
       if (unmountReadyRef.current) {
         await unmountReadyRef.current;
+        unmountReadyRef.current = undefined;
       }
       mountReadyRef.current = mount();
     };
     run();
     return () => {
       // Need to update this state before Cesium is updated,
       // because Viewer's children must be hidden before it's unmounted.
-      setMounted(false);
+      beforeUnmount();
       unmountReadyRef.current = unmount();
     };
-  }, [mount, unmount]);
+  }, [mount, unmount, beforeUnmount]);
 
   // Update properties of cesium element
   useEffect(() => {
@@ -299,9 +308,13 @@ export const useCesiumComponent = <Element, Props extends RootComponentInternalP
   }, [ctx.__$internal, mounted, props, updateProperties]);
 
   // Expose cesium element
-  useImperativeHandle(ref, () => ({
-    cesiumElement: element.current,
-  }));
+  useImperativeHandle(
+    ref,
+    () => ({
+      cesiumElement: mounted ? element.current : null,
+    }),
+    [mounted],
+  );
 
   return [provided.current, mounted, wrapperRef];
 };