Skip to content

Commit

Permalink
Don't reset work loop until stack is unwound
Browse files Browse the repository at this point in the history 
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.
  • Loading branch information
@acdlite
acdlite committed Nov 17, 2022
1 parent 9dfbd9f commit 4a2d86b
Show file tree
Hide file tree
Showing 5 changed files with 119 additions and 22 deletions.
9 changes: 9 additions & 0 deletions packages/react-reconciler/src/ReactFiberHooks.new.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -732,10 +732,19 @@ export function bailoutHooks(
}

export function resetHooksAfterThrow(): void {
// This is called immediaetly after a throw. It shouldn't reset the entire
// module state, because the work loop might decide to replay the component
// again without rewinding.
//
// It should only reset things like the current dispatcher, to prevent hooks
// from being called outside of a component.

// We can assume the previous dispatcher is always this one, since we set it
// at the beginning of the render phase and there's no re-entrance.
ReactCurrentDispatcher.current = ContextOnlyDispatcher;
}

export function resetHooksOnUnwind(): void {
if (didScheduleRenderPhaseUpdate) {
// There were render phase updates. These are only valid for this render
// phase, which we are now aborting. Remove the updates from the queues so
Expand Down
9 changes: 9 additions & 0 deletions packages/react-reconciler/src/ReactFiberHooks.old.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -732,10 +732,19 @@ export function bailoutHooks(
}

export function resetHooksAfterThrow(): void {
// This is called immediaetly after a throw. It shouldn't reset the entire
// module state, because the work loop might decide to replay the component
// again without rewinding.
//
// It should only reset things like the current dispatcher, to prevent hooks
// from being called outside of a component.

// We can assume the previous dispatcher is always this one, since we set it
// at the beginning of the render phase and there's no re-entrance.
ReactCurrentDispatcher.current = ContextOnlyDispatcher;
}

export function resetHooksOnUnwind(): void {
if (didScheduleRenderPhaseUpdate) {
// There were render phase updates. These are only valid for this render
// phase, which we are now aborting. Remove the updates from the queues so
Expand Down
46 changes: 35 additions & 11 deletions packages/react-reconciler/src/ReactFiberWorkLoop.new.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -210,6 +210,7 @@ import {enqueueUpdate} from './ReactFiberClassUpdateQueue.new';
import {resetContextDependencies} from './ReactFiberNewContext.new';
import {
resetHooksAfterThrow,
resetHooksOnUnwind,
ContextOnlyDispatcher,
} from './ReactFiberHooks.new';
import {DefaultCacheDispatcher} from './ReactFiberCache.new';
Expand Down Expand Up @@ -1719,10 +1720,17 @@ function prepareFreshStack(root: FiberRoot, lanes: Lanes): Fiber {
}

if (workInProgress !== null) {
let interruptedWork =
workInProgressSuspendedReason === NotSuspended
? workInProgress.return
: workInProgress;
let interruptedWork;
if (workInProgressSuspendedReason === NotSuspended) {
// Normal case. Work-in-progress hasn't started yet. Unwind all
// its parents.
interruptedWork = workInProgress.return;
} else {
// Work-in-progress is in suspended state. Reset the work loop and unwind
// both the suspended fiber and all its parents.
resetSuspendedWorkLoopOnUnwind();
interruptedWork = workInProgress;
}
while (interruptedWork !== null) {
const current = interruptedWork.alternate;
unwindInterruptedWork(
Expand Down Expand Up @@ -1759,13 +1767,30 @@ function prepareFreshStack(root: FiberRoot, lanes: Lanes): Fiber {
return rootWorkInProgress;
}

function handleThrow(root, thrownValue): void {
function resetSuspendedWorkLoopOnUnwind() {
// Reset module-level state that was set during the render phase.
resetContextDependencies();
resetHooksOnUnwind();
}

function handleThrow(root, thrownValue): void {
// A component threw an exception. Usually this is because it suspended, but
// it also includes regular program errors.
//
// We're either going to unwind the stack to show a Suspense or error
// boundary, or we're going to replay the component again. Like after a
// promise resolves.
//
// Until we decide whether we're going to unwind or replay, we should preserve
// the current state of the work loop without resetting anything.
//
// If we do decide to unwind the stack, module-level variables will be reset
// in resetSuspendedWorkLoopOnUnwind.

// These should be reset immediately because they're only supposed to be set
// when React is executing user code.
resetHooksAfterThrow();
resetCurrentDebugFiberInDEV();
// TODO: I found and added this missing line while investigating a
// separate issue. Write a regression test using string refs.
ReactCurrentOwner.current = null;

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

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

const returnFiber = unitOfWork.return;
if (returnFiber === null || workInProgressRoot === null) {
Expand Down Expand Up @@ -3707,14 +3734,11 @@ if (__DEV__ && replayFailedUnitOfWorkWithInvokeGuardedCallback) {
throw originalError;
}

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

// Unwind the failed stack frame
resetSuspendedWorkLoopOnUnwind();
unwindInterruptedWork(current, unitOfWork, workInProgressRootRenderLanes);

// Restore the original properties of the fiber.
Expand Down
46 changes: 35 additions & 11 deletions packages/react-reconciler/src/ReactFiberWorkLoop.old.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -210,6 +210,7 @@ import {enqueueUpdate} from './ReactFiberClassUpdateQueue.old';
import {resetContextDependencies} from './ReactFiberNewContext.old';
import {
resetHooksAfterThrow,
resetHooksOnUnwind,
ContextOnlyDispatcher,
} from './ReactFiberHooks.old';
import {DefaultCacheDispatcher} from './ReactFiberCache.old';
Expand Down Expand Up @@ -1719,10 +1720,17 @@ function prepareFreshStack(root: FiberRoot, lanes: Lanes): Fiber {
}

if (workInProgress !== null) {
let interruptedWork =
workInProgressSuspendedReason === NotSuspended
? workInProgress.return
: workInProgress;
let interruptedWork;
if (workInProgressSuspendedReason === NotSuspended) {
// Normal case. Work-in-progress hasn't started yet. Unwind all
// its parents.
interruptedWork = workInProgress.return;
} else {
// Work-in-progress is in suspended state. Reset the work loop and unwind
// both the suspended fiber and all its parents.
resetSuspendedWorkLoopOnUnwind();
interruptedWork = workInProgress;
}
while (interruptedWork !== null) {
const current = interruptedWork.alternate;
unwindInterruptedWork(
Expand Down Expand Up @@ -1759,13 +1767,30 @@ function prepareFreshStack(root: FiberRoot, lanes: Lanes): Fiber {
return rootWorkInProgress;
}

function handleThrow(root, thrownValue): void {
function resetSuspendedWorkLoopOnUnwind() {
// Reset module-level state that was set during the render phase.
resetContextDependencies();
resetHooksOnUnwind();
}

function handleThrow(root, thrownValue): void {
// A component threw an exception. Usually this is because it suspended, but
// it also includes regular program errors.
//
// We're either going to unwind the stack to show a Suspense or error
// boundary, or we're going to replay the component again. Like after a
// promise resolves.
//
// Until we decide whether we're going to unwind or replay, we should preserve
// the current state of the work loop without resetting anything.
//
// If we do decide to unwind the stack, module-level variables will be reset
// in resetSuspendedWorkLoopOnUnwind.

// These should be reset immediately because they're only supposed to be set
// when React is executing user code.
resetHooksAfterThrow();
resetCurrentDebugFiberInDEV();
// TODO: I found and added this missing line while investigating a
// separate issue. Write a regression test using string refs.
ReactCurrentOwner.current = null;

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

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

const returnFiber = unitOfWork.return;
if (returnFiber === null || workInProgressRoot === null) {
Expand Down Expand Up @@ -3707,14 +3734,11 @@ if (__DEV__ && replayFailedUnitOfWorkWithInvokeGuardedCallback) {
throw originalError;
}

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

// Unwind the failed stack frame
resetSuspendedWorkLoopOnUnwind();
unwindInterruptedWork(current, unitOfWork, workInProgressRootRenderLanes);

// Restore the original properties of the fiber.
Expand Down
31 changes: 31 additions & 0 deletions packages/react-reconciler/src/__tests__/ReactThenable-test.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -643,6 +643,37 @@ describe('ReactThenable', () => {
expect(Scheduler).toHaveYielded(['Something different']);
});

// @gate enableUseHook
test('while suspended, hooks cannot be called (i.e. current dispatcher is unset correctly)', async () => {
function App() {
return <Text text={use(getAsyncText('Will never resolve'))} />;
}

const root = ReactNoop.createRoot();
await act(async () => {
root.render(<Suspense fallback={<Text text="Loading..." />} />);
});

await act(async () => {
startTransition(() => {
root.render(
<Suspense fallback={<Text text="Loading..." />}>
<App />
</Suspense>,
);
});
});
expect(Scheduler).toHaveYielded([
'Async text requested [Will never resolve]',
]);

// Calling a hook should error because we're oustide of a component.
expect(useState).toThrow(
'Invalid hook call. Hooks can only be called inside of the body of a ' +
'function component.',
);
});

// @gate enableUseHook
test('unwraps thenable that fulfills synchronously without suspending', async () => {
function App() {
Expand Down

0 comments on commit 4a2d86b

Please sign in to comment.