Discard render phase updates if component throws

When aborting a render, we also need to throw out render phase updates.
Remove the updates from the queues so they do not persist to the next
render. We already did a single pass through the whole list of hooks, so
we know that any pending updates must have been dispatched during the
render phase. The ones that were dispatched before we started rendering
were already transferred to the current hook's queue.

Author: acdlite

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -136,7 +136,7 @@ import {
   calculateChangedBits,
   scheduleWorkOnParentPath,
 } from './ReactFiberNewContext';
-import {resetHooks, renderWithHooks, bailoutHooks} from './ReactFiberHooks';
+import {renderWithHooks, bailoutHooks} from './ReactFiberHooks';
 import {stopProfilerTimerIfRunning} from './ReactProfilerTimer';
 import {
   getMaskedContext,
@@ -1319,7 +1319,8 @@ function mountIndeterminateComponent(
     workInProgress.tag = ClassComponent;
 
     // Throw out any hooks that were used.
-    resetHooks();
+    workInProgress.memoizedState = null;
+    workInProgress.updateQueue = null;
 
     // Push context providers early to prevent context stack mismatches.
     // During mounting we don't know the child context yet as the instance doesn't exist.

packages/react-reconciler/src/ReactFiberHooks.js

@@ -504,14 +504,31 @@ export function bailoutHooks(
   }
 }
 
-export function resetHooks(): void {
+export function resetHooksAfterThrow(): void {
   // 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-entrancy.
   ReactCurrentDispatcher.current = ContextOnlyDispatcher;
 
-  // This is used to reset the state of this module when a component throws.
-  // It's also called inside mountIndeterminateComponent if we determine the
-  // component is a module-style component.
+  if (didScheduleRenderPhaseUpdate) {
+    const current = (currentlyRenderingFiber: any).alternate;
+    if (current !== null) {
+      // There were render phase updates. These are only valid for this render
+      // pass, which we are now aborting. Remove the updates from the queues so
+      // they do not persist to the next render. We already did a single pass
+      // through the whole list of hooks, so we know that any pending updates
+      // must have been dispatched during the render phase. The ones that were
+      // dispatched before we started rendering were already transferred to the
+      // current hook's queue.
+      let hook: Hook | null = current.memoizedState;
+      while (hook !== null) {
+        const queue = hook.queue;
+        if (queue !== null) {
+          queue.pending = null;
+        }
+        hook = hook.next;
+      }
+    }
+  }
 
   renderExpirationTime = NoWork;
   currentlyRenderingFiber = (null: any);

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -140,7 +140,7 @@ import {
 } from './ReactFiberCommitWork';
 import {enqueueUpdate} from './ReactUpdateQueue';
 import {resetContextDependencies} from './ReactFiberNewContext';
-import {resetHooks, ContextOnlyDispatcher} from './ReactFiberHooks';
+import {resetHooksAfterThrow, ContextOnlyDispatcher} from './ReactFiberHooks';
 import {createCapturedValue} from './ReactCapturedValue';
 
 import {
@@ -1281,7 +1281,7 @@ function handleError(root, thrownValue) {
     try {
       // Reset module-level state that was set during the render phase.
       resetContextDependencies();
-      resetHooks();
+      resetHooksAfterThrow();
       resetCurrentDebugFiberInDEV();
 
       if (workInProgress === null || workInProgress.return === null) {
@@ -2635,7 +2635,7 @@ if (__DEV__ && replayFailedUnitOfWorkWithInvokeGuardedCallback) {
       // Keep this code in sync with handleError; any changes here must have
       // corresponding changes there.
       resetContextDependencies();
-      resetHooks();
+      resetHooksAfterThrow();
       // Don't reset current debug fiber, since we're about to work on the
       // same fiber again.
 

packages/react-reconciler/src/__tests__/ReactHooksWithNoopRenderer-test.internal.js

@@ -542,6 +542,50 @@ describe('ReactHooksWithNoopRenderer', () => {
       ]);
       expect(ReactNoop.getChildren()).toEqual([span(22)]);
     });
+
+    it('discards render phase updates if something suspends', () => {
+      const thenable = {then() {}};
+      function Foo({signal}) {
+        return (
+          <Suspense fallback="Loading...">
+            <Bar signal={signal} />
+          </Suspense>
+        );
+      }
+
+      function Bar({signal: newSignal}) {
+        let [counter, setCounter] = useState(0);
+        let [signal, setSignal] = useState(true);
+
+        // Increment a counter every time the signal changes
+        if (signal !== newSignal) {
+          setCounter(c => c + 1);
+          setSignal(newSignal);
+          if (counter === 0) {
+            // We're suspending during a render that includes render phase
+            // updates. Those updates should not persist to the next render.
+            Scheduler.unstable_yieldValue('Suspend!');
+            throw thenable;
+          }
+        }
+
+        return <Text text={counter} />;
+      }
+
+      const root = ReactNoop.createRoot();
+      root.render(<Foo signal={true} />);
+
+      expect(Scheduler).toFlushAndYield([0]);
+      expect(root).toMatchRenderedOutput(<span prop={0} />);
+
+      root.render(<Foo signal={false} />);
+      expect(Scheduler).toFlushAndYield(['Suspend!']);
+      expect(root).toMatchRenderedOutput(<span prop={0} />);
+
+      // Rendering again should suspend again.
+      root.render(<Foo signal={false} />);
+      expect(Scheduler).toFlushAndYield(['Suspend!']);
+    });
   });
 
   describe('useReducer', () => {