[New Scheduler] Fix: Suspending an expired update (#15326)

When an async update expires, React renders at the expiration time that
corresponds to the current time, not at the original update's expiration
time. That way, all the expired work in the tree is flushed in a
single batch.

This is implemented inside `renderRoot` by comparing the next render
expiration time to the current time. If the current time is later,
`renderRoot` will restart at the later time.

Because of poor factoring, the check is currently performed right before
entering the work loop. But the work loop is usually entered multiple
times in a single callback: each time a component throws or suspends.
This led to an infinite loop where React would detect that an update
expired, restart at the current time, make a bit of progress, suspend,
check for expired work again, and start the loop again.

I fixed this by moving the expired work check to the beginning of
`renderRoot`, so that it is not performed every time something suspends.
This isn't ideal, because you could technically still fall into a loop
if more than 10ms lapse in between exiting `renderRoot` and entering it
again. The proper fix is to lift the check outside of `renderRoot`
entirely so that the function can restart without checking for expired
work again. Since this is exceedingly unlikely (and this whole thing is
still behind a flag), I'll do the better fix in an already-planned
follow up to fork `renderRoot` into separate functions for sync and
async work.

packages/react-reconciler/src/ReactFiberScheduler.new.js

@@ -742,26 +742,38 @@ function renderRoot(
     }
 
     startWorkLoopTimer(workInProgress);
+
+    // TODO: Fork renderRoot into renderRootSync and renderRootAsync
+    if (isSync) {
+      if (expirationTime !== Sync) {
+        // An async update expired. There may be other expired updates on
+        // this root. We should render all the expired work in a
+        // single batch.
+        const currentTime = requestCurrentTime();
+        if (currentTime < expirationTime) {
+          // Restart at the current time.
+          workPhase = prevWorkPhase;
+          resetContextDependencies();
+          ReactCurrentDispatcher.current = prevDispatcher;
+          if (enableSchedulerTracing) {
+            __interactionsRef.current = ((prevInteractions: any): Set<
+              Interaction,
+            >);
+          }
+          return renderRoot.bind(null, root, currentTime);
+        }
+      }
+    } else {
+      // Since we know we're in a React event, we can clear the current
+      // event time. The next update will compute a new event time.
+      currentEventTime = NoWork;
+    }
+
     do {
       try {
         if (isSync) {
-          if (expirationTime !== Sync) {
-            // An async update expired. There may be other expired updates on
-            // this root. We should render all the expired work in a
-            // single batch.
-            const currentTime = requestCurrentTime();
-            if (currentTime < expirationTime) {
-              // Restart at the current time.
-              workPhase = prevWorkPhase;
-              ReactCurrentDispatcher.current = prevDispatcher;
-              return renderRoot.bind(null, root, currentTime);
-            }
-          }
           workLoopSync();
         } else {
-          // Since we know we're in a React event, we can clear the current
-          // event time. The next update will compute a new event time.
-          currentEventTime = NoWork;
           workLoop();
         }
         break;

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

@@ -827,6 +827,50 @@ describe('ReactSuspenseWithNoopRenderer', () => {
     expect(ReactNoop.getChildren()).toEqual([span('goodbye')]);
   });
 
+  it('a suspended update that expires', async () => {
+    // Regression test. This test used to fall into an infinite loop.
+    function ExpensiveText({text}) {
+      // This causes the update to expire.
+      Scheduler.advanceTime(10000);
+      // Then something suspends.
+      return <AsyncText text={text} ms={200000} />;
+    }
+
+    function App() {
+      return (
+        <Suspense fallback="Loading...">
+          <ExpensiveText text="A" />
+          <ExpensiveText text="B" />
+          <ExpensiveText text="C" />
+        </Suspense>
+      );
+    }
+
+    ReactNoop.render(<App />);
+    expect(Scheduler).toFlushAndYield([
+      'Suspend! [A]',
+      'Suspend! [B]',
+      'Suspend! [C]',
+    ]);
+    expect(ReactNoop).toMatchRenderedOutput('Loading...');
+
+    await advanceTimers(200000);
+    expect(Scheduler).toHaveYielded([
+      'Promise resolved [A]',
+      'Promise resolved [B]',
+      'Promise resolved [C]',
+    ]);
+
+    expect(Scheduler).toFlushAndYield(['A', 'B', 'C']);
+    expect(ReactNoop).toMatchRenderedOutput(
+      <React.Fragment>
+        <span prop="A" />
+        <span prop="B" />
+        <span prop="C" />
+      </React.Fragment>,
+    );
+  });
+
   describe('sync mode', () => {
     it('times out immediately', async () => {
       function App() {