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
Joe Savona committed
Sep 6, 2024 at 14:10 UTC
727b3615287074ddaa28069bfbd2dfee8cf73575
2 files changed
+42
-24
packages/react-reconciler/src/ReactFiberHooks.js
+32
-10
@@ -770,6 +770,12 @@ export function replaySuspendedComponentWithHooks<Props, SecondArg>(
770
ignorePreviousDependencies =
771
current !== null && current.type !== workInProgress.type;
772
}
773
+ // renderWithHooks only resets the updateQueue but does not clear it, since
774
+ // it needs to work for both this case (suspense replay) as well as for double
775
+ // renders in dev and setState-in-render. However, for the suspense replay case
776
+ // we need to reset the updateQueue to correctly handle unmount effects, so we
777
+ // clear the queue here
778
+ workInProgress.updateQueue = null;
779
const children = renderWithHooksAgain(
780
workInProgress,
781
Component,
@@ -828,7 +834,9 @@ function renderWithHooksAgain<Props, SecondArg>(
834
currentHook = null;
835
workInProgressHook = null;
836
831
- workInProgress.updateQueue = null;
837
+ if (workInProgress.updateQueue != null) {
838
+ resetFunctionComponentUpdateQueue((workInProgress.updateQueue: any));
839
+ }
840
841
if (__DEV__) {
842
// Also validate hook order for cascading updates.
@@ -1101,6 +1109,22 @@ if (enableUseMemoCacheHook) {
1109
};
1110
}
1111
1112
+function resetFunctionComponentUpdateQueue(
1113
+ updateQueue: FunctionComponentUpdateQueue,
1114
+): void {
1115
+ updateQueue.lastEffect = null;
1116
+ updateQueue.events = null;
1117
+ updateQueue.stores = null;
1118
+ if (enableUseMemoCacheHook) {
1119
+ if (updateQueue.memoCache != null) {
1120
+ // NOTE: this function intentionally does not reset memoCache data. We reuse updateQueue for the memo
1121
+ // cache to avoid increasing the size of fibers that don't need a cache, but we don't want to reset
1122
+ // the cache when other properties are reset.
1123
+ updateQueue.memoCache.index = 0;
1124
+ }
1125
+ }
1126
+}
1127
+
1128
function useThenable<T>(thenable: Thenable<T>): T {
1129
// Track the position of the thenable within this fiber.
1130
const index = thenableIndexCounter;
@@ -2496,17 +2520,15 @@ function pushEffect(
2520
if (componentUpdateQueue === null) {
2521
componentUpdateQueue = createFunctionComponentUpdateQueue();
2522
currentlyRenderingFiber.updateQueue = (componentUpdateQueue: any);
2523
+ }
2524
+ const lastEffect = componentUpdateQueue.lastEffect;
2525
+ if (lastEffect === null) {
2526
componentUpdateQueue.lastEffect = effect.next = effect;
2527
} else {
2501
- const lastEffect = componentUpdateQueue.lastEffect;
2502
- if (lastEffect === null) {
2503
- componentUpdateQueue.lastEffect = effect.next = effect;
2504
- } else {
2505
- const firstEffect = lastEffect.next;
2506
- lastEffect.next = effect;
2507
- effect.next = firstEffect;
2508
- componentUpdateQueue.lastEffect = effect;
2509
- }
2528
+ const firstEffect = lastEffect.next;
2529
+ lastEffect.next = effect;
2530
+ effect.next = firstEffect;
2531
+ componentUpdateQueue.lastEffect = effect;
2532
}
2533
return effect;
2534
}
packages/react-reconciler/src/__tests__/useMemoCache-test.js
+10
-14
@@ -16,7 +16,6 @@ let assertLog;
16
let useMemo;
17
let useState;
18
let useMemoCache;
19
-let waitForThrow;
19
let MemoCacheSentinel;
20
let ErrorBoundary;
21
@@ -32,7 +31,6 @@ describe('useMemoCache()', () => {
31
useMemo = React.useMemo;
32
useMemoCache = require('react/compiler-runtime').c;
33
useState = React.useState;
35
- waitForThrow = require('internal-test-utils').waitForThrow;
34
MemoCacheSentinel = Symbol.for('react.memo_cache_sentinel');
35
36
class _ErrorBoundary extends React.Component {
@@ -667,7 +665,7 @@ describe('useMemoCache()', () => {
665
}
666
667
// Baseline / source code
670
- function useUserMemo(value) {
668
+ function useManualMemo(value) {
669
return useMemo(() => [value], [value]);
670
}
671
@@ -683,24 +681,22 @@ describe('useMemoCache()', () => {
681
}
682
683
/**
686
- * Test case: note that the initial render never completes
684
+ * Test with useMemoCache
685
*/
686
let root = ReactNoop.createRoot();
689
- const IncorrectInfiniteComponent = makeComponent(useCompilerMemo);
690
- root.render(<IncorrectInfiniteComponent value={2} />);
691
- await waitForThrow(
692
- 'Too many re-renders. React limits the number of renders to prevent ' +
693
- 'an infinite loop.',
694
- );
687
+ const CompilerMemoComponent = makeComponent(useCompilerMemo);
688
+ await act(() => {
689
+ root.render(<CompilerMemoComponent value={2} />);
690
+ });
691
+ expect(root).toMatchRenderedOutput(<div>2</div>);
692
693
/**
697
- * Baseline test: initial render is expected to complete after a retry
698
- * (triggered by the setState)
694
+ * Test with useMemo
695
*/
696
root = ReactNoop.createRoot();
701
- const CorrectComponent = makeComponent(useUserMemo);
697
+ const HookMemoComponent = makeComponent(useManualMemo);
698
await act(() => {
703
- root.render(<CorrectComponent value={2} />);
699
+ root.render(<HookMemoComponent value={2} />);
700
});
701
expect(root).toMatchRenderedOutput(<div>2</div>);
702
});