Suspensey commits in prerendered trees

Prerendering a tree (i.e. with Offscreen) should not suspend the commit
phase, because the content is not yet visible. However, when revealing
a prerendered tree, we should suspend the commit phase if resources in
the prerendered tree haven't finished loading yet.

To do this properly, we need to visit all the visible nodes in the tree
that might possibly suspend. This includes nodes in the current tree,
because even though they were already "mounted", the resources might not
have loaded yet, because we didn't suspend when it was prerendered.

We will need to add this capability to the Offscreen component's
"manual" mode, too. Something like a `ready()` method that returns a
promise that resolves when the tree has fully loaded.

packages/react-reconciler/src/ReactFiberCommitWork.js

@@ -94,7 +94,8 @@ import {
   LayoutMask,
   PassiveMask,
   Visibility,
-  SuspenseyCommit,
+  ShouldSuspendCommit,
+  MaySuspendCommit,
 } from './ReactFiberFlags';
 import getComponentNameFromFiber from 'react-reconciler/src/getComponentNameFromFiber';
 import {
@@ -4064,21 +4065,62 @@ export function commitPassiveUnmountEffects(finishedWork: Fiber): void {
   resetCurrentDebugFiberInDEV();
 }
 
+// If we're inside a brand new tree, or a tree that was already visible, then we
+// should only suspend host components that have a ShouldSuspendCommit flag.
+// Components without it haven't changed since the last commit, so we can skip
+// over those.
+//
+// When we enter a tree that is being revealed (going from hidden -> visible),
+// we need to suspend _any_ component that _may_ suspend. Even if they're
+// already in the "current" tree. Because their visibility has changed, the
+// browser may not have prerendered them yet. So we check the MaySuspendCommit
+// flag instead.
+let suspenseyCommitFlag = ShouldSuspendCommit;
 export function recursivelyAccumulateSuspenseyCommit(parentFiber: Fiber): void {
-  if (parentFiber.subtreeFlags & SuspenseyCommit) {
+  if (parentFiber.subtreeFlags & suspenseyCommitFlag) {
     let child = parentFiber.child;
     while (child !== null) {
-      recursivelyAccumulateSuspenseyCommit(child);
       switch (child.tag) {
+        case OffscreenComponent: {
+          const isHidden =
+            (child.memoizedState: OffscreenState | null) !== null;
+          if (isHidden) {
+            // Don't suspend in hidden trees
+          } else {
+            const current = child.alternate;
+            const wasHidden =
+              current !== null &&
+              (current.memoizedState: OffscreenState | null) !== null;
+            if (wasHidden) {
+              // This tree is being revealed. Visit all newly visible suspensey
+              // instances, even if they're in the current tree.
+              const prevFlags = suspenseyCommitFlag;
+              suspenseyCommitFlag = MaySuspendCommit;
+              recursivelyAccumulateSuspenseyCommit(child);
+              suspenseyCommitFlag = prevFlags;
+            } else {
+              recursivelyAccumulateSuspenseyCommit(child);
+            }
+          }
+          break;
+        }
         case HostComponent:
         case HostHoistable: {
-          if (child.flags & SuspenseyCommit) {
+          recursivelyAccumulateSuspenseyCommit(child);
+          if (child.flags & suspenseyCommitFlag) {
             const type = child.type;
             const props = child.memoizedProps;
-            suspendInstance(type, props);
+            try {
+              suspendInstance(type, props);
+            } catch (error) {
+              captureCommitPhaseError(child, child.return, error);
+            }
           }
           break;
         }
+        default: {
+          recursivelyAccumulateSuspenseyCommit(child);
+        }
       }
       child = child.sibling;
     }

packages/react-reconciler/src/ReactFiberCompleteWork.js

@@ -86,8 +86,9 @@ import {
   MutationMask,
   Passive,
   ForceClientRender,
-  SuspenseyCommit,
+  MaySuspendCommit,
   ScheduleRetry,
+  ShouldSuspendCommit,
 } from './ReactFiberFlags';
 
 import {
@@ -151,6 +152,7 @@ import {
   getRenderTargetTime,
   getWorkInProgressTransitions,
   shouldRemainOnPreviousScreen,
+  getWorkInProgressRootRenderLanes,
 } from './ReactFiberWorkLoop';
 import {
   OffscreenLane,
@@ -527,26 +529,31 @@ function preloadInstanceAndSuspendIfNeeded(
     // safest thing to do is for maySuspendCommit to always return true, but
     // if the renderer is reasonably confident that the underlying resource
     // won't be evicted, it can return false as a performance optimization.
-    workInProgress.flags &= ~SuspenseyCommit;
+    workInProgress.flags &= ~MaySuspendCommit;
     return;
   }
 
-  // Mark this fiber with a flag. We use this right before the commit phase to
-  // find all the fibers that might need to suspend the commit. In the future
-  // we'll also use it when revealing a hidden tree. It gets set even if we
-  // don't end up suspending this particular commit, because if this tree ever
-  // becomes hidden, we might want to suspend before revealing it again.
-  workInProgress.flags |= SuspenseyCommit;
+  // Mark this fiber with a flag. This gets set on all host components that
+  // might possibly suspend, even if they don't need to suspend currently. We
+  // use this when revealing a prerendered tree, because even though the tree
+  // has "mounted", its resources might not have loaded yet.
+  workInProgress.flags |= MaySuspendCommit;
 
   // Check if we're rendering at a "non-urgent" priority. This is the same
   // check that `useDeferredValue` does to determine whether it needs to
   // defer. This is partly for gradual adoption purposes (i.e. shouldn't start
   // suspending until you opt in with startTransition or Suspense) but it
   // also happens to be the desired behavior for the concrete use cases we've
   // thought of so far, like CSS loading, fonts, images, etc.
+  //
+  // We check the "root" render lanes here rather than the "subtree" render
+  // because during a retry or offscreen prerender, the "subtree" render
+  // lanes may include additional "base" lanes that were deferred during
+  // a previous render.
   // TODO: We may decide to expose a way to force a fallback even during a
   // sync update.
-  if (!includesOnlyNonUrgentLanes(renderLanes)) {
+  const rootRenderLanes = getWorkInProgressRootRenderLanes();
+  if (!includesOnlyNonUrgentLanes(rootRenderLanes)) {
     // This is an urgent render. Don't suspend or show a fallback. Also,
     // there's no need to preload, because we're going to commit this
     // synchronously anyway.
@@ -560,7 +567,9 @@ function preloadInstanceAndSuspendIfNeeded(
     const isReady = preloadInstance(type, props);
     if (!isReady) {
       if (shouldRemainOnPreviousScreen()) {
-        // It's OK to suspend. Continue rendering.
+        // It's OK to suspend. Mark the fiber so we know to suspend before the
+        // commit phase. Then continue rendering.
+        workInProgress.flags |= ShouldSuspendCommit;
       } else {
         // Trigger a fallback rather than block the render.
         suspendCommit();

packages/react-reconciler/src/ReactFiberFlags.js

@@ -36,10 +36,11 @@ export const Passive = /*                      */ 0b0000000000000000100000000000
 export const Visibility = /*                   */ 0b0000000000000010000000000000;
 export const StoreConsistency = /*             */ 0b0000000000000100000000000000;
 
-// It's OK to reuse this bit because these flags are mutually exclusive for
+// It's OK to reuse these bits because these flags are mutually exclusive for
 // different fiber types. We should really be doing this for as many flags as
 // possible, because we're about to run out of bits.
 export const ScheduleRetry = StoreConsistency;
+export const ShouldSuspendCommit = Visibility;
 
 export const LifecycleEffectMask =
   Passive | Update | Callback | Ref | Snapshot | StoreConsistency;
@@ -63,7 +64,7 @@ export const Forked = /*                       */ 0b0000000100000000000000000000
 export const RefStatic = /*                    */ 0b0000001000000000000000000000;
 export const LayoutStatic = /*                 */ 0b0000010000000000000000000000;
 export const PassiveStatic = /*                */ 0b0000100000000000000000000000;
-export const SuspenseyCommit = /*              */ 0b0001000000000000000000000000;
+export const MaySuspendCommit = /*             */ 0b0001000000000000000000000000;
 
 // Flag used to identify newly inserted fibers. It isn't reset after commit unlike `Placement`.
 export const PlacementDEV = /*                 */ 0b0010000000000000000000000000;
@@ -103,4 +104,4 @@ export const PassiveMask = Passive | Visibility | ChildDeletion;
 // This allows certain concepts to persist without recalculating them,
 // e.g. whether a subtree contains passive effects or portals.
 export const StaticMask =
-  LayoutStatic | PassiveStatic | RefStatic | SuspenseyCommit;
+  LayoutStatic | PassiveStatic | RefStatic | MaySuspendCommit;

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -169,6 +169,8 @@ import {
   addTransitionToLanesMap,
   getTransitionsForLanes,
   includesOnlyNonUrgentLanes,
+  includesSomeLane,
+  OffscreenLane,
 } from './ReactFiberLane';
 import {
   DiscreteEventPriority,
@@ -1997,17 +1999,27 @@ export function shouldRemainOnPreviousScreen(): boolean {
     // parent Suspense boundary, even outside a transition. Somehow. Otherwise,
     // an uncached promise can fall into an infinite loop.
   } else {
-    if (includesOnlyRetries(workInProgressRootRenderLanes)) {
+    if (
+      includesOnlyRetries(workInProgressRootRenderLanes) ||
+      // In this context, an OffscreenLane counts as a Retry
+      // TODO: It's become increasingly clear that Retries and Offscreen are
+      // deeply connected. They probably can be unified further.
+      includesSomeLane(workInProgressRootRenderLanes, OffscreenLane)
+    ) {
       // During a retry, we can suspend rendering if the nearest Suspense boundary
       // is the boundary of the "shell", because we're guaranteed not to block
       // any new content from appearing.
+      //
+      // The reason we must check if this is a retry is because it guarantees
+      // that suspending the work loop won't block an actual update, because
+      // retries don't "update" anything; they fill in fallbacks that were left
+      // behind by a previous transition.
       return handler === getShellBoundary();
     }
   }
 
   // For all other Lanes besides Transitions and Retries, we should not wait
   // for the data to load.
-  // TODO: We should wait during Offscreen prerendering, too.
   return false;
 }
 

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

@@ -4,7 +4,9 @@ let ReactNoop;
 let resolveSuspenseyThing;
 let getSuspenseyThingStatus;
 let Suspense;
+let Offscreen;
 let SuspenseList;
+let useMemo;
 let Scheduler;
 let act;
 let assertLog;
@@ -18,10 +20,11 @@ describe('ReactSuspenseyCommitPhase', () => {
     ReactNoop = require('react-noop-renderer');
     Scheduler = require('scheduler');
     Suspense = React.Suspense;
-    SuspenseList = React.SuspenseList;
     if (gate(flags => flags.enableSuspenseList)) {
       SuspenseList = React.SuspenseList;
     }
+    Offscreen = React.unstable_Offscreen;
+    useMemo = React.useMemo;
     startTransition = React.startTransition;
     resolveSuspenseyThing = ReactNoop.resolveSuspenseyThing;
     getSuspenseyThingStatus = ReactNoop.getSuspenseyThingStatus;
@@ -279,4 +282,67 @@ describe('ReactSuspenseyCommitPhase', () => {
       </>,
     );
   });
+
+  // @gate enableOffscreen
+  test("host instances don't suspend during prerendering, but do suspend when they are revealed", async () => {
+    function More() {
+      Scheduler.log('More');
+      return <SuspenseyImage src="More" />;
+    }
+
+    function Details({showMore}) {
+      Scheduler.log('Details');
+      const more = useMemo(() => <More />, []);
+      return (
+        <>
+          <div>Main Content</div>
+          <Offscreen mode={showMore ? 'visible' : 'hidden'}>{more}</Offscreen>
+        </>
+      );
+    }
+
+    const root = ReactNoop.createRoot();
+    await act(async () => {
+      root.render(<Details showMore={false} />);
+      // First render the outer component, without the hidden content
+      await waitForPaint(['Details']);
+      expect(root).toMatchRenderedOutput(<div>Main Content</div>);
+    });
+    // Then prerender the hidden content.
+    assertLog(['More', 'Image requested [More]']);
+    // The prerender should commit even though the image is still loading,
+    // because it's hidden.
+    expect(root).toMatchRenderedOutput(
+      <>
+        <div>Main Content</div>
+        <suspensey-thing hidden={true} src="More" />
+      </>,
+    );
+
+    // Reveal the prerendered content. This update should suspend, because the
+    // image that is being revealed still hasn't loaded.
+    await act(() => {
+      startTransition(() => {
+        root.render(<Details showMore={true} />);
+      });
+    });
+    // The More component should not render again, because it was memoized,
+    // and it already prerendered.
+    assertLog(['Details']);
+    expect(root).toMatchRenderedOutput(
+      <>
+        <div>Main Content</div>
+        <suspensey-thing hidden={true} src="More" />
+      </>,
+    );
+
+    // Now resolve the image. The transition should complete.
+    resolveSuspenseyThing('More');
+    expect(root).toMatchRenderedOutput(
+      <>
+        <div>Main Content</div>
+        <suspensey-thing src="More" />
+      </>,
+    );
+  });
 });