Fix: Update while suspended fails to interrupt

packages/react-reconciler/src/ReactFiberHooks.js

@@ -711,6 +711,9 @@ function renderWithHooksAgain<Props, SecondArg>(
   //
   // Keep rendering in a loop for as long as render phase updates continue to
   // be scheduled. Use a counter to prevent infinite loops.
+
+  currentlyRenderingFiber = workInProgress;
+
   let numberOfReRenders: number = 0;
   let children;
   do {
@@ -826,13 +829,14 @@ export function resetHooksAfterThrow(): void {
   //
   // It should only reset things like the current dispatcher, to prevent hooks
   // from being called outside of a component.
+  currentlyRenderingFiber = (null: any);
 
   // 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 {
+export function resetHooksOnUnwind(workInProgress: Fiber): 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
@@ -842,7 +846,7 @@ export function resetHooksOnUnwind(): void {
     // Only reset the updates from the queue if it has a clone. If it does
     // not have a clone, that means it wasn't processed, and the updates were
     // scheduled before we entered the render phase.
-    let hook: Hook | null = currentlyRenderingFiber.memoizedState;
+    let hook: Hook | null = workInProgress.memoizedState;
     while (hook !== null) {
       const queue = hook.queue;
       if (queue !== null) {

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -1504,7 +1504,7 @@ function resetWorkInProgressStack() {
   } else {
     // Work-in-progress is in suspended state. Reset the work loop and unwind
     // both the suspended fiber and all its parents.
-    resetSuspendedWorkLoopOnUnwind();
+    resetSuspendedWorkLoopOnUnwind(workInProgress);
     interruptedWork = workInProgress;
   }
   while (interruptedWork !== null) {
@@ -1563,10 +1563,10 @@ function prepareFreshStack(root: FiberRoot, lanes: Lanes): Fiber {
   return rootWorkInProgress;
 }
 
-function resetSuspendedWorkLoopOnUnwind() {
+function resetSuspendedWorkLoopOnUnwind(fiber: Fiber) {
   // Reset module-level state that was set during the render phase.
   resetContextDependencies();
-  resetHooksOnUnwind();
+  resetHooksOnUnwind(fiber);
   resetChildReconcilerOnUnwind();
 }
 
@@ -2337,7 +2337,7 @@ function replaySuspendedUnitOfWork(unitOfWork: Fiber): void {
       // is to reuse uncached promises, but we happen to know that the only
       // promises that a host component might suspend on are definitely cached
       // because they are controlled by us. So don't bother.
-      resetHooksOnUnwind();
+      resetHooksOnUnwind(unitOfWork);
       // Fallthrough to the next branch.
     }
     default: {
@@ -2383,7 +2383,7 @@ function throwAndUnwindWorkLoop(unitOfWork: Fiber, thrownValue: mixed) {
   //
   // Return to the normal work loop. This will unwind the stack, and potentially
   // result in showing a fallback.
-  resetSuspendedWorkLoopOnUnwind();
+  resetSuspendedWorkLoopOnUnwind(unitOfWork);
 
   const returnFiber = unitOfWork.return;
   if (returnFiber === null || workInProgressRoot === null) {
@@ -3744,7 +3744,7 @@ if (__DEV__ && replayFailedUnitOfWorkWithInvokeGuardedCallback) {
       // same fiber again.
 
       // Unwind the failed stack frame
-      resetSuspendedWorkLoopOnUnwind();
+      resetSuspendedWorkLoopOnUnwind(unitOfWork);
       unwindInterruptedWork(current, unitOfWork, workInProgressRootRenderLanes);
 
       // Restore the original properties of the fiber.

packages/react-reconciler/src/__tests__/ReactUse-test.js

@@ -1580,4 +1580,40 @@ describe('ReactUse', () => {
       </>,
     );
   });
+
+  test('regression test: updates while component is suspended should not be mistaken for render phase updates', async () => {
+    const getCachedAsyncText = cache(getAsyncText);
+
+    let setState;
+    function App() {
+      const [state, _setState] = useState('A');
+      setState = _setState;
+      return <Text text={use(getCachedAsyncText(state))} />;
+    }
+
+    // Initial render
+    const root = ReactNoop.createRoot();
+    await act(() => root.render(<App />));
+    assertLog(['Async text requested [A]']);
+    expect(root).toMatchRenderedOutput(null);
+    await act(() => resolveTextRequests('A'));
+    assertLog(['A']);
+    expect(root).toMatchRenderedOutput('A');
+
+    // Update to B. This will suspend.
+    await act(() => startTransition(() => setState('B')));
+    assertLog(['Async text requested [B]']);
+    expect(root).toMatchRenderedOutput('A');
+
+    // While B is suspended, update to C. This should immediately interrupt
+    // the render for B. In the regression, this update was mistakenly treated
+    // as a render phase update.
+    ReactNoop.flushSync(() => setState('C'));
+    assertLog(['Async text requested [C]']);
+
+    // Finish rendering.
+    await act(() => resolveTextRequests('C'));
+    assertLog(['C']);
+    expect(root).toMatchRenderedOutput('C');
+  });
 });