Clamp suspense timeout to expiration time

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

@@ -895,8 +895,9 @@ function renderRoot(
         // track any event times. That can happen if we retried but nothing switched
         // from fallback to content. There's no reason to delay doing no work.
         if (workInProgressRootMostRecentEventTime !== Sync) {
-          const msUntilTimeout = computeMsUntilTimeout(
+          let msUntilTimeout = computeMsUntilTimeout(
             workInProgressRootMostRecentEventTime,
+            expirationTime,
           );
           // Don't bother with a very short suspense time.
           if (msUntilTimeout > 10) {
@@ -1843,7 +1844,10 @@ function jnd(timeElapsed: number) {
               : ceil(timeElapsed / 1960) * 1960;
 }
 
-function computeMsUntilTimeout(mostRecentEventTime: ExpirationTime) {
+function computeMsUntilTimeout(
+  mostRecentEventTime: ExpirationTime,
+  committedExpirationTime: ExpirationTime,
+) {
   if (disableYielding) {
     // Timeout immediately when yielding is disabled.
     return 0;
@@ -1853,7 +1857,19 @@ function computeMsUntilTimeout(mostRecentEventTime: ExpirationTime) {
   const currentTimeMs: number = now();
   const timeElapsed = currentTimeMs - eventTimeMs;
 
-  const msUntilTimeout = jnd(timeElapsed) - timeElapsed;
+  let msUntilTimeout = jnd(timeElapsed) - timeElapsed;
+
+  // Compute the time until this render pass would expire.
+  const timeUntilExpirationMs =
+    expirationTimeToMs(committedExpirationTime) + initialTimeMs - currentTimeMs;
+
+  // Clamp the timeout to the expiration time.
+  // TODO: Once the event time is exact instead of inferred from expiration time
+  // we don't need this.
+  if (timeUntilExpirationMs < msUntilTimeout) {
+    msUntilTimeout = timeUntilExpirationMs;
+  }
+
   // This is the value that is passed to `setTimeout`.
   return msUntilTimeout;
 }

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

@@ -1535,9 +1535,20 @@ function renderRoot(root: FiberRoot, isYieldy: boolean): void {
     const timeElapsed = currentTimeMs - eventTimeMs;
 
     let msUntilTimeout = jnd(timeElapsed) - timeElapsed;
+
     if (msUntilTimeout < 10) {
       // Don't bother with a very short suspense time.
       msUntilTimeout = 0;
+    } else {
+      // Compute the time until this render pass would expire.
+      const timeUntilExpirationMs =
+        expirationTimeToMs(suspendedExpirationTime) +
+        originalStartTimeMs -
+        currentTimeMs;
+      // Clamp the timeout to the expiration time.
+      if (timeUntilExpirationMs < msUntilTimeout) {
+        msUntilTimeout = timeUntilExpirationMs;
+      }
     }
 
     const rootExpirationTime = root.expirationTime;

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

@@ -1788,6 +1788,40 @@ describe('ReactSuspenseWithNoopRenderer', () => {
     expect(Scheduler).toFlushAndYield(['B']);
     expect(ReactNoop.getChildren()).toEqual([span('A'), span('B')]);
   });
+
+  it('does not suspend for very long after a higher priority update', async () => {
+    function Foo() {
+      Scheduler.yieldValue('Foo');
+      return (
+        <Suspense fallback={<Text text="Loading..." />}>
+          <AsyncText text="A" ms={5000} />
+        </Suspense>
+      );
+    }
+
+    ReactNoop.interactiveUpdates(() => ReactNoop.render(<Foo />));
+    Scheduler.advanceTime(200);
+    await advanceTimers(200);
+    // Start rendering
+    expect(Scheduler).toFlushAndYield([
+      'Foo',
+      // A suspends
+      'Suspend! [A]',
+      'Loading...',
+    ]);
+    // We're now suspended and we haven't shown anything yet.
+    expect(ReactNoop.getChildren()).toEqual([]);
+
+    // Flush some of the time
+    Scheduler.advanceTime(400);
+    await advanceTimers(400);
+    // We should have already shown the fallback.
+    // When we wrote this test, we inferred the start time of high priority
+    // updates as way earlier in the past. This test ensures that we don't
+    // use this assumption to add a very long JND.
+    expect(Scheduler).toFlushWithoutYielding();
+    expect(ReactNoop.getChildren()).toEqual([span('Loading...')]);
+  });
 });
 
 // TODO: