From 4a2d86bddbcef3e64bc404302cdbd9638af8801b Mon Sep 17 00:00:00 2001 From: Andrew Clark Date: Wed, 2 Nov 2022 16:19:22 -0400 Subject: [PATCH] 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. --- .../src/ReactFiberHooks.new.js | 9 ++++ .../src/ReactFiberHooks.old.js | 9 ++++ .../src/ReactFiberWorkLoop.new.js | 46 ++++++++++++++----- .../src/ReactFiberWorkLoop.old.js | 46 ++++++++++++++----- .../src/__tests__/ReactThenable-test.js | 31 +++++++++++++ 5 files changed, 119 insertions(+), 22 deletions(-) diff --git a/packages/react-reconciler/src/ReactFiberHooks.new.js b/packages/react-reconciler/src/ReactFiberHooks.new.js index 7538bb9b3c217..b07667527277c 100644 --- a/packages/react-reconciler/src/ReactFiberHooks.new.js +++ b/packages/react-reconciler/src/ReactFiberHooks.new.js @@ -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 diff --git a/packages/react-reconciler/src/ReactFiberHooks.old.js b/packages/react-reconciler/src/ReactFiberHooks.old.js index e98d684a78e6f..73479626e1fbc 100644 --- a/packages/react-reconciler/src/ReactFiberHooks.old.js +++ b/packages/react-reconciler/src/ReactFiberHooks.old.js @@ -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 diff --git a/packages/react-reconciler/src/ReactFiberWorkLoop.new.js b/packages/react-reconciler/src/ReactFiberWorkLoop.new.js index 9bc53e4f30e78..7e233448b5530 100644 --- a/packages/react-reconciler/src/ReactFiberWorkLoop.new.js +++ b/packages/react-reconciler/src/ReactFiberWorkLoop.new.js @@ -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'; @@ -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( @@ -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) { @@ -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); @@ -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) { @@ -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. diff --git a/packages/react-reconciler/src/ReactFiberWorkLoop.old.js b/packages/react-reconciler/src/ReactFiberWorkLoop.old.js index 12293f0210e0f..f1cc5a6600c1b 100644 --- a/packages/react-reconciler/src/ReactFiberWorkLoop.old.js +++ b/packages/react-reconciler/src/ReactFiberWorkLoop.old.js @@ -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'; @@ -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( @@ -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) { @@ -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); @@ -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) { @@ -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. diff --git a/packages/react-reconciler/src/__tests__/ReactThenable-test.js b/packages/react-reconciler/src/__tests__/ReactThenable-test.js index 06a11a2101225..600c6e3e5ba9e 100644 --- a/packages/react-reconciler/src/__tests__/ReactThenable-test.js +++ b/packages/react-reconciler/src/__tests__/ReactThenable-test.js @@ -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 ; + } + + const root = ReactNoop.createRoot(); + await act(async () => { + root.render(} />); + }); + + await act(async () => { + startTransition(() => { + root.render( + }> + + , + ); + }); + }); + 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() {