Add Suspense Boundary Context (and unstable_avoidThisFallback) (#15578)

* Avoidable suspense boundaries

* Move the context out of SuspenseComponent

* Use setDefaultShallowSuspenseContext instead of passing 0

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,16 @@ import {
   pushHostContextForEventComponent,
   pushHostContextForEventTarget,
 } from './ReactFiberHostContext';
+import {
+  suspenseStackCursor,
+  pushSuspenseContext,
+  popSuspenseContext,
+  InvisibleParentSuspenseContext,
+  ForceSuspenseFallback,
+  hasSuspenseContext,
+  setDefaultShallowSuspenseContext,
+  addSubtreeSuspenseContext,
+} from './ReactFiberSuspenseContext';
 import {
   pushProvider,
   propagateContextChange,
@@ -1394,32 +1405,62 @@ 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 parent.
+    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
+    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 = setDefaultShallowSuspenseContext(suspenseContext);
+
+  pushSuspenseContext(workInProgress, suspenseContext);
+
   if (__DEV__) {
     if ('maxDuration' in nextProps) {
       if (!didWarnAboutMaxDuration) {
@@ -1472,6 +1513,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 +1755,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 +1772,10 @@ function updateDehydratedSuspenseComponent(
   workInProgress: Fiber,
   renderExpirationTime: ExpirationTime,
 ) {
+  pushSuspenseContext(
+    workInProgress,
+    setDefaultShallowSuspenseContext(suspenseStackCursor.current),
+  );
   const suspenseInstance = (workInProgress.stateNode: SuspenseInstance);
   if (current === null) {
     // During the first pass, we'll bail out and not drill into the children.
@@ -2131,6 +2179,10 @@ function beginWork(
                 renderExpirationTime,
               );
             } else {
+              pushSuspenseContext(
+                workInProgress,
+                setDefaultShallowSuspenseContext(suspenseStackCursor.current),
+              );
               // The primary children do not have pending work with sufficient
               // priority. Bailout.
               const child = bailoutOnAlreadyFinishedWork(
@@ -2146,11 +2198,20 @@ function beginWork(
                 return null;
               }
             }
+          } else {
+            pushSuspenseContext(
+              workInProgress,
+              setDefaultShallowSuspenseContext(suspenseStackCursor.current),
+            );
           }
           break;
         }
         case DehydratedSuspenseComponent: {
           if (enableSuspenseServerRenderer) {
+            pushSuspenseContext(
+              workInProgress,
+              setDefaultShallowSuspenseContext(suspenseStackCursor.current),
+            );
             // 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

@@ -14,13 +14,30 @@ 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;
-  }
+export function shouldCaptureSuspense(
+  workInProgress: Fiber,
+  hasInvisibleParent: boolean,
+): boolean {
   // 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 (hasInvisibleParent) {
+    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,83 @@
+/**
+ * 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:
+
+// 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 setDefaultShallowSuspenseContext(
+  parentContext: SuspenseContext,
+): SuspenseContext {
+  return parentContext & SubtreeSuspenseContextMask;
+}
+
+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

@@ -13,6 +13,7 @@ import type {ExpirationTime} from './ReactFiberExpirationTime';
 import type {CapturedValue} from './ReactCapturedValue';
 import type {Update} from './ReactUpdateQueue';
 import type {Thenable} from './ReactFiberScheduler';
+import type {SuspenseContext} from './ReactFiberSuspenseContext';
 
 import {unstable_wrap as Schedule_tracing_wrap} from 'scheduler/tracing';
 import getComponentName from 'shared/getComponentName';
@@ -55,6 +56,12 @@ import {
 import {logError} from './ReactFiberCommitWork';
 import {getStackByFiberInDevAndProd} from './ReactCurrentFiber';
 import {popHostContainer, popHostContext} from './ReactFiberHostContext';
+import {
+  suspenseStackCursor,
+  InvisibleParentSuspenseContext,
+  hasSuspenseContext,
+  popSuspenseContext,
+} from './ReactFiberSuspenseContext';
 import {
   isContextProvider as isLegacyContextProvider,
   popContext as popLegacyContext,
@@ -206,12 +213,17 @@ function throwException(
 
     checkForWrongSuspensePriorityInDEV(sourceFiber);
 
+    let hasInvisibleParentBoundary = hasSuspenseContext(
+      suspenseStackCursor.current,
+      (InvisibleParentSuspenseContext: SuspenseContext),
+    );
+
     // Schedule the nearest Suspense to re-render the timed out view.
     let workInProgress = returnFiber;
     do {
       if (
         workInProgress.tag === SuspenseComponent &&
-        shouldCaptureSuspense(workInProgress)
+        shouldCaptureSuspense(workInProgress, hasInvisibleParentBoundary)
       ) {
         // Found the nearest boundary.
 
@@ -274,6 +286,13 @@ function throwException(
 
         workInProgress.effectTag |= ShouldCapture;
         workInProgress.expirationTime = renderExpirationTime;
+
+        if (!hasInvisibleParentBoundary) {
+          // TODO: If we're not in an invisible subtree, then we need to mark this render
+          // pass as needing to suspend for longer to avoid showing this fallback state.
+          // We could do it here or when we render the fallback.
+        }
+
         return;
       } else if (
         enableSuspenseServerRenderer &&
@@ -408,6 +427,7 @@ function unwindWork(
       return null;
     }
     case SuspenseComponent: {
+      popSuspenseContext(workInProgress);
       const effectTag = workInProgress.effectTag;
       if (effectTag & ShouldCapture) {
         workInProgress.effectTag = (effectTag & ~ShouldCapture) | DidCapture;
@@ -419,6 +439,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 +487,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;