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: d0508abb7babfcd7d96a166eabb4f8a7e19c9c54
Pull Request resolved: #30889

packages/react-reconciler/src/ReactFiberHooks.js

@@ -522,7 +522,9 @@ export function renderWithHooks<Props, SecondArg>(
   }
 
   workInProgress.memoizedState = null;
-  workInProgress.updateQueue = null;
+  if (workInProgress.updateQueue != null) {
+    resetFunctionComponentUpdateQueue(workInProgress.updateQueue);
+  }
   workInProgress.lanes = NoLanes;
 
   // The following should have already been reset
@@ -828,7 +830,9 @@ function renderWithHooksAgain<Props, SecondArg>(
     currentHook = null;
     workInProgressHook = null;
 
-    workInProgress.updateQueue = null;
+    if (workInProgress.updateQueue != null) {
+      resetFunctionComponentUpdateQueue(workInProgress.updateQueue);
+    }
 
     if (__DEV__) {
       // Also validate hook order for cascading updates.
@@ -1101,6 +1105,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;

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 {
@@ -577,11 +575,7 @@ describe('useMemoCache()', () => {
       'Some expensive processing... [A2]',
       'Suspend! [chunkB]',
 
-      ...(gate('enableSiblingPrerendering')
-        ? gate('enableNoCloningMemoCache')
-          ? ['Suspend! [chunkB]']
-          : ['Some expensive processing... [A2]', 'Suspend! [chunkB]']
-        : []),
+      ...(gate('enableSiblingPrerendering') ? ['Suspend! [chunkB]'] : []),
     ]);
 
     // The second chunk hasn't loaded yet, so we're still showing the
@@ -600,26 +594,13 @@ describe('useMemoCache()', () => {
     await act(() => setInput('hi!'));
 
     // Once the input has updated, we go back to rendering the transition.
-    if (gate(flags => flags.enableNoCloningMemoCache)) {
-      // We did not have process the first chunk again. We reused the
-      // computation from the earlier attempt.
-      assertLog([
-        'Suspend! [chunkB]',
-
-        ...(gate('enableSiblingPrerendering') ? ['Suspend! [chunkB]'] : []),
-      ]);
-    } else {
-      // Because we clone/reset the memo cache after every aborted attempt, we
-      // must process the first chunk again.
-      assertLog([
-        'Some expensive processing... [A2]',
-        'Suspend! [chunkB]',
-
-        ...(gate('enableSiblingPrerendering')
-          ? ['Some expensive processing... [A2]', 'Suspend! [chunkB]']
-          : []),
-      ]);
-    }
+    // We did not have process the first chunk again. We reused the
+    // computation from the earlier attempt.
+    assertLog([
+      'Suspend! [chunkB]',
+
+      ...(gate('enableSiblingPrerendering') ? ['Suspend! [chunkB]'] : []),
+    ]);
 
     expect(root).toMatchRenderedOutput(
       <>
@@ -630,18 +611,9 @@ describe('useMemoCache()', () => {
 
     // Finish loading the data.
     await act(() => updatedChunkB.resolve('B2'));
-    if (gate(flags => flags.enableNoCloningMemoCache)) {
-      // We did not have process the first chunk again. We reused the
-      // computation from the earlier attempt.
-      assertLog([]);
-    } else {
-      // Because we clone/reset the memo cache after every aborted attempt, we
-      // must process the first chunk again.
-      //
-      // That's three total times we've processed the first chunk, compared to
-      // just once when enableNoCloningMemoCache is on.
-      assertLog(['Some expensive processing... [A2]']);
-    }
+    // We did not have process the first chunk again. We reused the
+    // computation from the earlier attempt.
+    assertLog([]);
     expect(root).toMatchRenderedOutput(
       <>
         <div>Input: hi!</div>
@@ -667,7 +639,7 @@ describe('useMemoCache()', () => {
     }
 
     // Baseline / source code
-    function useUserMemo(value) {
+    function useManualMemo(value) {
       return useMemo(() => [value], [value]);
     }
 
@@ -683,24 +655,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>);
   });