Start prerendering Suspense retries immediately

When a component suspends and is replaced by a fallback, we should
start prerendering the fallback immediately, even before any new data
is received. During the retry, we can enter prerender mode directly
if we're sure that no new data was received since we last attempted
to render the boundary.

To do this, when completing the fallback, we leave behind a pending
retry lane on the Suspense boundary. Previously we only did this once
a promise resolved, but by assigning a lane during the complete phase,
we will know that there's speculative work to be done.

Then, upon committing the fallback, we mark the retry lane as suspended
— but only if nothing was pinged or updated in the meantime. That
allows us to immediately enter prerender mode (i.e. render without
skipping any siblings) when performing the retry.

Author: acdlite

packages/react-cache/src/__tests__/ReactCacheOld-test.internal.js

@@ -18,6 +18,7 @@ let Suspense;
 let TextResource;
 let textResourceShouldFail;
 let waitForAll;
+let waitForPaint;
 let assertLog;
 let waitForThrow;
 let act;
@@ -37,6 +38,7 @@ describe('ReactCache', () => {
     waitForAll = InternalTestUtils.waitForAll;
     assertLog = InternalTestUtils.assertLog;
     waitForThrow = InternalTestUtils.waitForThrow;
+    waitForPaint = InternalTestUtils.waitForPaint;
     act = InternalTestUtils.act;
 
     TextResource = createResource(
@@ -119,7 +121,12 @@ describe('ReactCache', () => {
     const root = ReactNoop.createRoot();
     root.render(<App />);
 
-    await waitForAll(['Suspend! [Hi]', 'Loading...']);
+    await waitForAll([
+      'Suspend! [Hi]',
+      'Loading...',
+
+      ...(gate('enableSiblingPrerendering') ? ['Suspend! [Hi]'] : []),
+    ]);
 
     jest.advanceTimersByTime(100);
     assertLog(['Promise resolved [Hi]']);
@@ -138,7 +145,12 @@ describe('ReactCache', () => {
     const root = ReactNoop.createRoot();
     root.render(<App />);
 
-    await waitForAll(['Suspend! [Hi]', 'Loading...']);
+    await waitForAll([
+      'Suspend! [Hi]',
+      'Loading...',
+
+      ...(gate('enableSiblingPrerendering') ? ['Suspend! [Hi]'] : []),
+    ]);
 
     textResourceShouldFail = true;
     let error;
@@ -179,12 +191,19 @@ describe('ReactCache', () => {
 
     if (__DEV__) {
       await expect(async () => {
-        await waitForAll(['App', 'Loading...']);
+        await waitForAll([
+          'App',
+          'Loading...',
+
+          ...(gate('enableSiblingPrerendering') ? ['App'] : []),
+        ]);
       }).toErrorDev([
         'Invalid key type. Expected a string, number, symbol, or ' +
           "boolean, but instead received: [ 'Hi', 100 ]\n\n" +
           'To use non-primitive values as keys, you must pass a hash ' +
           'function as the second argument to createResource().',
+
+        ...(gate('enableSiblingPrerendering') ? ['Invalid key type'] : []),
       ]);
     } else {
       await waitForAll(['App', 'Loading...']);
@@ -204,7 +223,7 @@ describe('ReactCache', () => {
         <AsyncText ms={100} text={3} />
       </Suspense>,
     );
-    await waitForAll(['Suspend! [1]', 'Loading...']);
+    await waitForPaint(['Suspend! [1]', 'Loading...']);
     jest.advanceTimersByTime(100);
     assertLog(['Promise resolved [1]']);
     await waitForAll([1, 'Suspend! [2]', 1, 'Suspend! [2]', 'Suspend! [3]']);
@@ -225,25 +244,18 @@ describe('ReactCache', () => {
       </Suspense>,
     );
 
-    await waitForAll([1, 'Suspend! [4]', 'Loading...']);
-
-    await act(() => jest.advanceTimersByTime(100));
-    assertLog([
-      'Promise resolved [4]',
-
+    await waitForAll([
       1,
-      4,
-      'Suspend! [5]',
+      'Suspend! [4]',
+      'Loading...',
       1,
-      4,
+      'Suspend! [4]',
       'Suspend! [5]',
-
-      'Promise resolved [5]',
-      1,
-      4,
-      5,
     ]);
 
+    await act(() => jest.advanceTimersByTime(100));
+    assertLog(['Promise resolved [4]', 'Promise resolved [5]', 1, 4, 5]);
+
     expect(root).toMatchRenderedOutput('145');
 
     // We've now rendered values 1, 2, 3, 4, 5, over our limit of 3. The least
@@ -263,24 +275,14 @@ describe('ReactCache', () => {
       // 2 and 3 suspend because they were evicted from the cache
       'Suspend! [2]',
       'Loading...',
-    ]);
 
-    await act(() => jest.advanceTimersByTime(100));
-    assertLog([
-      'Promise resolved [2]',
-
-      1,
-      2,
-      'Suspend! [3]',
       1,
-      2,
+      'Suspend! [2]',
       'Suspend! [3]',
-
-      'Promise resolved [3]',
-      1,
-      2,
-      3,
     ]);
+
+    await act(() => jest.advanceTimersByTime(100));
+    assertLog(['Promise resolved [2]', 'Promise resolved [3]', 1, 2, 3]);
     expect(root).toMatchRenderedOutput('123');
   });
 
@@ -355,7 +357,12 @@ describe('ReactCache', () => {
       </Suspense>,
     );
 
-    await waitForAll(['Suspend! [Hi]', 'Loading...']);
+    await waitForAll([
+      'Suspend! [Hi]',
+      'Loading...',
+
+      ...(gate('enableSiblingPrerendering') ? ['Suspend! [Hi]'] : []),
+    ]);
 
     resolveThenable('Hi');
     // This thenable improperly resolves twice. We should not update the

packages/react-dom/src/__tests__/ReactDOMForm-test.js

@@ -1459,13 +1459,23 @@ describe('ReactDOMForm', () => {
         </Suspense>,
       ),
     );
-    assertLog(['Suspend! [Count: 0]', 'Loading...']);
+    assertLog([
+      'Suspend! [Count: 0]',
+      'Loading...',
+
+      ...(gate('enableSiblingPrerendering') ? ['Suspend! [Count: 0]'] : []),
+    ]);
     await act(() => resolveText('Count: 0'));
     assertLog(['Count: 0']);
 
     // Dispatch outside of a transition. This will trigger a loading state.
     await act(() => dispatch());
-    assertLog(['Suspend! [Count: 1]', 'Loading...']);
+    assertLog([
+      'Suspend! [Count: 1]',
+      'Loading...',
+
+      ...(gate('enableSiblingPrerendering') ? ['Suspend! [Count: 1]'] : []),
+    ]);
     expect(container.textContent).toBe('Loading...');
 
     await act(() => resolveText('Count: 1'));

packages/react-dom/src/__tests__/ReactDOMSuspensePlaceholder-test.js

@@ -160,7 +160,13 @@ describe('ReactDOMSuspensePlaceholder', () => {
     });
 
     expect(container.textContent).toEqual('Loading...');
-    assertLog(['A', 'Suspend! [B]', 'Loading...']);
+    assertLog([
+      'A',
+      'Suspend! [B]',
+      'Loading...',
+
+      ...(gate('enableSiblingPrerendering') ? ['A', 'Suspend! [B]', 'C'] : []),
+    ]);
     await act(() => {
       resolveText('B');
     });

packages/react-dom/src/__tests__/ReactWrongReturnPointer-test.js

@@ -192,7 +192,13 @@ test('regression (#20932): return pointer is correct before entering deleted tre
   await act(() => {
     root.render(<App />);
   });
-  assertLog(['Suspend! [0]', 'Loading Async...', 'Loading Tail...']);
+  assertLog([
+    'Suspend! [0]',
+    'Loading Async...',
+    'Loading Tail...',
+
+    ...(gate('enableSiblingPrerendering') ? ['Suspend! [0]'] : []),
+  ]);
   await act(() => {
     resolveText(0);
   });
@@ -205,5 +211,7 @@ test('regression (#20932): return pointer is correct before entering deleted tre
     'Loading Async...',
     'Suspend! [1]',
     'Loading Async...',
+
+    ...(gate('enableSiblingPrerendering') ? ['Suspend! [1]'] : []),
   ]);
 });

packages/react-reconciler/src/ReactFiberCompleteWork.js

@@ -155,6 +155,7 @@ import {
   getRenderTargetTime,
   getWorkInProgressTransitions,
   shouldRemainOnPreviousScreen,
+  markSpawnedRetryLane,
 } from './ReactFiberWorkLoop';
 import {
   OffscreenLane,
@@ -600,25 +601,28 @@ function scheduleRetryEffect(
     // Schedule an effect to attach a retry listener to the promise.
     // TODO: Move to passive phase
     workInProgress.flags |= Update;
-  } else {
-    // This boundary suspended, but no wakeables were added to the retry
-    // queue. Check if the renderer suspended commit. If so, this means
-    // that once the fallback is committed, we can immediately retry
-    // rendering again, because rendering wasn't actually blocked. Only
-    // the commit phase.
-    // TODO: Consider a model where we always schedule an immediate retry, even
-    // for normal Suspense. That way the retry can partially render up to the
-    // first thing that suspends.
-    if (workInProgress.flags & ScheduleRetry) {
-      const retryLane =
-        // TODO: This check should probably be moved into claimNextRetryLane
-        // I also suspect that we need some further consolidation of offscreen
-        // and retry lanes.
-        workInProgress.tag !== OffscreenComponent
-          ? claimNextRetryLane()
-          : OffscreenLane;
-      workInProgress.lanes = mergeLanes(workInProgress.lanes, retryLane);
-    }
+  }
+
+  // Check if we need to schedule an immediate retry. This should happen
+  // whenever we unwind a suspended tree without fully rendering its siblings;
+  // we need to begin the retry so we can start prerendering them.
+  //
+  // We also use this mechanism for Suspensey Resources (e.g. stylesheets),
+  // because those don't actually block the render phase, only the commit phase.
+  // So we can start rendering even before the resources are ready.
+  if (workInProgress.flags & ScheduleRetry) {
+    const retryLane =
+      // TODO: This check should probably be moved into claimNextRetryLane
+      // I also suspect that we need some further consolidation of offscreen
+      // and retry lanes.
+      workInProgress.tag !== OffscreenComponent
+        ? claimNextRetryLane()
+        : OffscreenLane;
+    workInProgress.lanes = mergeLanes(workInProgress.lanes, retryLane);
+
+    // Track the lanes that have been scheduled for an immediate retry so that
+    // we can mark them as suspended upon committing the root.
+    markSpawnedRetryLane(retryLane);
   }
 }
 

packages/react-reconciler/src/ReactFiberLane.js

@@ -26,9 +26,11 @@ import {
   syncLaneExpirationMs,
   transitionLaneExpirationMs,
   retryLaneExpirationMs,
+  disableLegacyMode,
 } from 'shared/ReactFeatureFlags';
 import {isDevToolsPresent} from './ReactFiberDevToolsHook';
 import {clz32} from './clz32';
+import {LegacyRoot} from './ReactRootTags';
 
 // Lane values below should be kept in sync with getLabelForLane(), used by react-devtools-timeline.
 // If those values are changed that package should be rebuilt and redeployed.
@@ -753,10 +755,14 @@ export function markRootPinged(root: FiberRoot, pingedLanes: Lanes) {
 
 export function markRootFinished(
   root: FiberRoot,
+  finishedLanes: Lanes,
   remainingLanes: Lanes,
   spawnedLane: Lane,
+  updatedLanes: Lanes,
+  suspendedRetryLanes: Lanes,
 ) {
-  const noLongerPendingLanes = root.pendingLanes & ~remainingLanes;
+  const previouslyPendingLanes = root.pendingLanes;
+  const noLongerPendingLanes = previouslyPendingLanes & ~remainingLanes;
 
   root.pendingLanes = remainingLanes;
 
@@ -812,6 +818,37 @@ export function markRootFinished(
       NoLanes,
     );
   }
+
+  // suspendedRetryLanes represents the retry lanes spawned by new Suspense
+  // boundaries during this render that were not later pinged.
+  //
+  // These lanes were marked as pending on their associated Suspense boundary
+  // fiber during the render phase so that we could start rendering them
+  // before new data streams in. As soon as the fallback commits, we can try
+  // to render them again.
+  //
+  // But since we know they're still suspended, we can skip straight to the
+  // "prerender" mode (i.e. don't skip over siblings after something
+  // suspended) instead of the regular mode (i.e. unwind and skip the siblings
+  // as soon as something suspends to unblock the rest of the update).
+  if (
+    suspendedRetryLanes !== NoLanes &&
+    // Note that we only do this if there were no updates since we started
+    // rendering. This mirrors the logic in markRootUpdated — whenever we
+    // receive an update, we reset all the suspended and pinged lanes.
+    updatedLanes === NoLanes &&
+    !(disableLegacyMode && root.tag === LegacyRoot)
+  ) {
+    // We also need to avoid marking a retry lane as suspended if it was already
+    // pending before this render. We can't say these are now suspended if they
+    // weren't included in our attempt.
+    const freshlySpawnedRetryLanes =
+      suspendedRetryLanes &
+      // Remove any retry lane that was already pending before our just-finished
+      // attempt, and also wasn't included in that attempt.
+      ~(previouslyPendingLanes & ~finishedLanes);
+    root.suspendedLanes |= freshlySpawnedRetryLanes;
+  }
 }
 
 function markSpawnedDeferredLane(