Bugfix: Render phase update causes remaining updates in same component to be dropped

packages/react-reconciler/src/ReactFiberHooks.js

@@ -230,6 +230,11 @@ let workInProgressHook: Hook | null = null;
 // finished evaluating this component. This is an optimization so we know
 // whether we need to clear render phase updates after a throw.
 let didScheduleRenderPhaseUpdate: boolean = false;
+// Where an update was scheduled only during the current render pass. This
+// gets reset after each attempt.
+// TODO: Maybe there's some way to consolidate this with
+// `didScheduleRenderPhaseUpdate`. Or with `numberOfReRenders`.
+let didScheduleRenderPhaseUpdateDuringThisPass: boolean = false;
 
 const RE_RENDER_LIMIT = 25;
 
@@ -455,13 +460,12 @@ export function renderWithHooks<Props, SecondArg>(
   let children = Component(props, secondArg);
 
   // Check if there was a render phase update
-  if (workInProgress.expirationTime === renderExpirationTime) {
+  if (didScheduleRenderPhaseUpdateDuringThisPass) {
     // Keep rendering in a loop for as long as render phase updates continue to
     // be scheduled. Use a counter to prevent infinite loops.
     let numberOfReRenders: number = 0;
     do {
-      workInProgress.expirationTime = NoWork;
-
+      didScheduleRenderPhaseUpdateDuringThisPass = false;
       invariant(
         numberOfReRenders < RE_RENDER_LIMIT,
         'Too many re-renders. React limits the number of renders to prevent ' +
@@ -491,7 +495,7 @@ export function renderWithHooks<Props, SecondArg>(
         : HooksDispatcherOnRerender;
 
       children = Component(props, secondArg);
-    } while (workInProgress.expirationTime === renderExpirationTime);
+    } while (didScheduleRenderPhaseUpdateDuringThisPass);
   }
 
   // We can assume the previous dispatcher is always this one, since we set it
@@ -564,6 +568,7 @@ export function resetHooksAfterThrow(): void {
       }
       hook = hook.next;
     }
+    didScheduleRenderPhaseUpdate = false;
   }
 
   renderExpirationTime = NoWork;
@@ -581,7 +586,7 @@ export function resetHooksAfterThrow(): void {
     isUpdatingOpaqueValueInRenderPhase = false;
   }
 
-  didScheduleRenderPhaseUpdate = false;
+  didScheduleRenderPhaseUpdateDuringThisPass = false;
 }
 
 function mountWorkInProgressHook(): Hook {
@@ -1726,9 +1731,8 @@ function dispatchAction<S, A>(
     // This is a render phase update. Stash it in a lazily-created map of
     // queue -> linked list of updates. After this render pass, we'll restart
     // and apply the stashed updates on top of the work-in-progress hook.
-    didScheduleRenderPhaseUpdate = true;
+    didScheduleRenderPhaseUpdateDuringThisPass = didScheduleRenderPhaseUpdate = true;
     update.expirationTime = renderExpirationTime;
-    currentlyRenderingFiber.expirationTime = renderExpirationTime;
   } else {
     if (
       fiber.expirationTime === NoWork &&

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

@@ -750,6 +750,41 @@ describe('ReactHooksWithNoopRenderer', () => {
       expect(root).toMatchRenderedOutput(<span prop="B:0" />);
     });
 
+    it('regression: render phase updates cause lower pri work to be dropped', async () => {
+      let setRow;
+      function ScrollView() {
+        const [row, _setRow] = useState(10);
+        setRow = _setRow;
+
+        const [scrollDirection, setScrollDirection] = useState('Up');
+        const [prevRow, setPrevRow] = useState(null);
+
+        if (prevRow !== row) {
+          setScrollDirection(prevRow !== null && row > prevRow ? 'Down' : 'Up');
+          setPrevRow(row);
+        }
+
+        return <Text text={scrollDirection} />;
+      }
+
+      const root = ReactNoop.createRoot();
+
+      await act(async () => {
+        root.render(<ScrollView row={10} />);
+      });
+      expect(Scheduler).toHaveYielded(['Up']);
+      expect(root).toMatchRenderedOutput(<span prop="Up" />);
+
+      await act(async () => {
+        ReactNoop.discreteUpdates(() => {
+          setRow(5);
+        });
+        setRow(20);
+      });
+      expect(Scheduler).toHaveYielded(['Up', 'Down']);
+      expect(root).toMatchRenderedOutput(<span prop="Down" />);
+    });
+
     // TODO: This should probably warn
     it.experimental('calling startTransition inside render phase', async () => {
       let startTransition;

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

@@ -3758,4 +3758,119 @@ describe('ReactSuspenseWithNoopRenderer', () => {
       </>,
     );
   });
+
+  // Regression: https://github.com/facebook/react/issues/18486
+  it.experimental(
+    'does not get stuck in pending state with render phase updates',
+    async () => {
+      let setTextWithTransition;
+
+      function App() {
+        const [startTransition, isPending] = React.useTransition({
+          timeoutMs: 30000,
+        });
+        const [text, setText] = React.useState('');
+        const [mirror, setMirror] = React.useState('');
+
+        if (text !== mirror) {
+          // Render phase update was needed to repro the bug.
+          setMirror(text);
+        }
+
+        setTextWithTransition = value => {
+          startTransition(() => {
+            setText(value);
+          });
+        };
+
+        return (
+          <>
+            {isPending ? <Text text="Pending..." /> : null}
+            {text !== '' ? <AsyncText text={text} /> : <Text text={text} />}
+          </>
+        );
+      }
+
+      function Root() {
+        return (
+          <Suspense fallback={<Text text="Loading..." />}>
+            <App />
+          </Suspense>
+        );
+      }
+
+      const root = ReactNoop.createRoot();
+      await ReactNoop.act(async () => {
+        root.render(<Root />);
+      });
+      expect(Scheduler).toHaveYielded(['']);
+      expect(root).toMatchRenderedOutput(<span prop="" />);
+
+      // Update to "a". That will suspend.
+      await ReactNoop.act(async () => {
+        setTextWithTransition('a');
+        // Let it expire. This is important for the repro.
+        Scheduler.unstable_advanceTime(1000);
+        expect(Scheduler).toFlushAndYield([
+          'Pending...',
+          '',
+          'Suspend! [a]',
+          'Loading...',
+        ]);
+      });
+      expect(Scheduler).toHaveYielded([]);
+      expect(root).toMatchRenderedOutput(
+        <>
+          <span prop="Pending..." />
+          <span prop="" />
+        </>,
+      );
+
+      // Update to "b". That will suspend, too.
+      await ReactNoop.act(async () => {
+        setTextWithTransition('b');
+        expect(Scheduler).toFlushAndYield([
+          // Neither is resolved yet.
+          'Pending...',
+          'Suspend! [a]',
+          'Loading...',
+          'Suspend! [b]',
+          'Loading...',
+        ]);
+      });
+      expect(Scheduler).toHaveYielded([]);
+      expect(root).toMatchRenderedOutput(
+        <>
+          <span prop="Pending..." />
+          <span prop="" />
+        </>,
+      );
+
+      // Resolve "a". But "b" is still pending.
+      await ReactNoop.act(async () => {
+        await resolveText('a');
+      });
+      expect(Scheduler).toHaveYielded([
+        'Promise resolved [a]',
+        'Pending...',
+        'a',
+        'Suspend! [b]',
+        'Loading...',
+      ]);
+      expect(root).toMatchRenderedOutput(
+        <>
+          <span prop="Pending..." />
+          <span prop="a" />
+        </>,
+      );
+
+      // Resolve "b". This should remove the pending state.
+      await ReactNoop.act(async () => {
+        await resolveText('b');
+      });
+      expect(Scheduler).toHaveYielded(['Promise resolved [b]', 'b']);
+      // The bug was that the pending state got stuck forever.
+      expect(root).toMatchRenderedOutput(<span prop="b" />);
+    },
+  );
 });