Reuse hooks when replaying a suspended component
When a component suspends, under some conditions, we can wait for the data to resolve and replay the component without unwinding the stack or showing a fallback in the interim. When we do this, we reuse the promises that were unwrapped during the previous attempts, so that if they aren't memoized, the result can still be used. We should do the same for all hooks. That way, if you _do_ memoize an async function call with useMemo, it won't be called again during the replay. This effectively gives you a local version of the functionality provided by `cache`, using the normal memoization patterns that have long existed in React.
Andrew Clark committed
Nov 2, 2022 at 21:43 UTC
33e3d2878e9ec82c65468316ffcc473e5288bb87
7 files changed
+411
-20
packages/react-reconciler/src/ReactFiberBeginWork.new.js
+53
@@ -38,6 +38,8 @@ import type {
38
import type {UpdateQueue} from './ReactFiberClassUpdateQueue.new';
39
import type {RootState} from './ReactFiberRoot.new';
40
import type {TracingMarkerInstance} from './ReactFiberTracingMarkerComponent.new';
41
+import type {ThenableState} from './ReactFiberThenable.new';
42
+
43
import checkPropTypes from 'shared/checkPropTypes';
44
import {
45
markComponentRenderStarted,
@@ -203,6 +205,7 @@ import {
205
renderWithHooks,
206
checkDidRenderIdHook,
207
bailoutHooks,
208
+ replaySuspendedComponentWithHooks,
209
} from './ReactFiberHooks.new';
210
import {stopProfilerTimerIfRunning} from './ReactProfilerTimer.new';
211
import {
@@ -1159,6 +1162,56 @@ function updateFunctionComponent(
1162
return workInProgress.child;
1163
}
1164
1165
+export function replayFunctionComponent(
1166
+ current: Fiber | null,
1167
+ workInProgress: Fiber,
1168
+ nextProps: any,
1169
+ Component: any,
1170
+ prevThenableState: ThenableState,
1171
+ renderLanes: Lanes,
1172
+): Fiber | null {
1173
+ // This function is used to replay a component that previously suspended,
1174
+ // after its data resolves. It's a simplified version of
1175
+ // updateFunctionComponent that reuses the hooks from the previous attempt.
1176
+
1177
+ let context;
1178
+ if (!disableLegacyContext) {
1179
+ const unmaskedContext = getUnmaskedContext(workInProgress, Component, true);
1180
+ context = getMaskedContext(workInProgress, unmaskedContext);
1181
+ }
1182
+
1183
+ prepareToReadContext(workInProgress, renderLanes);
1184
+ if (enableSchedulingProfiler) {
1185
+ markComponentRenderStarted(workInProgress);
1186
+ }
1187
+ const nextChildren = replaySuspendedComponentWithHooks(
1188
+ current,
1189
+ workInProgress,
1190
+ Component,
1191
+ nextProps,
1192
+ context,
1193
+ prevThenableState,
1194
+ );
1195
+ const hasId = checkDidRenderIdHook();
1196
+ if (enableSchedulingProfiler) {
1197
+ markComponentRenderStopped();
1198
+ }
1199
+
1200
+ if (current !== null && !didReceiveUpdate) {
1201
+ bailoutHooks(current, workInProgress, renderLanes);
1202
+ return bailoutOnAlreadyFinishedWork(current, workInProgress, renderLanes);
1203
+ }
1204
+
1205
+ if (getIsHydrating() && hasId) {
1206
+ pushMaterializedTreeId(workInProgress);
1207
+ }
1208
+
1209
+ // React DevTools reads this flag.
1210
+ workInProgress.flags |= PerformedWork;
1211
+ reconcileChildren(current, workInProgress, nextChildren, renderLanes);
1212
+ return workInProgress.child;
1213
+}
1214
+
1215
function updateClassComponent(
1216
current: Fiber | null,
1217
workInProgress: Fiber,
packages/react-reconciler/src/ReactFiberBeginWork.old.js
+53
@@ -38,6 +38,8 @@ import type {
38
import type {UpdateQueue} from './ReactFiberClassUpdateQueue.old';
39
import type {RootState} from './ReactFiberRoot.old';
40
import type {TracingMarkerInstance} from './ReactFiberTracingMarkerComponent.old';
41
+import type {ThenableState} from './ReactFiberThenable.old';
42
+
43
import checkPropTypes from 'shared/checkPropTypes';
44
import {
45
markComponentRenderStarted,
@@ -203,6 +205,7 @@ import {
205
renderWithHooks,
206
checkDidRenderIdHook,
207
bailoutHooks,
208
+ replaySuspendedComponentWithHooks,
209
} from './ReactFiberHooks.old';
210
import {stopProfilerTimerIfRunning} from './ReactProfilerTimer.old';
211
import {
@@ -1159,6 +1162,56 @@ function updateFunctionComponent(
1162
return workInProgress.child;
1163
}
1164
1165
+export function replayFunctionComponent(
1166
+ current: Fiber | null,
1167
+ workInProgress: Fiber,
1168
+ nextProps: any,
1169
+ Component: any,
1170
+ prevThenableState: ThenableState,
1171
+ renderLanes: Lanes,
1172
+): Fiber | null {
1173
+ // This function is used to replay a component that previously suspended,
1174
+ // after its data resolves. It's a simplified version of
1175
+ // updateFunctionComponent that reuses the hooks from the previous attempt.
1176
+
1177
+ let context;
1178
+ if (!disableLegacyContext) {
1179
+ const unmaskedContext = getUnmaskedContext(workInProgress, Component, true);
1180
+ context = getMaskedContext(workInProgress, unmaskedContext);
1181
+ }
1182
+
1183
+ prepareToReadContext(workInProgress, renderLanes);
1184
+ if (enableSchedulingProfiler) {
1185
+ markComponentRenderStarted(workInProgress);
1186
+ }
1187
+ const nextChildren = replaySuspendedComponentWithHooks(
1188
+ current,
1189
+ workInProgress,
1190
+ Component,
1191
+ nextProps,
1192
+ context,
1193
+ prevThenableState,
1194
+ );
1195
+ const hasId = checkDidRenderIdHook();
1196
+ if (enableSchedulingProfiler) {
1197
+ markComponentRenderStopped();
1198
+ }
1199
+
1200
+ if (current !== null && !didReceiveUpdate) {
1201
+ bailoutHooks(current, workInProgress, renderLanes);
1202
+ return bailoutOnAlreadyFinishedWork(current, workInProgress, renderLanes);
1203
+ }
1204
+
1205
+ if (getIsHydrating() && hasId) {
1206
+ pushMaterializedTreeId(workInProgress);
1207
+ }
1208
+
1209
+ // React DevTools reads this flag.
1210
+ workInProgress.flags |= PerformedWork;
1211
+ reconcileChildren(current, workInProgress, nextChildren, renderLanes);
1212
+ return workInProgress.child;
1213
+}
1214
+
1215
function updateClassComponent(
1216
current: Fiber | null,
1217
workInProgress: Fiber,
packages/react-reconciler/src/ReactFiberHooks.new.js
+40
@@ -545,6 +545,12 @@ export function renderWithHooks<Props, SecondArg>(
545
}
546
}
547
548
+ finishRenderingHooks(current, workInProgress);
549
+
550
+ return children;
551
+}
552
+
553
+function finishRenderingHooks(current: Fiber | null, workInProgress: Fiber) {
554
// We can assume the previous dispatcher is always this one, since we set it
555
// at the beginning of the render phase and there's no re-entrance.
556
ReactCurrentDispatcher.current = ContextOnlyDispatcher;
@@ -638,7 +644,41 @@ export function renderWithHooks<Props, SecondArg>(
644
}
645
}
646
}
647
+}
648
649
+export function replaySuspendedComponentWithHooks<Props, SecondArg>(
650
+ current: Fiber | null,
651
+ workInProgress: Fiber,
652
+ Component: (p: Props, arg: SecondArg) => any,
653
+ props: Props,
654
+ secondArg: SecondArg,
655
+ prevThenableState: ThenableState | null,
656
+): any {
657
+ // This function is used to replay a component that previously suspended,
658
+ // after its data resolves.
659
+ //
660
+ // It's a simplified version of renderWithHooks, but it doesn't need to do
661
+ // most of the set up work because they weren't reset when we suspended; they
662
+ // only get reset when the component either completes (finishRenderingHooks)
663
+ // or unwinds (resetHooksOnUnwind).
664
+ if (__DEV__) {
665
+ hookTypesDev =
666
+ current !== null
667
+ ? ((current._debugHookTypes: any): Array<HookType>)
668
+ : null;
669
+ hookTypesUpdateIndexDev = -1;
670
+ // Used for hot reloading:
671
+ ignorePreviousDependencies =
672
+ current !== null && current.type !== workInProgress.type;
673
+ }
674
+ const children = renderWithHooksAgain(
675
+ workInProgress,
676
+ Component,
677
+ props,
678
+ secondArg,
679
+ prevThenableState,
680
+ );
681
+ finishRenderingHooks(current, workInProgress);
682
return children;
683
}
684
packages/react-reconciler/src/ReactFiberHooks.old.js
+40
@@ -545,6 +545,12 @@ export function renderWithHooks<Props, SecondArg>(
545
}
546
}
547
548
+ finishRenderingHooks(current, workInProgress);
549
+
550
+ return children;
551
+}
552
+
553
+function finishRenderingHooks(current: Fiber | null, workInProgress: Fiber) {
554
// We can assume the previous dispatcher is always this one, since we set it
555
// at the beginning of the render phase and there's no re-entrance.
556
ReactCurrentDispatcher.current = ContextOnlyDispatcher;
@@ -638,7 +644,41 @@ export function renderWithHooks<Props, SecondArg>(
644
}
645
}
646
}
647
+}
648
649
+export function replaySuspendedComponentWithHooks<Props, SecondArg>(
650
+ current: Fiber | null,
651
+ workInProgress: Fiber,
652
+ Component: (p: Props, arg: SecondArg) => any,
653
+ props: Props,
654
+ secondArg: SecondArg,
655
+ prevThenableState: ThenableState | null,
656
+): any {
657
+ // This function is used to replay a component that previously suspended,
658
+ // after its data resolves.
659
+ //
660
+ // It's a simplified version of renderWithHooks, but it doesn't need to do
661
+ // most of the set up work because they weren't reset when we suspended; they
662
+ // only get reset when the component either completes (finishRenderingHooks)
663
+ // or unwinds (resetHooksOnUnwind).
664
+ if (__DEV__) {
665
+ hookTypesDev =
666
+ current !== null
667
+ ? ((current._debugHookTypes: any): Array<HookType>)
668
+ : null;
669
+ hookTypesUpdateIndexDev = -1;
670
+ // Used for hot reloading:
671
+ ignorePreviousDependencies =
672
+ current !== null && current.type !== workInProgress.type;
673
+ }
674
+ const children = renderWithHooksAgain(
675
+ workInProgress,
676
+ Component,
677
+ props,
678
+ secondArg,
679
+ prevThenableState,
680
+ );
681
+ finishRenderingHooks(current, workInProgress);
682
return children;
683
}
684
packages/react-reconciler/src/ReactFiberWorkLoop.new.js
+69
-10
@@ -181,6 +181,7 @@ import {requestCurrentTransition, NoTransition} from './ReactFiberTransition';
181
import {
182
SelectiveHydrationException,
183
beginWork as originalBeginWork,
184
+ replayFunctionComponent,
185
} from './ReactFiberBeginWork.new';
186
import {completeWork} from './ReactFiberCompleteWork.new';
187
import {unwindWork, unwindInterruptedWork} from './ReactFiberUnwindWork.new';
@@ -282,6 +283,7 @@ import {
283
getSuspenseHandler,
284
isBadSuspenseFallback,
285
} from './ReactFiberSuspenseContext.new';
286
+import {resolveDefaultProps} from './ReactFiberLazyComponent.new';
287
288
const ceil = Math.ceil;
289
@@ -2353,22 +2355,79 @@ function replaySuspendedUnitOfWork(
2355
// This is a fork of performUnitOfWork specifcally for replaying a fiber that
2356
// just suspended.
2357
//
2356
- // Instead of unwinding the stack and potentially showing a fallback, unwind
2357
- // only the last stack frame, reset the fiber, and try rendering it again.
2358
const current = unitOfWork.alternate;
2359
- resetSuspendedWorkLoopOnUnwind();
2360
- unwindInterruptedWork(current, unitOfWork, workInProgressRootRenderLanes);
2361
- unitOfWork = workInProgress = resetWorkInProgress(unitOfWork, renderLanes);
2362
-
2359
setCurrentDebugFiberInDEV(unitOfWork);
2360
2361
let next;
2366
- if (enableProfilerTimer && (unitOfWork.mode & ProfileMode) !== NoMode) {
2362
+ setCurrentDebugFiberInDEV(unitOfWork);
2363
+ const isProfilingMode =
2364
+ enableProfilerTimer && (unitOfWork.mode & ProfileMode) !== NoMode;
2365
+ if (isProfilingMode) {
2366
startProfilerTimer(unitOfWork);
2368
- next = beginWork(current, unitOfWork, renderLanes);
2367
+ }
2368
+ switch (unitOfWork.tag) {
2369
+ case IndeterminateComponent: {
2370
+ // Because it suspended with `use`, we can assume it's a
2371
+ // function component.
2372
+ unitOfWork.tag = FunctionComponent;
2373
+ // Fallthrough to the next branch.
2374
+ }
2375
+ // eslint-disable-next-line no-fallthrough
2376
+ case FunctionComponent:
2377
+ case ForwardRef: {
2378
+ // Resolve `defaultProps`. This logic is copied from `beginWork`.
2379
+ // TODO: Consider moving this switch statement into that module. Also,
2380
+ // could maybe use this as an opportunity to say `use` doesn't work with
2381
+ // `defaultProps` :)
2382
+ const Component = unitOfWork.type;
2383
+ const unresolvedProps = unitOfWork.pendingProps;
2384
+ const resolvedProps =
2385
+ unitOfWork.elementType === Component
2386
+ ? unresolvedProps
2387
+ : resolveDefaultProps(Component, unresolvedProps);
2388
+ next = replayFunctionComponent(
2389
+ current,
2390
+ unitOfWork,
2391
+ resolvedProps,
2392
+ Component,
2393
+ thenableState,
2394
+ workInProgressRootRenderLanes,
2395
+ );
2396
+ break;
2397
+ }
2398
+ case SimpleMemoComponent: {
2399
+ const Component = unitOfWork.type;
2400
+ const nextProps = unitOfWork.pendingProps;
2401
+ next = replayFunctionComponent(
2402
+ current,
2403
+ unitOfWork,
2404
+ nextProps,
2405
+ Component,
2406
+ thenableState,
2407
+ workInProgressRootRenderLanes,
2408
+ );
2409
+ break;
2410
+ }
2411
+ default: {
2412
+ if (__DEV__) {
2413
+ console.error(
2414
+ 'Unexpected type of work: %s, Currently only function ' +
2415
+ 'components are replayed after suspending. This is a bug in React.',
2416
+ unitOfWork.tag,
2417
+ );
2418
+ }
2419
+ resetSuspendedWorkLoopOnUnwind();
2420
+ unwindInterruptedWork(current, unitOfWork, workInProgressRootRenderLanes);
2421
+ unitOfWork = workInProgress = resetWorkInProgress(
2422
+ unitOfWork,
2423
+ renderLanes,
2424
+ );
2425
+ next = beginWork(current, unitOfWork, renderLanes);
2426
+ break;
2427
+ }
2428
+ }
2429
+ if (isProfilingMode) {
2430
stopProfilerTimerIfRunningAndRecordDelta(unitOfWork, true);
2370
- } else {
2371
- next = beginWork(current, unitOfWork, renderLanes);
2431
}
2432
2433
// The begin phase finished successfully without suspending. Reset the state
packages/react-reconciler/src/ReactFiberWorkLoop.old.js
+69
-10
@@ -181,6 +181,7 @@ import {requestCurrentTransition, NoTransition} from './ReactFiberTransition';
181
import {
182
SelectiveHydrationException,
183
beginWork as originalBeginWork,
184
+ replayFunctionComponent,
185
} from './ReactFiberBeginWork.old';
186
import {completeWork} from './ReactFiberCompleteWork.old';
187
import {unwindWork, unwindInterruptedWork} from './ReactFiberUnwindWork.old';
@@ -282,6 +283,7 @@ import {
283
getSuspenseHandler,
284
isBadSuspenseFallback,
285
} from './ReactFiberSuspenseContext.old';
286
+import {resolveDefaultProps} from './ReactFiberLazyComponent.old';
287
288
const ceil = Math.ceil;
289
@@ -2353,22 +2355,79 @@ function replaySuspendedUnitOfWork(
2355
// This is a fork of performUnitOfWork specifcally for replaying a fiber that
2356
// just suspended.
2357
//
2356
- // Instead of unwinding the stack and potentially showing a fallback, unwind
2357
- // only the last stack frame, reset the fiber, and try rendering it again.
2358
const current = unitOfWork.alternate;
2359
- resetSuspendedWorkLoopOnUnwind();
2360
- unwindInterruptedWork(current, unitOfWork, workInProgressRootRenderLanes);
2361
- unitOfWork = workInProgress = resetWorkInProgress(unitOfWork, renderLanes);
2362
-
2359
setCurrentDebugFiberInDEV(unitOfWork);
2360
2361
let next;
2366
- if (enableProfilerTimer && (unitOfWork.mode & ProfileMode) !== NoMode) {
2362
+ setCurrentDebugFiberInDEV(unitOfWork);
2363
+ const isProfilingMode =
2364
+ enableProfilerTimer && (unitOfWork.mode & ProfileMode) !== NoMode;
2365
+ if (isProfilingMode) {
2366
startProfilerTimer(unitOfWork);
2368
- next = beginWork(current, unitOfWork, renderLanes);
2367
+ }
2368
+ switch (unitOfWork.tag) {
2369
+ case IndeterminateComponent: {
2370
+ // Because it suspended with `use`, we can assume it's a
2371
+ // function component.
2372
+ unitOfWork.tag = FunctionComponent;
2373
+ // Fallthrough to the next branch.
2374
+ }
2375
+ // eslint-disable-next-line no-fallthrough
2376
+ case FunctionComponent:
2377
+ case ForwardRef: {
2378
+ // Resolve `defaultProps`. This logic is copied from `beginWork`.
2379
+ // TODO: Consider moving this switch statement into that module. Also,
2380
+ // could maybe use this as an opportunity to say `use` doesn't work with
2381
+ // `defaultProps` :)
2382
+ const Component = unitOfWork.type;
2383
+ const unresolvedProps = unitOfWork.pendingProps;
2384
+ const resolvedProps =
2385
+ unitOfWork.elementType === Component
2386
+ ? unresolvedProps
2387
+ : resolveDefaultProps(Component, unresolvedProps);
2388
+ next = replayFunctionComponent(
2389
+ current,
2390
+ unitOfWork,
2391
+ resolvedProps,
2392
+ Component,
2393
+ thenableState,
2394
+ workInProgressRootRenderLanes,
2395
+ );
2396
+ break;
2397
+ }
2398
+ case SimpleMemoComponent: {
2399
+ const Component = unitOfWork.type;
2400
+ const nextProps = unitOfWork.pendingProps;
2401
+ next = replayFunctionComponent(
2402
+ current,
2403
+ unitOfWork,
2404
+ nextProps,
2405
+ Component,
2406
+ thenableState,
2407
+ workInProgressRootRenderLanes,
2408
+ );
2409
+ break;
2410
+ }
2411
+ default: {
2412
+ if (__DEV__) {
2413
+ console.error(
2414
+ 'Unexpected type of work: %s, Currently only function ' +
2415
+ 'components are replayed after suspending. This is a bug in React.',
2416
+ unitOfWork.tag,
2417
+ );
2418
+ }
2419
+ resetSuspendedWorkLoopOnUnwind();
2420
+ unwindInterruptedWork(current, unitOfWork, workInProgressRootRenderLanes);
2421
+ unitOfWork = workInProgress = resetWorkInProgress(
2422
+ unitOfWork,
2423
+ renderLanes,
2424
+ );
2425
+ next = beginWork(current, unitOfWork, renderLanes);
2426
+ break;
2427
+ }
2428
+ }
2429
+ if (isProfilingMode) {
2430
stopProfilerTimerIfRunningAndRecordDelta(unitOfWork, true);
2370
- } else {
2371
- next = beginWork(current, unitOfWork, renderLanes);
2431
}
2432
2433
// The begin phase finished successfully without suspending. Reset the state
packages/react-reconciler/src/__tests__/ReactThenable-test.js
+87
@@ -6,6 +6,7 @@ let Scheduler;
6
let act;
7
let use;
8
let useState;
9
+let useMemo;
10
let Suspense;
11
let startTransition;
12
let pendingTextRequests;
@@ -20,6 +21,7 @@ describe('ReactThenable', () => {
21
act = require('jest-react').act;
22
use = React.use;
23
useState = React.useState;
24
+ useMemo = React.useMemo;
25
Suspense = React.Suspense;
26
startTransition = React.startTransition;
27
@@ -789,4 +791,89 @@ describe('ReactThenable', () => {
791
]);
792
expect(root).toMatchRenderedOutput('(empty)');
793
});
794
+
795
+ test('when replaying a suspended component, reuses the hooks computed during the previous attempt', async () => {
796
+ function ExcitingText({text}) {
797
+ // This computes the uppercased version of some text. Pretend it's an
798
+ // expensive operation that we want to reuse.
799
+ const uppercaseText = useMemo(() => {
800
+ Scheduler.unstable_yieldValue('Compute uppercase: ' + text);
801
+ return text.toUpperCase();
802
+ }, [text]);
803
+
804
+ // This adds an exclamation point to the text. Pretend it's an async
805
+ // operation that is sent to a service for processing.
806
+ const exclamatoryText = use(getAsyncText(uppercaseText + '!'));
807
+
808
+ // This surrounds the text with sparkle emojis. The purpose in this test
809
+ // is to show that you can suspend in the middle of a sequence of hooks
810
+ // without breaking anything.
811
+ const sparklingText = useMemo(() => {
812
+ Scheduler.unstable_yieldValue('Add sparkles: ' + exclamatoryText);
813
+ return `✨ ${exclamatoryText} ✨`;
814
+ }, [exclamatoryText]);
815
+
816
+ return <Text text={sparklingText} />;
817
+ }
818
+
819
+ const root = ReactNoop.createRoot();
820
+ await act(async () => {
821
+ startTransition(() => {
822
+ root.render(<ExcitingText text="Hello" />);
823
+ });
824
+ });
825
+ // Suspends while we wait for the async service to respond.
826
+ expect(Scheduler).toHaveYielded([
827
+ 'Compute uppercase: Hello',
828
+ 'Async text requested [HELLO!]',
829
+ ]);
830
+ expect(root).toMatchRenderedOutput(null);
831
+
832
+ // The data is received.
833
+ await act(async () => {
834
+ resolveTextRequests('HELLO!');
835
+ });
836
+ expect(Scheduler).toHaveYielded([
837
+ // We shouldn't run the uppercase computation again, because we can reuse
838
+ // the computation from the previous attempt.
839
+ // 'Compute uppercase: Hello',
840
+
841
+ 'Async text requested [HELLO!]',
842
+ 'Add sparkles: HELLO!',
843
+ '✨ HELLO! ✨',
844
+ ]);
845
+ });
846
+
847
+ // @gate enableUseHook
848
+ test(
849
+ 'wrap an async function with useMemo to skip running the function ' +
850
+ 'twice when loading new data',
851
+ async () => {
852
+ function App({text}) {
853
+ const promiseForText = useMemo(async () => getAsyncText(text), [text]);
854
+ const asyncText = use(promiseForText);
855
+ return <Text text={asyncText} />;
856
+ }
857
+
858
+ const root = ReactNoop.createRoot();
859
+ await act(async () => {
860
+ startTransition(() => {
861
+ root.render(<App text="Hello" />);
862
+ });
863
+ });
864
+ expect(Scheduler).toHaveYielded(['Async text requested [Hello]']);
865
+ expect(root).toMatchRenderedOutput(null);
866
+
867
+ await act(async () => {
868
+ resolveTextRequests('Hello');
869
+ });
870
+ expect(Scheduler).toHaveYielded([
871
+ // We shouldn't request async text again, because the async function
872
+ // was memoized
873
+ // 'Async text requested [Hello]'
874
+
875
+ 'Hello',
876
+ ]);
877
+ },
878
+ );
879
});