Don't reset work loop until stack is unwound
When replaying a suspended function components, we want to reuse the hooks that were computed during the original render. Currently we reset the state of the hooks right after the component suspends (or throws an error). This is too early because it doesn't give us an opportunity to wait for the promise to resolve. This refactors the work loop to reset the hooks right before unwinding instead of right after throwing. It doesn't include any other changes yet, so there should be no observable behavioral change.
Andrew Clark committed
Nov 2, 2022 at 16:19 UTC
4a2d86bddbcef3e64bc404302cdbd9638af8801b
5 files changed
+119
-22
packages/react-reconciler/src/ReactFiberHooks.new.js
+9
@@ -732,10 +732,19 @@ export function bailoutHooks(
732
}
733
734
export function resetHooksAfterThrow(): void {
735
+ // This is called immediaetly after a throw. It shouldn't reset the entire
736
+ // module state, because the work loop might decide to replay the component
737
+ // again without rewinding.
738
+ //
739
+ // It should only reset things like the current dispatcher, to prevent hooks
740
+ // from being called outside of a component.
741
+
742
// We can assume the previous dispatcher is always this one, since we set it
743
// at the beginning of the render phase and there's no re-entrance.
744
ReactCurrentDispatcher.current = ContextOnlyDispatcher;
745
+}
746
747
+export function resetHooksOnUnwind(): void {
748
if (didScheduleRenderPhaseUpdate) {
749
// There were render phase updates. These are only valid for this render
750
// phase, which we are now aborting. Remove the updates from the queues so
packages/react-reconciler/src/ReactFiberHooks.old.js
+9
@@ -732,10 +732,19 @@ export function bailoutHooks(
732
}
733
734
export function resetHooksAfterThrow(): void {
735
+ // This is called immediaetly after a throw. It shouldn't reset the entire
736
+ // module state, because the work loop might decide to replay the component
737
+ // again without rewinding.
738
+ //
739
+ // It should only reset things like the current dispatcher, to prevent hooks
740
+ // from being called outside of a component.
741
+
742
// We can assume the previous dispatcher is always this one, since we set it
743
// at the beginning of the render phase and there's no re-entrance.
744
ReactCurrentDispatcher.current = ContextOnlyDispatcher;
745
+}
746
747
+export function resetHooksOnUnwind(): void {
748
if (didScheduleRenderPhaseUpdate) {
749
// There were render phase updates. These are only valid for this render
750
// phase, which we are now aborting. Remove the updates from the queues so
packages/react-reconciler/src/ReactFiberWorkLoop.new.js
+35
-11
@@ -210,6 +210,7 @@ import {enqueueUpdate} from './ReactFiberClassUpdateQueue.new';
210
import {resetContextDependencies} from './ReactFiberNewContext.new';
211
import {
212
resetHooksAfterThrow,
213
+ resetHooksOnUnwind,
214
ContextOnlyDispatcher,
215
} from './ReactFiberHooks.new';
216
import {DefaultCacheDispatcher} from './ReactFiberCache.new';
@@ -1719,10 +1720,17 @@ function prepareFreshStack(root: FiberRoot, lanes: Lanes): Fiber {
1720
}
1721
1722
if (workInProgress !== null) {
1722
- let interruptedWork =
1723
- workInProgressSuspendedReason === NotSuspended
1724
- ? workInProgress.return
1725
- : workInProgress;
1723
+ let interruptedWork;
1724
+ if (workInProgressSuspendedReason === NotSuspended) {
1725
+ // Normal case. Work-in-progress hasn't started yet. Unwind all
1726
+ // its parents.
1727
+ interruptedWork = workInProgress.return;
1728
+ } else {
1729
+ // Work-in-progress is in suspended state. Reset the work loop and unwind
1730
+ // both the suspended fiber and all its parents.
1731
+ resetSuspendedWorkLoopOnUnwind();
1732
+ interruptedWork = workInProgress;
1733
+ }
1734
while (interruptedWork !== null) {
1735
const current = interruptedWork.alternate;
1736
unwindInterruptedWork(
@@ -1759,13 +1767,30 @@ function prepareFreshStack(root: FiberRoot, lanes: Lanes): Fiber {
1767
return rootWorkInProgress;
1768
}
1769
1762
-function handleThrow(root, thrownValue): void {
1770
+function resetSuspendedWorkLoopOnUnwind() {
1771
// Reset module-level state that was set during the render phase.
1772
resetContextDependencies();
1773
+ resetHooksOnUnwind();
1774
+}
1775
+
1776
+function handleThrow(root, thrownValue): void {
1777
+ // A component threw an exception. Usually this is because it suspended, but
1778
+ // it also includes regular program errors.
1779
+ //
1780
+ // We're either going to unwind the stack to show a Suspense or error
1781
+ // boundary, or we're going to replay the component again. Like after a
1782
+ // promise resolves.
1783
+ //
1784
+ // Until we decide whether we're going to unwind or replay, we should preserve
1785
+ // the current state of the work loop without resetting anything.
1786
+ //
1787
+ // If we do decide to unwind the stack, module-level variables will be reset
1788
+ // in resetSuspendedWorkLoopOnUnwind.
1789
+
1790
+ // These should be reset immediately because they're only supposed to be set
1791
+ // when React is executing user code.
1792
resetHooksAfterThrow();
1793
resetCurrentDebugFiberInDEV();
1767
- // TODO: I found and added this missing line while investigating a
1768
- // separate issue. Write a regression test using string refs.
1794
ReactCurrentOwner.current = null;
1795
1796
if (thrownValue === SuspenseException) {
@@ -2317,6 +2342,7 @@ function replaySuspendedUnitOfWork(
2342
// Instead of unwinding the stack and potentially showing a fallback, unwind
2343
// only the last stack frame, reset the fiber, and try rendering it again.
2344
const current = unitOfWork.alternate;
2345
+ resetSuspendedWorkLoopOnUnwind();
2346
unwindInterruptedWork(current, unitOfWork, workInProgressRootRenderLanes);
2347
unitOfWork = workInProgress = resetWorkInProgress(unitOfWork, renderLanes);
2348
@@ -2355,6 +2381,7 @@ function unwindSuspendedUnitOfWork(unitOfWork: Fiber, thrownValue: mixed) {
2381
// Return to the normal work loop. This will unwind the stack, and potentially
2382
// result in showing a fallback.
2383
workInProgressSuspendedThenableState = null;
2384
+ resetSuspendedWorkLoopOnUnwind();
2385
2386
const returnFiber = unitOfWork.return;
2387
if (returnFiber === null || workInProgressRoot === null) {
@@ -3707,14 +3734,11 @@ if (__DEV__ && replayFailedUnitOfWorkWithInvokeGuardedCallback) {
3734
throw originalError;
3735
}
3736
3710
- // Keep this code in sync with handleThrow; any changes here must have
3711
- // corresponding changes there.
3712
- resetContextDependencies();
3713
- resetHooksAfterThrow();
3737
// Don't reset current debug fiber, since we're about to work on the
3738
// same fiber again.
3739
3740
// Unwind the failed stack frame
3741
+ resetSuspendedWorkLoopOnUnwind();
3742
unwindInterruptedWork(current, unitOfWork, workInProgressRootRenderLanes);
3743
3744
// Restore the original properties of the fiber.
packages/react-reconciler/src/ReactFiberWorkLoop.old.js
+35
-11
@@ -210,6 +210,7 @@ import {enqueueUpdate} from './ReactFiberClassUpdateQueue.old';
210
import {resetContextDependencies} from './ReactFiberNewContext.old';
211
import {
212
resetHooksAfterThrow,
213
+ resetHooksOnUnwind,
214
ContextOnlyDispatcher,
215
} from './ReactFiberHooks.old';
216
import {DefaultCacheDispatcher} from './ReactFiberCache.old';
@@ -1719,10 +1720,17 @@ function prepareFreshStack(root: FiberRoot, lanes: Lanes): Fiber {
1720
}
1721
1722
if (workInProgress !== null) {
1722
- let interruptedWork =
1723
- workInProgressSuspendedReason === NotSuspended
1724
- ? workInProgress.return
1725
- : workInProgress;
1723
+ let interruptedWork;
1724
+ if (workInProgressSuspendedReason === NotSuspended) {
1725
+ // Normal case. Work-in-progress hasn't started yet. Unwind all
1726
+ // its parents.
1727
+ interruptedWork = workInProgress.return;
1728
+ } else {
1729
+ // Work-in-progress is in suspended state. Reset the work loop and unwind
1730
+ // both the suspended fiber and all its parents.
1731
+ resetSuspendedWorkLoopOnUnwind();
1732
+ interruptedWork = workInProgress;
1733
+ }
1734
while (interruptedWork !== null) {
1735
const current = interruptedWork.alternate;
1736
unwindInterruptedWork(
@@ -1759,13 +1767,30 @@ function prepareFreshStack(root: FiberRoot, lanes: Lanes): Fiber {
1767
return rootWorkInProgress;
1768
}
1769
1762
-function handleThrow(root, thrownValue): void {
1770
+function resetSuspendedWorkLoopOnUnwind() {
1771
// Reset module-level state that was set during the render phase.
1772
resetContextDependencies();
1773
+ resetHooksOnUnwind();
1774
+}
1775
+
1776
+function handleThrow(root, thrownValue): void {
1777
+ // A component threw an exception. Usually this is because it suspended, but
1778
+ // it also includes regular program errors.
1779
+ //
1780
+ // We're either going to unwind the stack to show a Suspense or error
1781
+ // boundary, or we're going to replay the component again. Like after a
1782
+ // promise resolves.
1783
+ //
1784
+ // Until we decide whether we're going to unwind or replay, we should preserve
1785
+ // the current state of the work loop without resetting anything.
1786
+ //
1787
+ // If we do decide to unwind the stack, module-level variables will be reset
1788
+ // in resetSuspendedWorkLoopOnUnwind.
1789
+
1790
+ // These should be reset immediately because they're only supposed to be set
1791
+ // when React is executing user code.
1792
resetHooksAfterThrow();
1793
resetCurrentDebugFiberInDEV();
1767
- // TODO: I found and added this missing line while investigating a
1768
- // separate issue. Write a regression test using string refs.
1794
ReactCurrentOwner.current = null;
1795
1796
if (thrownValue === SuspenseException) {
@@ -2317,6 +2342,7 @@ function replaySuspendedUnitOfWork(
2342
// Instead of unwinding the stack and potentially showing a fallback, unwind
2343
// only the last stack frame, reset the fiber, and try rendering it again.
2344
const current = unitOfWork.alternate;
2345
+ resetSuspendedWorkLoopOnUnwind();
2346
unwindInterruptedWork(current, unitOfWork, workInProgressRootRenderLanes);
2347
unitOfWork = workInProgress = resetWorkInProgress(unitOfWork, renderLanes);
2348
@@ -2355,6 +2381,7 @@ function unwindSuspendedUnitOfWork(unitOfWork: Fiber, thrownValue: mixed) {
2381
// Return to the normal work loop. This will unwind the stack, and potentially
2382
// result in showing a fallback.
2383
workInProgressSuspendedThenableState = null;
2384
+ resetSuspendedWorkLoopOnUnwind();
2385
2386
const returnFiber = unitOfWork.return;
2387
if (returnFiber === null || workInProgressRoot === null) {
@@ -3707,14 +3734,11 @@ if (__DEV__ && replayFailedUnitOfWorkWithInvokeGuardedCallback) {
3734
throw originalError;
3735
}
3736
3710
- // Keep this code in sync with handleThrow; any changes here must have
3711
- // corresponding changes there.
3712
- resetContextDependencies();
3713
- resetHooksAfterThrow();
3737
// Don't reset current debug fiber, since we're about to work on the
3738
// same fiber again.
3739
3740
// Unwind the failed stack frame
3741
+ resetSuspendedWorkLoopOnUnwind();
3742
unwindInterruptedWork(current, unitOfWork, workInProgressRootRenderLanes);
3743
3744
// Restore the original properties of the fiber.
packages/react-reconciler/src/__tests__/ReactThenable-test.js
+31
@@ -643,6 +643,37 @@ describe('ReactThenable', () => {
643
expect(Scheduler).toHaveYielded(['Something different']);
644
});
645
646
+ // @gate enableUseHook
647
+ test('while suspended, hooks cannot be called (i.e. current dispatcher is unset correctly)', async () => {
648
+ function App() {
649
+ return <Text text={use(getAsyncText('Will never resolve'))} />;
650
+ }
651
+
652
+ const root = ReactNoop.createRoot();
653
+ await act(async () => {
654
+ root.render(<Suspense fallback={<Text text="Loading..." />} />);
655
+ });
656
+
657
+ await act(async () => {
658
+ startTransition(() => {
659
+ root.render(
660
+ <Suspense fallback={<Text text="Loading..." />}>
661
+ <App />
662
+ </Suspense>,
663
+ );
664
+ });
665
+ });
666
+ expect(Scheduler).toHaveYielded([
667
+ 'Async text requested [Will never resolve]',
668
+ ]);
669
+
670
+ // Calling a hook should error because we're oustide of a component.
671
+ expect(useState).toThrow(
672
+ 'Invalid hook call. Hooks can only be called inside of the body of a ' +
673
+ 'function component.',
674
+ );
675
+ });
676
+
677
// @gate enableUseHook
678
test('unwraps thenable that fulfills synchronously without suspending', async () => {
679
function App() {