[New Scheduler] Fix: Suspending an expired update

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() {