@samitouri / QOS-React / commits / 4a2d86bddb

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

@@ -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

@@ -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() {