From 727b3615287074ddaa28069bfbd2dfee8cf73575 Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Fri, 6 Sep 2024 14:10:56 -0700
Subject: [PATCH] Fix useMemoCache with setState in render

Fixes the bug that @alexmckenley and @mofeiZ found where setState-in-render can reset useMemoCache and cause an infinite loop. The bug was that renderWithHooksAgain() was not resetting hook state when rerendering (so useMemo values were preserved) but was resetting the updateQueue. This meant that the entire memo cache was cleared on a setState-in-render.

The fix here is to call a new helper function to clear the update queue. It nulls out other properties, but for memoCache it just sets the index back to zero.

ghstack-source-id: fc0947ce219334117075df6a4e33b39975af2bc4
Pull Request resolved: https://github.com/facebook/react/pull/30889
---
 .../react-reconciler/src/ReactFiberHooks.js   | 42 ++++++++++++++-----
 .../src/__tests__/useMemoCache-test.js        | 24 +++++------
 2 files changed, 42 insertions(+), 24 deletions(-)

diff --git a/packages/react-reconciler/src/ReactFiberHooks.js b/packages/react-reconciler/src/ReactFiberHooks.js
index 83b9d10fbf..ca02af3393 100644
--- a/packages/react-reconciler/src/ReactFiberHooks.js
+++ b/packages/react-reconciler/src/ReactFiberHooks.js
@@ -770,6 +770,12 @@ export function replaySuspendedComponentWithHooks<Props, SecondArg>(
     ignorePreviousDependencies =
       current !== null && current.type !== workInProgress.type;
   }
+  // renderWithHooks only resets the updateQueue but does not clear it, since
+  // it needs to work for both this case (suspense replay) as well as for double
+  // renders in dev and setState-in-render. However, for the suspense replay case
+  // we need to reset the updateQueue to correctly handle unmount effects, so we
+  // clear the queue here
+  workInProgress.updateQueue = null;
   const children = renderWithHooksAgain(
     workInProgress,
     Component,
@@ -828,7 +834,9 @@ function renderWithHooksAgain<Props, SecondArg>(
     currentHook = null;
     workInProgressHook = null;
 
-    workInProgress.updateQueue = null;
+    if (workInProgress.updateQueue != null) {
+      resetFunctionComponentUpdateQueue((workInProgress.updateQueue: any));
+    }
 
     if (__DEV__) {
       // Also validate hook order for cascading updates.
@@ -1101,6 +1109,22 @@ if (enableUseMemoCacheHook) {
   };
 }
 
+function resetFunctionComponentUpdateQueue(
+  updateQueue: FunctionComponentUpdateQueue,
+): void {
+  updateQueue.lastEffect = null;
+  updateQueue.events = null;
+  updateQueue.stores = null;
+  if (enableUseMemoCacheHook) {
+    if (updateQueue.memoCache != null) {
+      // NOTE: this function intentionally does not reset memoCache data. We reuse updateQueue for the memo
+      // cache to avoid increasing the size of fibers that don't need a cache, but we don't want to reset
+      // the cache when other properties are reset.
+      updateQueue.memoCache.index = 0;
+    }
+  }
+}
+
 function useThenable<T>(thenable: Thenable<T>): T {
   // Track the position of the thenable within this fiber.
   const index = thenableIndexCounter;
@@ -2496,17 +2520,15 @@ function pushEffect(
   if (componentUpdateQueue === null) {
     componentUpdateQueue = createFunctionComponentUpdateQueue();
     currentlyRenderingFiber.updateQueue = (componentUpdateQueue: any);
+  }
+  const lastEffect = componentUpdateQueue.lastEffect;
+  if (lastEffect === null) {
     componentUpdateQueue.lastEffect = effect.next = effect;
   } else {
-    const lastEffect = componentUpdateQueue.lastEffect;
-    if (lastEffect === null) {
-      componentUpdateQueue.lastEffect = effect.next = effect;
-    } else {
-      const firstEffect = lastEffect.next;
-      lastEffect.next = effect;
-      effect.next = firstEffect;
-      componentUpdateQueue.lastEffect = effect;
-    }
+    const firstEffect = lastEffect.next;
+    lastEffect.next = effect;
+    effect.next = firstEffect;
+    componentUpdateQueue.lastEffect = effect;
   }
   return effect;
 }
diff --git a/packages/react-reconciler/src/__tests__/useMemoCache-test.js b/packages/react-reconciler/src/__tests__/useMemoCache-test.js
index 8a1bd5dcdf..d5008fbae6 100644
--- a/packages/react-reconciler/src/__tests__/useMemoCache-test.js
+++ b/packages/react-reconciler/src/__tests__/useMemoCache-test.js
@@ -16,7 +16,6 @@ let assertLog;
 let useMemo;
 let useState;
 let useMemoCache;
-let waitForThrow;
 let MemoCacheSentinel;
 let ErrorBoundary;
 
@@ -32,7 +31,6 @@ describe('useMemoCache()', () => {
     useMemo = React.useMemo;
     useMemoCache = require('react/compiler-runtime').c;
     useState = React.useState;
-    waitForThrow = require('internal-test-utils').waitForThrow;
     MemoCacheSentinel = Symbol.for('react.memo_cache_sentinel');
 
     class _ErrorBoundary extends React.Component {
@@ -667,7 +665,7 @@ describe('useMemoCache()', () => {
     }
 
     // Baseline / source code
-    function useUserMemo(value) {
+    function useManualMemo(value) {
       return useMemo(() => [value], [value]);
     }
 
@@ -683,24 +681,22 @@ describe('useMemoCache()', () => {
     }
 
     /**
-     * Test case: note that the initial render never completes
+     * Test with useMemoCache
      */
     let root = ReactNoop.createRoot();
-    const IncorrectInfiniteComponent = makeComponent(useCompilerMemo);
-    root.render(<IncorrectInfiniteComponent value={2} />);
-    await waitForThrow(
-      'Too many re-renders. React limits the number of renders to prevent ' +
-        'an infinite loop.',
-    );
+    const CompilerMemoComponent = makeComponent(useCompilerMemo);
+    await act(() => {
+      root.render(<CompilerMemoComponent value={2} />);
+    });
+    expect(root).toMatchRenderedOutput(<div>2</div>);
 
     /**
-     * Baseline test: initial render is expected to complete after a retry
-     * (triggered by the setState)
+     * Test with useMemo
      */
     root = ReactNoop.createRoot();
-    const CorrectComponent = makeComponent(useUserMemo);
+    const HookMemoComponent = makeComponent(useManualMemo);
     await act(() => {
-      root.render(<CorrectComponent value={2} />);
+      root.render(<HookMemoComponent value={2} />);
     });
     expect(root).toMatchRenderedOutput(<div>2</div>);
   });