Avoidable suspense boundaries

Author: sebmarkbage

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -12,6 +12,7 @@ import type {Fiber} from './ReactFiber';
 import type {FiberRoot} from './ReactFiberRoot';
 import type {ExpirationTime} from './ReactFiberExpirationTime';
 import type {SuspenseState} from './ReactFiberSuspenseComponent';
+import type {SuspenseContext} from './ReactFiberSuspenseContext';
 
 import checkPropTypes from 'prop-types/checkPropTypes';
 
@@ -104,6 +105,17 @@ import {
   pushHostContextForEventComponent,
   pushHostContextForEventTarget,
 } from './ReactFiberHostContext';
+import {
+  suspenseStackCursor,
+  pushSuspenseContext,
+  popSuspenseContext,
+  DefaultShallowSuspenseContext,
+  InvisibleParentSuspenseContext,
+  ForceSuspenseFallback,
+  hasSuspenseContext,
+  setShallowSuspenseContext,
+  addSubtreeSuspenseContext,
+} from './ReactFiberSuspenseContext';
 import {
   pushProvider,
   propagateContextChange,
@@ -1394,32 +1406,65 @@ function updateSuspenseComponent(
   const mode = workInProgress.mode;
   const nextProps = workInProgress.pendingProps;
 
+  // This is used by DevTools to force a boundary to suspend.
   if (__DEV__) {
     if (shouldSuspend(workInProgress)) {
       workInProgress.effectTag |= DidCapture;
     }
   }
 
-  // We should attempt to render the primary children unless this boundary
-  // already suspended during this render (`alreadyCaptured` is true).
-  let nextState: SuspenseState | null = workInProgress.memoizedState;
+  let suspenseContext: SuspenseContext = suspenseStackCursor.current;
 
-  let nextDidTimeout;
-  if ((workInProgress.effectTag & DidCapture) === NoEffect) {
-    // This is the first attempt.
-    nextState = null;
-    nextDidTimeout = false;
-  } else {
+  let nextState = null;
+  let nextDidTimeout = false;
+
+  if (
+    (workInProgress.effectTag & DidCapture) !== NoEffect ||
+    hasSuspenseContext(
+      suspenseContext,
+      (ForceSuspenseFallback: SuspenseContext),
+    )
+  ) {
+    // This either already captured or is a new mount that was forced into its fallback
+    // state by a parernt.
+    const attemptedState: SuspenseState | null = workInProgress.memoizedState;
     // Something in this boundary's subtree already suspended. Switch to
     // rendering the fallback children.
     nextState = {
       fallbackExpirationTime:
-        nextState !== null ? nextState.fallbackExpirationTime : NoWork,
+        attemptedState !== null
+          ? attemptedState.fallbackExpirationTime
+          : NoWork,
     };
     nextDidTimeout = true;
     workInProgress.effectTag &= ~DidCapture;
+  } else {
+    // Attempting the main content the main content
+    if (current === null || current.memoizedState !== null) {
+      // This is a new mount or this boundary is already showing a fallback state.
+      // Mark this subtree context as having at least one invisible parent that could
+      // handle the fallback state.
+      // Boundaries without fallbacks or should be avoided are not considered since
+      // they cannot handle preferred fallback states.
+      if (
+        nextProps.fallback !== undefined &&
+        nextProps.unstable_avoidThisFallback !== true
+      ) {
+        suspenseContext = addSubtreeSuspenseContext(
+          suspenseContext,
+          InvisibleParentSuspenseContext,
+        );
+      }
+    }
   }
 
+  suspenseContext = setShallowSuspenseContext(
+    suspenseContext,
+    DefaultShallowSuspenseContext,
+  );
+
+  pushSuspenseContext(workInProgress, suspenseContext);
+
   if (__DEV__) {
     if ('maxDuration' in nextProps) {
       if (!didWarnAboutMaxDuration) {
@@ -1472,6 +1517,7 @@ function updateSuspenseComponent(
         tryToClaimNextHydratableInstance(workInProgress);
         // This could've changed the tag if this was a dehydrated suspense component.
         if (workInProgress.tag === DehydratedSuspenseComponent) {
+          popSuspenseContext(workInProgress);
           return updateDehydratedSuspenseComponent(
             null,
             workInProgress,
@@ -1713,6 +1759,8 @@ function retrySuspenseComponentWithoutHydrating(
   current.nextEffect = null;
   current.effectTag = Deletion;
 
+  popSuspenseContext(workInProgress);
+
   // Upgrade this work in progress to a real Suspense component.
   workInProgress.tag = SuspenseComponent;
   workInProgress.stateNode = null;
@@ -1728,6 +1776,13 @@ function updateDehydratedSuspenseComponent(
   workInProgress: Fiber,
   renderExpirationTime: ExpirationTime,
 ) {
+  pushSuspenseContext(
+    workInProgress,
+    setShallowSuspenseContext(
+      suspenseStackCursor.current,
+      DefaultShallowSuspenseContext,
+    ),
+  );
   const suspenseInstance = (workInProgress.stateNode: SuspenseInstance);
   if (current === null) {
     // During the first pass, we'll bail out and not drill into the children.
@@ -2131,6 +2186,13 @@ function beginWork(
                 renderExpirationTime,
               );
             } else {
+              pushSuspenseContext(
+                workInProgress,
+                setShallowSuspenseContext(
+                  suspenseStackCursor.current,
+                  DefaultShallowSuspenseContext,
+                ),
+              );
               // The primary children do not have pending work with sufficient
               // priority. Bailout.
               const child = bailoutOnAlreadyFinishedWork(
@@ -2146,11 +2208,26 @@ function beginWork(
                 return null;
               }
             }
+          } else {
+            pushSuspenseContext(
+              workInProgress,
+              setShallowSuspenseContext(
+                suspenseStackCursor.current,
+                DefaultShallowSuspenseContext,
+              ),
+            );
           }
           break;
         }
         case DehydratedSuspenseComponent: {
           if (enableSuspenseServerRenderer) {
+            pushSuspenseContext(
+              workInProgress,
+              setShallowSuspenseContext(
+                suspenseStackCursor.current,
+                DefaultShallowSuspenseContext,
+              ),
+            );
             // We know that this component will suspend again because if it has
             // been unsuspended it has committed as a regular Suspense component.
             // If it needs to be retried, it should have work scheduled on it.

packages/react-reconciler/src/ReactFiberCompleteWork.js

@@ -77,6 +77,7 @@ import {
   getHostContext,
   popHostContainer,
 } from './ReactFiberHostContext';
+import {popSuspenseContext} from './ReactFiberSuspenseContext';
 import {
   isContextProvider as isLegacyContextProvider,
   popContext as popLegacyContext,
@@ -667,6 +668,7 @@ function completeWork(
     case ForwardRef:
       break;
     case SuspenseComponent: {
+      popSuspenseContext(workInProgress);
       const nextState: null | SuspenseState = workInProgress.memoizedState;
       if ((workInProgress.effectTag & DidCapture) !== NoEffect) {
         // Something suspended. Re-render with the fallback children.
@@ -777,6 +779,7 @@ function completeWork(
     }
     case DehydratedSuspenseComponent: {
       if (enableSuspenseServerRenderer) {
+        popSuspenseContext(workInProgress);
         if (current === null) {
           let wasHydrated = popHydrationState(workInProgress);
           invariant(

packages/react-reconciler/src/ReactFiberSuspenseComponent.js

@@ -10,17 +10,43 @@
 import type {Fiber} from './ReactFiber';
 import type {ExpirationTime} from './ReactFiberExpirationTime';
 
+import type {SuspenseContext} from './ReactFiberSuspenseContext';
+import {
+  suspenseStackCursor,
+  InvisibleParentSuspenseContext,
+  hasSuspenseContext,
+} from './ReactFiberSuspenseContext';
+
 export type SuspenseState = {|
   fallbackExpirationTime: ExpirationTime,
 |};
 
 export function shouldCaptureSuspense(workInProgress: Fiber): boolean {
-  // In order to capture, the Suspense component must have a fallback prop.
-  if (workInProgress.memoizedProps.fallback === undefined) {
-    return false;
-  }
   // If it was the primary children that just suspended, capture and render the
   // fallback. Otherwise, don't capture and bubble to the next boundary.
   const nextState: SuspenseState | null = workInProgress.memoizedState;
-  return nextState === null;
+  if (nextState !== null) {
+    return false;
+  }
+  const props = workInProgress.memoizedProps;
+  // In order to capture, the Suspense component must have a fallback prop.
+  if (props.fallback === undefined) {
+    return false;
+  }
+  // Regular boundaries always capture.
+  if (props.unstable_avoidThisFallback !== true) {
+    return true;
+  }
+  // If it's a boundary we should avoid, then we prefer to bubble up to the
+  // parent boundary if it is currently invisible.
+  if (
+    hasSuspenseContext(
+      suspenseStackCursor.current,
+      (InvisibleParentSuspenseContext: SuspenseContext),
+    )
+  ) {
+    return false;
+  }
+  // If the parent is not able to handle it, we must handle it.
+  return true;
 }

packages/react-reconciler/src/ReactFiberSuspenseContext.js

@@ -0,0 +1,79 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @flow
+ */
+
+import type {Fiber} from './ReactFiber';
+import type {StackCursor} from './ReactFiberStack';
+
+import {createCursor, push, pop} from './ReactFiberStack';
+
+export opaque type SuspenseContext = number;
+export opaque type SubtreeSuspenseContext: SuspenseContext = number;
+export opaque type ShallowSuspenseContext: SuspenseContext = number;
+
+const DefaultSuspenseContext: SuspenseContext = 0b00;
+
+// The Suspense Context is split into two parts. The lower bits is
+// inherited deeply down the subtree. The upper bits only affect
+// this immediate suspense boundary and gets reset each new
+// boundary or suspense list.
+const SubtreeSuspenseContextMask: SuspenseContext = 0b01;
+
+// Subtree Flags:
+
+// InvisibleParentSuspenseContext indicates that one of our parent Suspense
+// boundaries is not currently showing visible main content.
+// Either because it is already showing a fallback or is not mounted at all.
+// We can use this to determine if it is desirable to trigger a fallback at
+// the parent. If not, then we might need to trigger undesirable boundaries
+// and/or suspend the commit to avoid hiding the parent content.
+export const InvisibleParentSuspenseContext: SubtreeSuspenseContext = 0b01;
+
+// Shallow Flags:
+
+export const DefaultShallowSuspenseContext: ShallowSuspenseContext = 0b00;
+
+// ForceSuspenseFallback can be used by SuspenseList to force newly added
+// items into their fallback state during one of the render passes.
+export const ForceSuspenseFallback: ShallowSuspenseContext = 0b10;
+
+export const suspenseStackCursor: StackCursor<SuspenseContext> = createCursor(
+  DefaultSuspenseContext,
+);
+
+export function hasSuspenseContext(
+  parentContext: SuspenseContext,
+  flag: SuspenseContext,
+): boolean {
+  return (parentContext & flag) !== 0;
+}
+
+export function setShallowSuspenseContext(
+  parentContext: SuspenseContext,
+  shallowContext: ShallowSuspenseContext,
+): SuspenseContext {
+  return (parentContext & SubtreeSuspenseContextMask) | shallowContext;
+}
+
+export function addSubtreeSuspenseContext(
+  parentContext: SuspenseContext,
+  subtreeContext: SubtreeSuspenseContext,
+): SuspenseContext {
+  return parentContext | subtreeContext;
+}
+
+export function pushSuspenseContext(
+  fiber: Fiber,
+  newContext: SuspenseContext,
+): void {
+  push(suspenseStackCursor, newContext, fiber);
+}
+
+export function popSuspenseContext(fiber: Fiber): void {
+  pop(suspenseStackCursor, fiber);
+}

packages/react-reconciler/src/ReactFiberUnwindWork.js

@@ -55,6 +55,7 @@ import {
 import {logError} from './ReactFiberCommitWork';
 import {getStackByFiberInDevAndProd} from './ReactCurrentFiber';
 import {popHostContainer, popHostContext} from './ReactFiberHostContext';
+import {popSuspenseContext} from './ReactFiberSuspenseContext';
 import {
   isContextProvider as isLegacyContextProvider,
   popContext as popLegacyContext,
@@ -206,6 +207,9 @@ function throwException(
 
     checkForWrongSuspensePriorityInDEV(sourceFiber);
 
+    // TODO: If we're not in an invisible subtree, then we need to mark this render
+    // as needing to suspend for longer to avoid showing this fallback state.
+
     // Schedule the nearest Suspense to re-render the timed out view.
     let workInProgress = returnFiber;
     do {
@@ -408,6 +412,7 @@ function unwindWork(
       return null;
     }
     case SuspenseComponent: {
+      popSuspenseContext(workInProgress);
       const effectTag = workInProgress.effectTag;
       if (effectTag & ShouldCapture) {
         workInProgress.effectTag = (effectTag & ~ShouldCapture) | DidCapture;
@@ -419,6 +424,7 @@ function unwindWork(
     case DehydratedSuspenseComponent: {
       if (enableSuspenseServerRenderer) {
         // TODO: popHydrationState
+        popSuspenseContext(workInProgress);
         const effectTag = workInProgress.effectTag;
         if (effectTag & ShouldCapture) {
           workInProgress.effectTag = (effectTag & ~ShouldCapture) | DidCapture;
@@ -466,6 +472,15 @@ function unwindInterruptedWork(interruptedWork: Fiber) {
     case HostPortal:
       popHostContainer(interruptedWork);
       break;
+    case SuspenseComponent:
+      popSuspenseContext(interruptedWork);
+      break;
+    case DehydratedSuspenseComponent:
+      if (enableSuspenseServerRenderer) {
+        // TODO: popHydrationState
+        popSuspenseContext(interruptedWork);
+      }
+      break;
     case ContextProvider:
       popProvider(interruptedWork);
       break;