Skip to content

Commit 4a2d86b

Browse files
acdliteacdlite
committed
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.
1 parent 9dfbd9f commit 4a2d86b

File tree

5 files changed

+119
-22
lines changed

5 files changed

+119
-22
lines changed

‎packages/react-reconciler/src/ReactFiberHooks.new.js‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -732,10 +732,19 @@ export function bailoutHooks(
732732
}
733733

734734
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+
735742
// We can assume the previous dispatcher is always this one, since we set it
736743
// at the beginning of the render phase and there's no re-entrance.
737744
ReactCurrentDispatcher.current = ContextOnlyDispatcher;
745+
}
738746

747+
export function resetHooksOnUnwind(): void {
739748
if (didScheduleRenderPhaseUpdate) {
740749
// There were render phase updates. These are only valid for this render
741750
// phase, which we are now aborting. Remove the updates from the queues so

‎packages/react-reconciler/src/ReactFiberHooks.old.js‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -732,10 +732,19 @@ export function bailoutHooks(
732732
}
733733

734734
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+
735742
// We can assume the previous dispatcher is always this one, since we set it
736743
// at the beginning of the render phase and there's no re-entrance.
737744
ReactCurrentDispatcher.current = ContextOnlyDispatcher;
745+
}
738746

747+
export function resetHooksOnUnwind(): void {
739748
if (didScheduleRenderPhaseUpdate) {
740749
// There were render phase updates. These are only valid for this render
741750
// phase, which we are now aborting. Remove the updates from the queues so

‎packages/react-reconciler/src/ReactFiberWorkLoop.new.js‎

Lines changed: 35 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -210,6 +210,7 @@ import {enqueueUpdate} from './ReactFiberClassUpdateQueue.new';
210210
import {resetContextDependencies} from './ReactFiberNewContext.new';
211211
import {
212212
resetHooksAfterThrow,
213+
resetHooksOnUnwind,
213214
ContextOnlyDispatcher,
214215
} from './ReactFiberHooks.new';
215216
import {DefaultCacheDispatcher} from './ReactFiberCache.new';
@@ -1719,10 +1720,17 @@ function prepareFreshStack(root: FiberRoot, lanes: Lanes): Fiber {
17191720
}
17201721

17211722
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+
}
17261734
while (interruptedWork !== null) {
17271735
const current = interruptedWork.alternate;
17281736
unwindInterruptedWork(
@@ -1759,13 +1767,30 @@ function prepareFreshStack(root: FiberRoot, lanes: Lanes): Fiber {
17591767
return rootWorkInProgress;
17601768
}
17611769

1762-
function handleThrow(root, thrownValue): void {
1770+
function resetSuspendedWorkLoopOnUnwind() {
17631771
// Reset module-level state that was set during the render phase.
17641772
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.
17651792
resetHooksAfterThrow();
17661793
resetCurrentDebugFiberInDEV();
1767-
// TODO: I found and added this missing line while investigating a
1768-
// separate issue. Write a regression test using string refs.
17691794
ReactCurrentOwner.current = null;
17701795

17711796
if (thrownValue === SuspenseException) {
@@ -2317,6 +2342,7 @@ function replaySuspendedUnitOfWork(
23172342
// Instead of unwinding the stack and potentially showing a fallback, unwind
23182343
// only the last stack frame, reset the fiber, and try rendering it again.
23192344
const current = unitOfWork.alternate;
2345+
resetSuspendedWorkLoopOnUnwind();
23202346
unwindInterruptedWork(current, unitOfWork, workInProgressRootRenderLanes);
23212347
unitOfWork = workInProgress = resetWorkInProgress(unitOfWork, renderLanes);
23222348

@@ -2355,6 +2381,7 @@ function unwindSuspendedUnitOfWork(unitOfWork: Fiber, thrownValue: mixed) {
23552381
// Return to the normal work loop. This will unwind the stack, and potentially
23562382
// result in showing a fallback.
23572383
workInProgressSuspendedThenableState = null;
2384+
resetSuspendedWorkLoopOnUnwind();
23582385

23592386
const returnFiber = unitOfWork.return;
23602387
if (returnFiber === null || workInProgressRoot === null) {
@@ -3707,14 +3734,11 @@ if (__DEV__ && replayFailedUnitOfWorkWithInvokeGuardedCallback) {
37073734
throw originalError;
37083735
}
37093736

3710-
// Keep this code in sync with handleThrow; any changes here must have
3711-
// corresponding changes there.
3712-
resetContextDependencies();
3713-
resetHooksAfterThrow();
37143737
// Don't reset current debug fiber, since we're about to work on the
37153738
// same fiber again.
37163739

37173740
// Unwind the failed stack frame
3741+
resetSuspendedWorkLoopOnUnwind();
37183742
unwindInterruptedWork(current, unitOfWork, workInProgressRootRenderLanes);
37193743

37203744
// Restore the original properties of the fiber.

‎packages/react-reconciler/src/ReactFiberWorkLoop.old.js‎

Lines changed: 35 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -210,6 +210,7 @@ import {enqueueUpdate} from './ReactFiberClassUpdateQueue.old';
210210
import {resetContextDependencies} from './ReactFiberNewContext.old';
211211
import {
212212
resetHooksAfterThrow,
213+
resetHooksOnUnwind,
213214
ContextOnlyDispatcher,
214215
} from './ReactFiberHooks.old';
215216
import {DefaultCacheDispatcher} from './ReactFiberCache.old';
@@ -1719,10 +1720,17 @@ function prepareFreshStack(root: FiberRoot, lanes: Lanes): Fiber {
17191720
}
17201721

17211722
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+
}
17261734
while (interruptedWork !== null) {
17271735
const current = interruptedWork.alternate;
17281736
unwindInterruptedWork(
@@ -1759,13 +1767,30 @@ function prepareFreshStack(root: FiberRoot, lanes: Lanes): Fiber {
17591767
return rootWorkInProgress;
17601768
}
17611769

1762-
function handleThrow(root, thrownValue): void {
1770+
function resetSuspendedWorkLoopOnUnwind() {
17631771
// Reset module-level state that was set during the render phase.
17641772
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.
17651792
resetHooksAfterThrow();
17661793
resetCurrentDebugFiberInDEV();
1767-
// TODO: I found and added this missing line while investigating a
1768-
// separate issue. Write a regression test using string refs.
17691794
ReactCurrentOwner.current = null;
17701795

17711796
if (thrownValue === SuspenseException) {
@@ -2317,6 +2342,7 @@ function replaySuspendedUnitOfWork(
23172342
// Instead of unwinding the stack and potentially showing a fallback, unwind
23182343
// only the last stack frame, reset the fiber, and try rendering it again.
23192344
const current = unitOfWork.alternate;
2345+
resetSuspendedWorkLoopOnUnwind();
23202346
unwindInterruptedWork(current, unitOfWork, workInProgressRootRenderLanes);
23212347
unitOfWork = workInProgress = resetWorkInProgress(unitOfWork, renderLanes);
23222348

@@ -2355,6 +2381,7 @@ function unwindSuspendedUnitOfWork(unitOfWork: Fiber, thrownValue: mixed) {
23552381
// Return to the normal work loop. This will unwind the stack, and potentially
23562382
// result in showing a fallback.
23572383
workInProgressSuspendedThenableState = null;
2384+
resetSuspendedWorkLoopOnUnwind();
23582385

23592386
const returnFiber = unitOfWork.return;
23602387
if (returnFiber === null || workInProgressRoot === null) {
@@ -3707,14 +3734,11 @@ if (__DEV__ && replayFailedUnitOfWorkWithInvokeGuardedCallback) {
37073734
throw originalError;
37083735
}
37093736

3710-
// Keep this code in sync with handleThrow; any changes here must have
3711-
// corresponding changes there.
3712-
resetContextDependencies();
3713-
resetHooksAfterThrow();
37143737
// Don't reset current debug fiber, since we're about to work on the
37153738
// same fiber again.
37163739

37173740
// Unwind the failed stack frame
3741+
resetSuspendedWorkLoopOnUnwind();
37183742
unwindInterruptedWork(current, unitOfWork, workInProgressRootRenderLanes);
37193743

37203744
// Restore the original properties of the fiber.

‎packages/react-reconciler/src/__tests__/ReactThenable-test.js‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -643,6 +643,37 @@ describe('ReactThenable', () => {
643643
expect(Scheduler).toHaveYielded(['Something different']);
644644
});
645645

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+
646677
// @gate enableUseHook
647678
test('unwraps thenable that fulfills synchronously without suspending', async () => {
648679
function App() {

0 commit comments

Comments
 (0)