Update on "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-poisoned]

packages/react-reconciler/src/ReactFiberHooks.js

@@ -522,9 +522,7 @@ export function renderWithHooks<Props, SecondArg>(
   }
 
   workInProgress.memoizedState = null;
-  if (workInProgress.updateQueue != null) {
-    resetFunctionComponentUpdateQueue(workInProgress.updateQueue);
-  }
+  workInProgress.updateQueue = null;
   workInProgress.lanes = NoLanes;
 
   // The following should have already been reset
@@ -772,6 +770,7 @@ export function replaySuspendedComponentWithHooks<Props, SecondArg>(
     ignorePreviousDependencies =
       current !== null && current.type !== workInProgress.type;
   }
+  workInProgress.updateQueue = null;
   const children = renderWithHooksAgain(
     workInProgress,
     Component,
@@ -2516,17 +2515,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;
 }

packages/react-reconciler/src/__tests__/useMemoCache-test.js

@@ -575,7 +575,11 @@ describe('useMemoCache()', () => {
       'Some expensive processing... [A2]',
       'Suspend! [chunkB]',
 
-      ...(gate('enableSiblingPrerendering') ? ['Suspend! [chunkB]'] : []),
+      ...(gate('enableSiblingPrerendering')
+        ? gate('enableNoCloningMemoCache')
+          ? ['Suspend! [chunkB]']
+          : ['Some expensive processing... [A2]', 'Suspend! [chunkB]']
+        : []),
     ]);
 
     // The second chunk hasn't loaded yet, so we're still showing the
@@ -596,11 +600,26 @@ describe('useMemoCache()', () => {
     // Once the input has updated, we go back to rendering the transition.
     // We did not have process the first chunk again. We reused the
     // computation from the earlier attempt.
-    assertLog([
-      'Suspend! [chunkB]',
-
-      ...(gate('enableSiblingPrerendering') ? ['Suspend! [chunkB]'] : []),
-    ]);
+    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]']
+          : []),
+      ]);
+    }
 
     expect(root).toMatchRenderedOutput(
       <>
@@ -613,7 +632,18 @@ describe('useMemoCache()', () => {
     await act(() => updatedChunkB.resolve('B2'));
     // We did not have process the first chunk again. We reused the
     // computation from the earlier attempt.
-    assertLog([]);
+    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]']);
+    }
     expect(root).toMatchRenderedOutput(
       <>
         <div>Input: hi!</div>