fix: should not race with useLayoutEffect (#123)

MadCcc · web-flow · commit f4b8d014c4b5 · 2023-07-05T14:59:10.000+08:00
* fix: should not race with useLayoutEffect

* chore: optimize

* chore: add test case

* chore: rename
diff --git a/src/hooks/useCompatibleInsertionEffect.tsx b/src/hooks/useCompatibleInsertionEffect.tsx
@@ -0,0 +1,37 @@
+// import canUseDom from 'rc-util/lib/Dom/canUseDom';
+import useLayoutEffect from 'rc-util/lib/hooks/useLayoutEffect';
+import type { EffectCallback } from 'react';
+import * as React from 'react';
+
+// We need fully clone React function here
+// to avoid webpack warning React 17 do not export `useId`
+const fullClone = {
+  ...React,
+};
+const { useInsertionEffect } = fullClone;
+
+type UseCompatibleInsertionEffect = (
+  renderEffect: EffectCallback,
+  effect: EffectCallback,
+  deps?: React.DependencyList,
+) => void;
+
+const useInsertionEffectPolyfill: UseCompatibleInsertionEffect = (
+  renderEffect,
+  effect,
+  deps,
+) => {
+  React.useMemo(renderEffect, deps);
+  useLayoutEffect(effect, deps);
+};
+
+const useCompatibleInsertionEffect: UseCompatibleInsertionEffect =
+  useInsertionEffect
+    ? (renderEffect, effect, deps) =>
+        useInsertionEffect(() => {
+          renderEffect();
+          return effect();
+        }, deps)
+    : useInsertionEffectPolyfill;
+
+export default useCompatibleInsertionEffect;
diff --git a/src/hooks/useGlobalCache.tsx b/src/hooks/useGlobalCache.tsx
@@ -1,8 +1,8 @@
 import * as React from 'react';
 import type { KeyType } from '../Cache';
 import StyleContext from '../StyleContext';
+import useCompatibleInsertionEffect from './useCompatibleInsertionEffect';
 import useHMR from './useHMR';
-import useInsertionEffect from './useInsertionEffect';
 
 export default function useGlobalCache<CacheType>(
   prefix: string,
@@ -51,28 +51,32 @@ export default function useGlobalCache<CacheType>(
   const cacheContent = globalCache.get(fullPath)![1];
 
   // Remove if no need anymore
-  useInsertionEffect(() => {
-    onCacheEffect?.(cacheContent);
-
-    // It's bad to call build again in effect.
-    // But we have to do this since StrictMode will call effect twice
-    // which will clear cache on the first time.
-    buildCache(([times, cache]) => [times + 1, cache]);
-
-    return () => {
-      globalCache.update(fullPath, (prevCache) => {
-        const [times = 0, cache] = prevCache || [];
-        const nextCount = times - 1;
-
-        if (nextCount === 0) {
-          onCacheRemove?.(cache, false);
-          return null;
-        }
-
-        return [times - 1, cache];
-      });
-    };
-  }, [deps]);
+  useCompatibleInsertionEffect(
+    () => {
+      onCacheEffect?.(cacheContent);
+    },
+    () => {
+      // It's bad to call build again in effect.
+      // But we have to do this since StrictMode will call effect twice
+      // which will clear cache on the first time.
+      buildCache(([times, cache]) => [times + 1, cache]);
+
+      return () => {
+        globalCache.update(fullPath, (prevCache) => {
+          const [times = 0, cache] = prevCache || [];
+          const nextCount = times - 1;
+
+          if (nextCount === 0) {
+            onCacheRemove?.(cache, false);
+            return null;
+          }
+
+          return [times - 1, cache];
+        });
+      };
+    },
+    [deps],
+  );
 
   return cacheContent;
 }
diff --git a/src/hooks/useInsertionEffect.tsx b/src/hooks/useInsertionEffect.tsx
diff --git a/tests/legacy.spec.tsx b/tests/legacy.spec.tsx
@@ -1,5 +1,7 @@
 import { render } from '@testing-library/react';
 import * as React from 'react';
+import { useLayoutEffect } from 'react';
+import { expect } from 'vitest';
 import type { CSSInterpolation } from '../src';
 import { Theme, useCacheToken, useStyleRegister } from '../src';
 
@@ -50,14 +52,15 @@ describe('legacy React version', () => {
 
   interface BoxProps {
     propToken?: DesignToken;
+    children?: React.ReactNode;
   }
 
-  const Box = ({ propToken = baseToken }: BoxProps) => {
+  const Box = ({ propToken = baseToken, children }: BoxProps) => {
     const [token] = useCacheToken<DerivativeToken>(theme, [propToken]);
 
     useStyleRegister({ theme, token, path: ['.box'] }, () => [genStyle(token)]);
 
-    return <div className="box" />;
+    return <div className="box">{children}</div>;
   };
 
   // We will not remove style immediately,
@@ -106,4 +109,27 @@ describe('legacy React version', () => {
 
     test('StrictMode', (ele) => <React.StrictMode>{ele}</React.StrictMode>);
   });
+
+  it('should not race with other useLayoutEffect', () => {
+    let styleCount = 0;
+
+    const Child = () => {
+      useLayoutEffect(() => {
+        styleCount = document.head.querySelectorAll('style').length;
+      }, []);
+
+      return null;
+    };
+    const Demo = () => {
+      return (
+        <Box>
+          <Child />
+        </Box>
+      );
+    };
+
+    render(<Demo />);
+
+    expect(styleCount).toBe(1);
+  });
 });
