Bugfix: Expired partial tree infinite loops (#17949)

* Bugfix: Expiring a partially completed tree (#17926)

* Failing test: Expiring a partially completed tree

We should not throw out a partially completed tree if it expires in the
middle of rendering. We should finish the rest of the tree without
yielding, then finish any remaining expired levels in a single batch.

* Check if there's a partial tree before restarting

If a partial render expires, we should stay in the concurrent path
(performConcurrentWorkOnRoot); we'll stop yielding, but the rest of the
behavior remains the same.

We will only revert to the sync path (performSyncWorkOnRoot) when
starting on a new level.

This approach prevents partially completed concurrent work from
being discarded.

* New test: retry after error during expired render

* Regression: Expired partial tree infinite loops

Adds regression tests that reproduce a scenario where a partially
completed tree expired then fell into an infinite loop.

The code change that exposed this bug made the assumption that if you
call Scheduler's `shouldYield` from inside an expired task, Scheduler
will always return `false`. But in fact, Scheduler sometimes returns
`true` in that scenario, which is a bug.

The reason it worked before is that once a task timed out, React would
always switch to a synchronous work loop without checking `shouldYield`.

My rationale for relying on `shouldYield` was to unify the code paths
between a partially concurrent render (i.e. expires midway through) and
a fully concurrent render, as opposed to a render that was synchronous
the whole time. However, this bug indicates that we need a stronger
guarantee within React for when tasks expire, given that the failure
case is so catastrophic. Instead of relying on the result of a dynamic
method call, we should use control flow to guarantee that the work is
synchronously executed.

(We should also fix the Scheduler bug so that `shouldYield` always
returns false inside an expired task, but I'll address that separately.)

* Always switch to sync work loop when task expires

Refactors the `didTimeout` check so that it always switches to the
synchronous work loop, like it did before the regression.

This breaks the error handling behavior that I added in 5f7361f (an
error during a partially concurrent render should retry once,
synchronously). I'll fix this next. I need to change that behavior,
anyway, to support retries that occur as a result of `flushSync`.

* Retry once after error even for sync renders

Except in legacy mode.

This is to support the `useOpaqueReference` hook, which uses an error
to trigger a retry at lower priority.

* Move partial tree check to performSyncWorkOnRoot

* Factor out render phase

Splits the work loop and its surrounding enter/exit code into their own
functions. Now we can do perform multiple render phase passes within a
single call to performConcurrentWorkOnRoot or performSyncWorkOnRoot.
This lets us get rid of the `didError` field.

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -90,6 +90,7 @@ import {
   SimpleMemoComponent,
   Block,
 } from 'shared/ReactWorkTags';
+import {LegacyRoot} from 'shared/ReactRootTags';
 import {
   NoEffect,
   PerformedWork,
@@ -643,6 +644,7 @@ function performConcurrentWorkOnRoot(root, didTimeout) {
   // event time. The next update will compute a new event time.
   currentEventTime = NoWork;
 
+  // Check if the render expired.
   if (didTimeout) {
     // The render task took too long to complete. Mark the current time as
     // expired to synchronously render all expired work in a single batch.
@@ -655,84 +657,52 @@ function performConcurrentWorkOnRoot(root, didTimeout) {
 
   // Determine the next expiration time to work on, using the fields stored
   // on the root.
-  const expirationTime = getNextRootExpirationTimeToWorkOn(root);
-  if (expirationTime !== NoWork) {
-    const originalCallbackNode = root.callbackNode;
-    invariant(
-      (executionContext & (RenderContext | CommitContext)) === NoContext,
-      'Should not already be working.',
-    );
+  let expirationTime = getNextRootExpirationTimeToWorkOn(root);
+  if (expirationTime === NoWork) {
+    return null;
+  }
+  const originalCallbackNode = root.callbackNode;
+  invariant(
+    (executionContext & (RenderContext | CommitContext)) === NoContext,
+    'Should not already be working.',
+  );
 
-    flushPassiveEffects();
+  flushPassiveEffects();
 
-    // If the root or expiration time have changed, throw out the existing stack
-    // and prepare a fresh one. Otherwise we'll continue where we left off.
-    if (
-      root !== workInProgressRoot ||
-      expirationTime !== renderExpirationTime
-    ) {
+  let exitStatus = renderRootConcurrent(root, expirationTime);
+
+  if (exitStatus !== RootIncomplete) {
+    if (exitStatus === RootErrored) {
+      // If something threw an error, try rendering one more time. We'll
+      // render synchronously to block concurrent data mutations, and we'll
+      // render at Idle (or lower) so that all pending updates are included.
+      // If it still fails after the second attempt, we'll give up and commit
+      // the resulting tree.
+      expirationTime = expirationTime > Idle ? Idle : expirationTime;
+      exitStatus = renderRootSync(root, expirationTime);
+    }
+
+    if (exitStatus === RootFatalErrored) {
+      const fatalError = workInProgressRootFatalError;
       prepareFreshStack(root, expirationTime);
-      startWorkOnPendingInteractions(root, expirationTime);
-    }
-
-    // If we have a work-in-progress fiber, it means there's still work to do
-    // in this root.
-    if (workInProgress !== null) {
-      const prevExecutionContext = executionContext;
-      executionContext |= RenderContext;
-      const prevDispatcher = pushDispatcher(root);
-      const prevInteractions = pushInteractions(root);
-      startWorkLoopTimer(workInProgress);
-      do {
-        try {
-          workLoopConcurrent();
-          break;
-        } catch (thrownValue) {
-          handleError(root, thrownValue);
-        }
-      } while (true);
-      resetContextDependencies();
-      executionContext = prevExecutionContext;
-      popDispatcher(prevDispatcher);
-      if (enableSchedulerTracing) {
-        popInteractions(((prevInteractions: any): Set<Interaction>));
-      }
-
-      if (workInProgressRootExitStatus === RootFatalErrored) {
-        const fatalError = workInProgressRootFatalError;
-        stopInterruptedWorkLoopTimer();
-        prepareFreshStack(root, expirationTime);
-        markRootSuspendedAtTime(root, expirationTime);
-        ensureRootIsScheduled(root);
-        throw fatalError;
-      }
-
-      if (workInProgress !== null) {
-        // There's still work left over. Exit without committing.
-        stopInterruptedWorkLoopTimer();
-      } else {
-        // We now have a consistent tree. The next step is either to commit it,
-        // or, if something suspended, wait to commit it after a timeout.
-        stopFinishedWorkLoopTimer();
-
-        const finishedWork: Fiber = ((root.finishedWork =
-          root.current.alternate): any);
-        root.finishedExpirationTime = expirationTime;
-        finishConcurrentRender(
-          root,
-          finishedWork,
-          workInProgressRootExitStatus,
-          expirationTime,
-        );
-      }
-
+      markRootSuspendedAtTime(root, expirationTime);
       ensureRootIsScheduled(root);
-      if (root.callbackNode === originalCallbackNode) {
-        // The task node scheduled for this root is the same one that's
-        // currently executed. Need to return a continuation.
-        return performConcurrentWorkOnRoot.bind(null, root);
-      }
+      throw fatalError;
     }
+
+    // We now have a consistent tree. The next step is either to commit it,
+    // or, if something suspended, wait to commit it after a timeout.
+    const finishedWork: Fiber = ((root.finishedWork =
+      root.current.alternate): any);
+    root.finishedExpirationTime = expirationTime;
+    finishConcurrentRender(root, finishedWork, exitStatus, expirationTime);
+  }
+
+  ensureRootIsScheduled(root);
+  if (root.callbackNode === originalCallbackNode) {
+    // The task node scheduled for this root is the same one that's
+    // currently executed. Need to return a continuation.
+    return performConcurrentWorkOnRoot.bind(null, root);
   }
   return null;
 }
@@ -743,9 +713,6 @@ function finishConcurrentRender(
   exitStatus,
   expirationTime,
 ) {
-  // Set this to null to indicate there's no in-progress render.
-  workInProgressRoot = null;
-
   switch (exitStatus) {
     case RootIncomplete:
     case RootFatalErrored: {
@@ -755,19 +722,9 @@ function finishConcurrentRender(
     // statement, but eslint doesn't know about invariant, so it complains
     // if I do. eslint-disable-next-line no-fallthrough
     case RootErrored: {
-      // If this was an async render, the error may have happened due to
-      // a mutation in a concurrent event. Try rendering one more time,
-      // synchronously, to see if the error goes away. If there are
-      // lower priority updates, let's include those, too, in case they
-      // fix the inconsistency. Render at Idle to include all updates.
-      // If it was Idle or Never or some not-yet-invented time, render
-      // at that time.
-      markRootExpiredAtTime(
-        root,
-        expirationTime > Idle ? Idle : expirationTime,
-      );
-      // We assume that this second render pass will be synchronous
-      // and therefore not hit this path again.
+      // We should have already attempted to retry this tree. If we reached
+      // this point, it errored again. Commit it.
+      commitRoot(root);
       break;
     }
     case RootSuspended: {
@@ -983,9 +940,6 @@ function finishConcurrentRender(
 // This is the entry point for synchronous tasks that don't go
 // through Scheduler
 function performSyncWorkOnRoot(root) {
-  // Check if there's expired work on this root. Otherwise, render at Sync.
-  const lastExpiredTime = root.lastExpiredTime;
-  const expirationTime = lastExpiredTime !== NoWork ? lastExpiredTime : Sync;
   invariant(
     (executionContext & (RenderContext | CommitContext)) === NoContext,
     'Should not already be working.',
@@ -993,76 +947,63 @@ function performSyncWorkOnRoot(root) {
 
   flushPassiveEffects();
 
-  // If the root or expiration time have changed, throw out the existing stack
-  // and prepare a fresh one. Otherwise we'll continue where we left off.
-  if (root !== workInProgressRoot || expirationTime !== renderExpirationTime) {
-    prepareFreshStack(root, expirationTime);
-    startWorkOnPendingInteractions(root, expirationTime);
-  }
+  const lastExpiredTime = root.lastExpiredTime;
 
-  // If we have a work-in-progress fiber, it means there's still work to do
-  // in this root.
-  if (workInProgress !== null) {
-    const prevExecutionContext = executionContext;
-    executionContext |= RenderContext;
-    const prevDispatcher = pushDispatcher(root);
-    const prevInteractions = pushInteractions(root);
-    startWorkLoopTimer(workInProgress);
-
-    do {
-      try {
-        workLoopSync();
-        break;
-      } catch (thrownValue) {
-        handleError(root, thrownValue);
-      }
-    } while (true);
-    resetContextDependencies();
-    executionContext = prevExecutionContext;
-    popDispatcher(prevDispatcher);
-    if (enableSchedulerTracing) {
-      popInteractions(((prevInteractions: any): Set<Interaction>));
-    }
-
-    if (workInProgressRootExitStatus === RootFatalErrored) {
-      const fatalError = workInProgressRootFatalError;
-      stopInterruptedWorkLoopTimer();
-      prepareFreshStack(root, expirationTime);
-      markRootSuspendedAtTime(root, expirationTime);
-      ensureRootIsScheduled(root);
-      throw fatalError;
-    }
-
-    if (workInProgress !== null) {
-      // This is a sync render, so we should have finished the whole tree.
-      invariant(
-        false,
-        'Cannot commit an incomplete root. This error is likely caused by a ' +
-          'bug in React. Please file an issue.',
-      );
+  let expirationTime;
+  if (lastExpiredTime !== NoWork) {
+    // There's expired work on this root. Check if we have a partial tree
+    // that we can reuse.
+    if (
+      root === workInProgressRoot &&
+      renderExpirationTime >= lastExpiredTime
+    ) {
+      // There's a partial tree with equal or greater than priority than the
+      // expired level. Finish rendering it before rendering the rest of the
+      // expired work.
+      expirationTime = renderExpirationTime;
     } else {
-      // We now have a consistent tree. Because this is a sync render, we
-      // will commit it even if something suspended.
-      stopFinishedWorkLoopTimer();
-      root.finishedWork = (root.current.alternate: any);
-      root.finishedExpirationTime = expirationTime;
-      finishSyncRender(root);
+      // Start a fresh tree.
+      expirationTime = lastExpiredTime;
     }
-
-    // Before exiting, make sure there's a callback scheduled for the next
-    // pending level.
-    ensureRootIsScheduled(root);
+  } else {
+    // There's no expired work. This must be a new, synchronous render.
+    expirationTime = Sync;
   }
 
+  let exitStatus = renderRootSync(root, expirationTime);
+
+  if (root.tag !== LegacyRoot && exitStatus === RootErrored) {
+    // If something threw an error, try rendering one more time. We'll
+    // render synchronously to block concurrent data mutations, and we'll
+    // render at Idle (or lower) so that all pending updates are included.
+    // If it still fails after the second attempt, we'll give up and commit
+    // the resulting tree.
+    expirationTime = expirationTime > Idle ? Idle : expirationTime;
+    exitStatus = renderRootSync(root, expirationTime);
+  }
+
+  if (exitStatus === RootFatalErrored) {
+    const fatalError = workInProgressRootFatalError;
+    prepareFreshStack(root, expirationTime);
+    markRootSuspendedAtTime(root, expirationTime);
+    ensureRootIsScheduled(root);
+    throw fatalError;
+  }
+
+  // We now have a consistent tree. Because this is a sync render, we
+  // will commit it even if something suspended.
+  root.finishedWork = (root.current.alternate: any);
+  root.finishedExpirationTime = expirationTime;
+
+  commitRoot(root);
+
+  // Before exiting, make sure there's a callback scheduled for the next
+  // pending level.
+  ensureRootIsScheduled(root);
+
   return null;
 }
 
-function finishSyncRender(root) {
-  // Set this to null to indicate there's no in-progress render.
-  workInProgressRoot = null;
-  commitRoot(root);
-}
-
 export function flushRoot(root: FiberRoot, expirationTime: ExpirationTime) {
   markRootExpiredAtTime(root, expirationTime);
   ensureRootIsScheduled(root);
@@ -1449,6 +1390,53 @@ function inferTimeFromExpirationTimeWithSuspenseConfig(
   );
 }
 
+function renderRootSync(root, expirationTime) {
+  const prevExecutionContext = executionContext;
+  executionContext |= RenderContext;
+  const prevDispatcher = pushDispatcher(root);
+
+  // If the root or expiration time have changed, throw out the existing stack
+  // and prepare a fresh one. Otherwise we'll continue where we left off.
+  if (root !== workInProgressRoot || expirationTime !== renderExpirationTime) {
+    prepareFreshStack(root, expirationTime);
+    startWorkOnPendingInteractions(root, expirationTime);
+  }
+
+  const prevInteractions = pushInteractions(root);
+  startWorkLoopTimer(workInProgress);
+  do {
+    try {
+      workLoopSync();
+      break;
+    } catch (thrownValue) {
+      handleError(root, thrownValue);
+    }
+  } while (true);
+  resetContextDependencies();
+  if (enableSchedulerTracing) {
+    popInteractions(((prevInteractions: any): Set<Interaction>));
+  }
+
+  executionContext = prevExecutionContext;
+  popDispatcher(prevDispatcher);
+
+  if (workInProgress !== null) {
+    // This is a sync render, so we should have finished the whole tree.
+    invariant(
+      false,
+      'Cannot commit an incomplete root. This error is likely caused by a ' +
+        'bug in React. Please file an issue.',
+    );
+  }
+
+  stopFinishedWorkLoopTimer();
+
+  // Set this to null to indicate there's no in-progress render.
+  workInProgressRoot = null;
+
+  return workInProgressRootExitStatus;
+}
+
 // The work loop is an extremely hot path. Tell Closure not to inline it.
 /** @noinline */
 function workLoopSync() {
@@ -1458,6 +1446,53 @@ function workLoopSync() {
   }
 }
 
+function renderRootConcurrent(root, expirationTime) {
+  const prevExecutionContext = executionContext;
+  executionContext |= RenderContext;
+  const prevDispatcher = pushDispatcher(root);
+
+  // If the root or expiration time have changed, throw out the existing stack
+  // and prepare a fresh one. Otherwise we'll continue where we left off.
+  if (root !== workInProgressRoot || expirationTime !== renderExpirationTime) {
+    prepareFreshStack(root, expirationTime);
+    startWorkOnPendingInteractions(root, expirationTime);
+  }
+
+  const prevInteractions = pushInteractions(root);
+  startWorkLoopTimer(workInProgress);
+  do {
+    try {
+      workLoopConcurrent();
+      break;
+    } catch (thrownValue) {
+      handleError(root, thrownValue);
+    }
+  } while (true);
+  resetContextDependencies();
+  if (enableSchedulerTracing) {
+    popInteractions(((prevInteractions: any): Set<Interaction>));
+  }
+
+  popDispatcher(prevDispatcher);
+  executionContext = prevExecutionContext;
+
+  // Check if the tree has completed.
+  if (workInProgress !== null) {
+    // Still work remaining.
+    stopInterruptedWorkLoopTimer();
+    return RootIncomplete;
+  } else {
+    // Completed the tree.
+    stopFinishedWorkLoopTimer();
+
+    // Set this to null to indicate there's no in-progress render.
+    workInProgressRoot = null;
+
+    // Return the final exit status.
+    return workInProgressRootExitStatus;
+  }
+}
+
 /** @noinline */
 function workLoopConcurrent() {
   // Perform work until Scheduler asks us to yield

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

@@ -381,6 +381,56 @@ describe('ReactIncrementalErrorHandling', () => {
     expect(ReactNoop.getChildren()).toEqual([]);
   });
 
+  it('retries one more time if an error occurs during a render that expires midway through the tree', () => {
+    function Oops() {
+      Scheduler.unstable_yieldValue('Oops');
+      throw new Error('Oops');
+    }
+
+    function Text({text}) {
+      Scheduler.unstable_yieldValue(text);
+      return text;
+    }
+
+    function App() {
+      return (
+        <>
+          <Text text="A" />
+          <Text text="B" />
+          <Oops />
+          <Text text="C" />
+          <Text text="D" />
+        </>
+      );
+    }
+
+    ReactNoop.render(<App />);
+
+    // Render part of the tree
+    expect(Scheduler).toFlushAndYieldThrough(['A', 'B']);
+
+    // Expire the render midway through
+    Scheduler.unstable_advanceTime(10000);
+    expect(() => Scheduler.unstable_flushExpired()).toThrow('Oops');
+
+    expect(Scheduler).toHaveYielded([
+      // The render expired, but we shouldn't throw out the partial work.
+      // Finish the current level.
+      'Oops',
+      'C',
+      'D',
+
+      // Since the error occured during a partially concurrent render, we should
+      // retry one more time, synchonrously.
+      'A',
+      'B',
+      'Oops',
+      'C',
+      'D',
+    ]);
+    expect(ReactNoop.getChildren()).toEqual([]);
+  });
+
   it('calls componentDidCatch multiple times for multiple errors', () => {
     let id = 0;
     class BadMount extends React.Component {
@@ -539,7 +589,12 @@ describe('ReactIncrementalErrorHandling', () => {
     expect(ops).toEqual([
       'ErrorBoundary render success',
       'BrokenRender',
-      // React doesn't retry because we're already rendering synchronously.
+
+      // React retries one more time
+      'ErrorBoundary render success',
+      'BrokenRender',
+
+      // Errored again on retry. Now handle it.
       'ErrorBoundary componentDidCatch',
       'ErrorBoundary render error',
     ]);
@@ -583,7 +638,12 @@ describe('ReactIncrementalErrorHandling', () => {
     expect(ops).toEqual([
       'ErrorBoundary render success',
       'BrokenRender',
-      // React doesn't retry because we're already rendering synchronously.
+
+      // React retries one more time
+      'ErrorBoundary render success',
+      'BrokenRender',
+
+      // Errored again on retry. Now handle it.
       'ErrorBoundary componentDidCatch',
       'ErrorBoundary render error',
     ]);
@@ -702,7 +762,12 @@ describe('ReactIncrementalErrorHandling', () => {
     expect(ops).toEqual([
       'RethrowErrorBoundary render',
       'BrokenRender',
-      // React doesn't retry because we're already rendering synchronously.
+
+      // React retries one more time
+      'RethrowErrorBoundary render',
+      'BrokenRender',
+
+      // Errored again on retry. Now handle it.
       'RethrowErrorBoundary componentDidCatch',
     ]);
     expect(ReactNoop.getChildren()).toEqual([]);
@@ -741,7 +806,12 @@ describe('ReactIncrementalErrorHandling', () => {
     expect(ops).toEqual([
       'RethrowErrorBoundary render',
       'BrokenRender',
-      // React doesn't retry because we're already rendering synchronously.
+
+      // React retries one more time
+      'RethrowErrorBoundary render',
+      'BrokenRender',
+
+      // Errored again on retry. Now handle it.
       'RethrowErrorBoundary componentDidCatch',
     ]);
     expect(ReactNoop.getChildren()).toEqual([]);

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

@@ -189,8 +189,14 @@ describe('ReactIncrementalErrorLogging', () => {
       [
         'render: 0',
         __DEV__ && 'render: 0', // replay
+
         'render: 1',
         __DEV__ && 'render: 1', // replay
+
+        // Retry one more time before handling error
+        'render: 1',
+        __DEV__ && 'render: 1', // replay
+
         'componentWillUnmount: 0',
       ].filter(Boolean),
     );

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

@@ -655,6 +655,116 @@ describe('ReactIncrementalUpdates', () => {
     });
   });
 
+  it('does not throw out partially completed tree if it expires midway through', () => {
+    function Text({text}) {
+      Scheduler.unstable_yieldValue(text);
+      return text;
+    }
+
+    function App({step}) {
+      return (
+        <>
+          <Text text={`A${step}`} />
+          <Text text={`B${step}`} />
+          <Text text={`C${step}`} />
+        </>
+      );
+    }
+
+    function interrupt() {
+      ReactNoop.flushSync(() => {
+        ReactNoop.renderToRootWithID(null, 'other-root');
+      });
+    }
+
+    // First, as a sanity check, assert what happens when four low pri
+    // updates in separate batches are all flushed in the same callback
+    ReactNoop.act(() => {
+      ReactNoop.render(<App step={1} />);
+      Scheduler.unstable_advanceTime(1000);
+      expect(Scheduler).toFlushAndYieldThrough(['A1']);
+
+      interrupt();
+
+      ReactNoop.render(<App step={2} />);
+      Scheduler.unstable_advanceTime(1000);
+      expect(Scheduler).toFlushAndYieldThrough(['A1']);
+
+      interrupt();
+
+      ReactNoop.render(<App step={3} />);
+      Scheduler.unstable_advanceTime(1000);
+      expect(Scheduler).toFlushAndYieldThrough(['A1']);
+
+      interrupt();
+
+      ReactNoop.render(<App step={4} />);
+      Scheduler.unstable_advanceTime(1000);
+      expect(Scheduler).toFlushAndYieldThrough(['A1']);
+
+      // Each update flushes in a separate commit.
+      // Note: This isn't necessarily the ideal behavior. It might be better to
+      // batch all of these updates together. The fact that they don't is an
+      // implementation detail. The important part of this unit test is what
+      // happens when they expire, in which case they really should be batched to
+      // avoid blocking the main thread for a long time.
+      expect(Scheduler).toFlushAndYield([
+        // A1 already completed. Finish rendering the first level.
+        'B1',
+        'C1',
+        // The remaining two levels complete sequentially.
+        'A2',
+        'B2',
+        'C2',
+        'A3',
+        'B3',
+        'C3',
+        'A4',
+        'B4',
+        'C4',
+      ]);
+    });
+
+    ReactNoop.act(() => {
+      // Now do the same thing over again, but this time, expire all the updates
+      // instead of flushing them normally.
+      ReactNoop.render(<App step={1} />);
+      Scheduler.unstable_advanceTime(1000);
+      expect(Scheduler).toFlushAndYieldThrough(['A1']);
+
+      interrupt();
+
+      ReactNoop.render(<App step={2} />);
+      Scheduler.unstable_advanceTime(1000);
+      expect(Scheduler).toFlushAndYieldThrough(['A1']);
+
+      interrupt();
+
+      ReactNoop.render(<App step={3} />);
+      Scheduler.unstable_advanceTime(1000);
+      expect(Scheduler).toFlushAndYieldThrough(['A1']);
+
+      interrupt();
+
+      ReactNoop.render(<App step={4} />);
+      Scheduler.unstable_advanceTime(1000);
+      expect(Scheduler).toFlushAndYieldThrough(['A1']);
+
+      // Expire all the updates
+      ReactNoop.expire(10000);
+
+      expect(Scheduler).toFlushExpired([
+        // A1 already completed. Finish rendering the first level.
+        'B1',
+        'C1',
+        // Then render the remaining two levels in a single batch
+        'A4',
+        'B4',
+        'C4',
+      ]);
+    });
+  });
+
   it('when rebasing, does not exclude updates that were already committed, regardless of priority', async () => {
     const {useState, useLayoutEffect} = React;
 

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

@@ -408,3 +408,123 @@ describe('ReactSchedulerIntegration', () => {
     );
   });
 });
+
+describe(
+  'regression test: does not infinite loop if `shouldYield` returns ' +
+    'true after a partial tree expires',
+  () => {
+    let logDuringShouldYield = false;
+
+    beforeEach(() => {
+      jest.resetModules();
+      ReactFeatureFlags = require('shared/ReactFeatureFlags');
+      ReactFeatureFlags.debugRenderPhaseSideEffectsForStrictMode = false;
+
+      jest.mock('scheduler', () => {
+        const actual = require.requireActual('scheduler/unstable_mock');
+        return {
+          ...actual,
+          unstable_shouldYield() {
+            if (logDuringShouldYield) {
+              actual.unstable_yieldValue('shouldYield');
+            }
+            return actual.unstable_shouldYield();
+          },
+        };
+      });
+
+      React = require('react');
+      ReactNoop = require('react-noop-renderer');
+      Scheduler = require('scheduler');
+
+      React = require('react');
+    });
+
+    afterEach(() => {
+      jest.mock('scheduler', () =>
+        require.requireActual('scheduler/unstable_mock'),
+      );
+    });
+
+    it('using public APIs to trigger real world scenario', async () => {
+      // This test reproduces a case where React's Scheduler task timed out but
+      // the `shouldYield` method returned true. The bug was that React fell
+      // into an infinite loop, because it would enter the work loop then
+      // immediately yield back to Scheduler.
+      //
+      // (The next test in this suite covers the same case. The difference is
+      // that this test only uses public APIs, whereas the next test mocks
+      // `shouldYield` to check when it is called.)
+      function Text({text}) {
+        return text;
+      }
+
+      function App({step}) {
+        return (
+          <>
+            <Text text="A" />
+            <TriggerErstwhileSchedulerBug />
+            <Text text="B" />
+            <TriggerErstwhileSchedulerBug />
+            <Text text="C" />
+          </>
+        );
+      }
+
+      function TriggerErstwhileSchedulerBug() {
+        // This triggers a once-upon-a-time bug in Scheduler that caused
+        // `shouldYield` to return true even though the current task expired.
+        Scheduler.unstable_advanceTime(10000);
+        Scheduler.unstable_requestPaint();
+        return null;
+      }
+
+      await ReactNoop.act(async () => {
+        ReactNoop.render(<App />);
+        expect(Scheduler).toFlushUntilNextPaint([]);
+        expect(Scheduler).toFlushUntilNextPaint([]);
+      });
+    });
+
+    it('mock Scheduler module to check if `shouldYield` is called', async () => {
+      // This test reproduces a bug where React's Scheduler task timed out but
+      // the `shouldYield` method returned true. Usually we try not to mock
+      // internal methods, but I've made an exception here since the point is
+      // specifically to test that React is reslient to the behavior of a
+      // Scheduler API. That being said, feel free to rewrite or delete this
+      // test if/when the API changes.
+      function Text({text}) {
+        Scheduler.unstable_yieldValue(text);
+        return text;
+      }
+
+      function App({step}) {
+        return (
+          <>
+            <Text text="A" />
+            <Text text="B" />
+            <Text text="C" />
+          </>
+        );
+      }
+
+      await ReactNoop.act(async () => {
+        // Partially render the tree, then yield
+        ReactNoop.render(<App />);
+        expect(Scheduler).toFlushAndYieldThrough(['A']);
+
+        // Start logging whenever shouldYield is called
+        logDuringShouldYield = true;
+        // Let's call it once to confirm the mock actually works
+        Scheduler.unstable_shouldYield();
+        expect(Scheduler).toHaveYielded(['shouldYield']);
+
+        // Expire the task
+        Scheduler.unstable_advanceTime(10000);
+        // Because the render expired, React should finish the tree without
+        // consulting `shouldYield` again
+        expect(Scheduler).toFlushExpired(['B', 'C']);
+      });
+    });
+  },
+);