Don't reset work loop until stack is unwound

acdlite · acdlite · commit 4a2d86bddbce · 2022-11-17T14:48:42.000-05:00
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.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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() {
