From 34c1d7b2da68b1f00f24547de17908c9eb4c0ecc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=BA=8C=E8=B4=A7=E7=88=B1=E5=90=83=E7=99=BD=E8=90=9D?=
 =?UTF-8?q?=E5=8D=9C?= <smith3816@gmail.com>
Date: Mon, 3 Jul 2023 17:23:19 +0800
Subject: [PATCH] fix: Suspense will make cache counter not sync (#120)

* fix: cache counter

* test: add StrictMode test

* fix: StrictMode broken

* chore: useEffect

* refactor: merg
---
 package.json                 |  2 +-
 src/hooks/useGlobalCache.tsx | 53 ++++++++++++++++++-------------
 tests/index.spec.tsx         | 61 ++++++++++++++++++++++--------------
 3 files changed, 71 insertions(+), 45 deletions(-)

diff --git a/package.json b/package.json
index 37c0b35..0b331de 100644
--- a/package.json
+++ b/package.json
@@ -43,7 +43,7 @@
     "@emotion/unitless": "^0.7.5",
     "classnames": "^2.3.1",
     "csstype": "^3.0.10",
-    "rc-util": "^5.27.0",
+    "rc-util": "^5.34.1",
     "stylis": "^4.0.13"
   },
   "devDependencies": {
diff --git a/src/hooks/useGlobalCache.tsx b/src/hooks/useGlobalCache.tsx
index e6a818a..e06a3af 100644
--- a/src/hooks/useGlobalCache.tsx
+++ b/src/hooks/useGlobalCache.tsx
@@ -1,6 +1,6 @@
 import * as React from 'react';
-import StyleContext from '../StyleContext';
 import type { KeyType } from '../Cache';
+import StyleContext from '../StyleContext';
 import useHMR from './useHMR';
 
 export default function useClientCache<CacheType>(
@@ -14,32 +14,44 @@ export default function useClientCache<CacheType>(
 
   const HMRUpdate = useHMR();
 
-  // Create cache
-  React.useMemo(
-    () => {
-      globalCache.update(fullPath, (prevCache) => {
-        const [times = 0, cache] = prevCache || [];
+  type UpdaterArgs = [times: number, cache: CacheType];
 
-        // HMR should always ignore cache since developer may change it
-        let tmpCache = cache;
-        if (process.env.NODE_ENV !== 'production' && cache && HMRUpdate) {
-          onCacheRemove?.(tmpCache, HMRUpdate);
-          tmpCache = null;
-        }
+  const buildCache = (updater?: (data: UpdaterArgs) => UpdaterArgs) => {
+    globalCache.update(fullPath, (prevCache) => {
+      const [times = 0, cache] = prevCache || [];
 
-        const mergedCache = tmpCache || cacheFn();
+      // HMR should always ignore cache since developer may change it
+      let tmpCache = cache;
+      if (process.env.NODE_ENV !== 'production' && cache && HMRUpdate) {
+        onCacheRemove?.(tmpCache, HMRUpdate);
+        tmpCache = null;
+      }
 
-        return [times + 1, mergedCache];
-      });
-    },
+      const mergedCache = tmpCache || cacheFn();
+
+      const data: UpdaterArgs = [times, mergedCache];
+
+      // Call updater if need additional logic
+      return updater ? updater(data) : data;
+    });
+  };
+
+  // Create cache
+  React.useMemo(
+    () => buildCache(),
     /* eslint-disable react-hooks/exhaustive-deps */
     [fullPath.join('_')],
     /* eslint-enable */
   );
 
   // Remove if no need anymore
-  React.useEffect(
-    () => () => {
+  React.useEffect(() => {
+    // 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;
@@ -51,9 +63,8 @@ export default function useClientCache<CacheType>(
 
         return [times - 1, cache];
       });
-    },
-    fullPath,
-  );
+    };
+  }, fullPath);
 
   return globalCache.get(fullPath)![1];
 }
diff --git a/tests/index.spec.tsx b/tests/index.spec.tsx
index e353a6c..e1a32b7 100644
--- a/tests/index.spec.tsx
+++ b/tests/index.spec.tsx
@@ -91,31 +91,46 @@ describe('csssinjs', () => {
 
     // We will not remove style immediately,
     // but remove when second style patched.
-    it('remove old style to ensure style set only exist one', () => {
-      const getBox = (props?: BoxProps) => <Box {...props} />;
-
-      const { rerender } = render(getBox());
-      expect(document.head.querySelectorAll('style')).toHaveLength(1);
+    describe('remove old style to ensure style set only exist one', () => {
+      function test(
+        name: string,
+        wrapperFn?: (node: React.ReactElement) => React.ReactElement,
+      ) {
+        it(name, () => {
+          const getBox = (props?: BoxProps) => {
+            const box: React.ReactElement = <Box {...props} />;
+
+            return wrapperFn?.(box) || box;
+          };
+
+          const { rerender } = render(getBox());
+          expect(document.head.querySelectorAll('style')).toHaveLength(1);
+
+          // First change
+          rerender(
+            getBox({
+              propToken: {
+                primaryColor: 'red',
+              },
+            }),
+          );
+          expect(document.head.querySelectorAll('style')).toHaveLength(1);
+
+          // Second change
+          rerender(
+            getBox({
+              propToken: {
+                primaryColor: 'green',
+              },
+            }),
+          );
+          expect(document.head.querySelectorAll('style')).toHaveLength(1);
+        });
+      }
 
-      // First change
-      rerender(
-        getBox({
-          propToken: {
-            primaryColor: 'red',
-          },
-        }),
-      );
-      expect(document.head.querySelectorAll('style')).toHaveLength(1);
+      test('normal');
 
-      // Second change
-      rerender(
-        getBox({
-          propToken: {
-            primaryColor: 'green',
-          },
-        }),
-      );
-      expect(document.head.querySelectorAll('style')).toHaveLength(1);
+      test('StrictMode', (ele) => <React.StrictMode>{ele}</React.StrictMode>);
     });
 
     it('remove style when unmount', () => {