Unhide Suspense trees without entanglement

When a Suspense boundary is in its fallback state, you cannot switch
back to the main content without also finishing any updates inside the
tree that might have been skipped. That would be a form of tearing.

Before we fixed this in #18411, the way this bug manifested was that a
boundary was suspended by an update that originated from a child
component (as opposed to props from a parent). While the fallback was
showing, it received another update, this time at high priority. React
would render the high priority update without also including the
original update. That would cause the fallback to switch back to the
main content, since the update that caused the tree to suspend was no
longer part of the render. But then, React would immediately try to
render the original update, which would again suspend and show the
fallback, leading to a momentary flicker in the UI.

The approach added in #18411 is, when receiving a high priority update
to a Suspense tree that's in its fallback state is to bail out, keep
showing the fallback and finish the update in the rest of the tree.
After that commits, render again at the original priority. Because low
priority expiration times are inclusive of higher priority expiration
times, this ensures that all the updates are committed together.

The new approach in this commit is to turn `renderExpirationTime` into a
context-like value that lives on the stack. Then, when unhiding the
Suspense boundary, we can push a new `renderExpirationTime` that is
inclusive of both the high pri update and the original update that
suspended. Then the boundary can be unblocked in a single render pass.

An advantage of the old approach is that by deferring the work of
unhiding, there's less work to do in the high priority update.

The key advantage of the new approach is that it solves the consistency
problem without having to entangle the entire root.

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

@@ -187,6 +187,7 @@ import {
   renderDidSuspendDelayIfPossible,
   markUnprocessedUpdateTime,
   getWorkInProgressRoot,
+  pushRenderExpirationTime,
 } from './ReactFiberWorkLoop.new';
 
 import {disableLogs, reenableLogs} from 'shared/ConsolePatchingDev';
@@ -569,17 +570,30 @@ function updateOffscreenComponent(
   const nextProps: OffscreenProps = workInProgress.pendingProps;
   const nextChildren = nextProps.children;
 
+  let subtreeRenderTime = renderExpirationTime;
   if (current !== null) {
     if (nextProps.mode === 'hidden') {
       // TODO: Should currently be unreachable because Offscreen is only used as
       // an implementation detail of Suspense. Once this is a public API, it
       // will need to create an OffscreenState.
     } else {
-      // Clear the offscreen state.
-      workInProgress.memoizedState = null;
+      const prevState: OffscreenState | null = current.memoizedState;
+      if (prevState !== null) {
+        const baseTime = prevState.baseTime;
+        subtreeRenderTime = !isSameOrHigherPriority(
+          baseTime,
+          renderExpirationTime,
+        )
+          ? baseTime
+          : renderExpirationTime;
+
+        // Since we're not hidden anymore, reset the state
+        workInProgress.memoizedState = null;
+      }
     }
   }
 
+  pushRenderExpirationTime(workInProgress, subtreeRenderTime);
   reconcileChildren(
     current,
     workInProgress,
@@ -1651,33 +1665,32 @@ function validateFunctionComponentInDev(workInProgress: Fiber, Component: any) {
   }
 }
 
-function mountSuspenseState(
+const SUSPENDED_MARKER: SuspenseState = {
+  dehydrated: null,
+  retryTime: NoWork,
+};
+
+function mountSuspenseOffscreenState(
   renderExpirationTime: ExpirationTimeOpaque,
-): SuspenseState {
+): OffscreenState {
   return {
-    dehydrated: null,
     baseTime: renderExpirationTime,
-    retryTime: NoWork,
   };
 }
 
-function updateSuspenseState(
-  prevSuspenseState: SuspenseState,
+function updateSuspenseOffscreenState(
+  prevOffscreenState: OffscreenState,
   renderExpirationTime: ExpirationTimeOpaque,
-): SuspenseState {
-  const prevSuspendedTime = prevSuspenseState.baseTime;
+): OffscreenState {
+  const prevBaseTime = prevOffscreenState.baseTime;
   return {
-    dehydrated: null,
+    // Choose whichever time is inclusive of the other one. This represents
+    // the union of all the levels that suspended.
     baseTime:
-      // Choose whichever time is inclusive of the other one. This represents
-      // the union of all the levels that suspended.
-      !isSameExpirationTime(
-        prevSuspendedTime,
-        (NoWork: ExpirationTimeOpaque),
-      ) && !isSameOrHigherPriority(prevSuspendedTime, renderExpirationTime)
-        ? prevSuspendedTime
+      !isSameExpirationTime(prevBaseTime, (NoWork: ExpirationTimeOpaque)) &&
+      !isSameOrHigherPriority(prevBaseTime, renderExpirationTime)
+        ? prevBaseTime
         : renderExpirationTime,
-    retryTime: NoWork,
   };
 }
 
@@ -1692,26 +1705,15 @@ function shouldRemainOnFallback(
   // For example, SuspenseList coordinates when nested content appears.
   if (current !== null) {
     const suspenseState: SuspenseState = current.memoizedState;
-    if (suspenseState !== null) {
-      // Currently showing a fallback. If the current render includes
-      // the level that triggered the fallback, we must continue showing it,
-      // regardless of what the Suspense context says.
-      const baseTime = suspenseState.baseTime;
-      if (
-        !isSameExpirationTime(baseTime, (NoWork: ExpirationTimeOpaque)) &&
-        !isSameOrHigherPriority(baseTime, renderExpirationTime)
-      ) {
-        return true;
-      }
-      // Otherwise, fall through to check the Suspense context.
-    } else {
+    if (suspenseState === null) {
       // Currently showing content. Don't hide it, even if ForceSuspenseFallack
       // is true. More precise name might be "ForceRemainSuspenseFallback".
       // Note: This is a factoring smell. Can't remain on a fallback if there's
       // no fallback to remain on.
       return false;
     }
   }
+
   // Not currently showing content. Consult the Suspense context.
   return hasSuspenseContext(
     suspenseContext,
@@ -1725,20 +1727,6 @@ function getRemainingWorkInPrimaryTree(
   renderExpirationTime,
 ) {
   const currentChildExpirationTime = current.childExpirationTime_opaque;
-  const currentSuspenseState: SuspenseState = current.memoizedState;
-  if (currentSuspenseState !== null) {
-    // This boundary already timed out. Check if this render includes the level
-    // that previously suspended.
-    const baseTime = currentSuspenseState.baseTime;
-    if (
-      !isSameExpirationTime(baseTime, (NoWork: ExpirationTimeOpaque)) &&
-      !isSameOrHigherPriority(baseTime, renderExpirationTime)
-    ) {
-      // There's pending work at a lower level that might now be unblocked.
-      return baseTime;
-    }
-  }
-
   if (
     !isSameOrHigherPriority(currentChildExpirationTime, renderExpirationTime)
   ) {
@@ -1880,8 +1868,10 @@ function updateSuspenseComponent(
         renderExpirationTime,
       );
       const primaryChildFragment: Fiber = (workInProgress.child: any);
-      primaryChildFragment.memoizedState = ({baseTime: NoWork}: OffscreenState);
-      workInProgress.memoizedState = mountSuspenseState(renderExpirationTime);
+      primaryChildFragment.memoizedState = mountSuspenseOffscreenState(
+        renderExpirationTime,
+      );
+      workInProgress.memoizedState = SUSPENDED_MARKER;
       return fallbackFragment;
     } else {
       const nextPrimaryChildren = nextProps.children;
@@ -1935,14 +1925,10 @@ function updateSuspenseComponent(
               renderExpirationTime,
             );
             const primaryChildFragment: Fiber = (workInProgress.child: any);
-            primaryChildFragment.memoizedState = ({
-              baseTime: NoWork,
-            }: OffscreenState);
-            workInProgress.memoizedState = updateSuspenseState(
-              current.memoizedState,
+            primaryChildFragment.memoizedState = mountSuspenseOffscreenState(
               renderExpirationTime,
             );
-
+            workInProgress.memoizedState = SUSPENDED_MARKER;
             return fallbackChildFragment;
           }
         }
@@ -1959,18 +1945,21 @@ function updateSuspenseComponent(
           renderExpirationTime,
         );
         const primaryChildFragment: Fiber = (workInProgress.child: any);
-        primaryChildFragment.memoizedState = ({
-          baseTime: NoWork,
-        }: OffscreenState);
+        const prevOffscreenState: OffscreenState | null = (current.child: any)
+          .memoizedState;
+        primaryChildFragment.memoizedState =
+          prevOffscreenState === null
+            ? mountSuspenseOffscreenState(renderExpirationTime)
+            : updateSuspenseOffscreenState(
+                prevOffscreenState,
+                renderExpirationTime,
+              );
         primaryChildFragment.childExpirationTime_opaque = getRemainingWorkInPrimaryTree(
           current,
           workInProgress,
           renderExpirationTime,
         );
-        workInProgress.memoizedState = updateSuspenseState(
-          current.memoizedState,
-          renderExpirationTime,
-        );
+        workInProgress.memoizedState = SUSPENDED_MARKER;
         return fallbackChildFragment;
       } else {
         const nextPrimaryChildren = nextProps.children;
@@ -1997,17 +1986,23 @@ function updateSuspenseComponent(
           renderExpirationTime,
         );
         const primaryChildFragment: Fiber = (workInProgress.child: any);
-        primaryChildFragment.memoizedState = ({
-          baseTime: NoWork,
-        }: OffscreenState);
+        const prevOffscreenState: OffscreenState | null = (current.child: any)
+          .memoizedState;
+        primaryChildFragment.memoizedState =
+          prevOffscreenState === null
+            ? mountSuspenseOffscreenState(renderExpirationTime)
+            : updateSuspenseOffscreenState(
+                prevOffscreenState,
+                renderExpirationTime,
+              );
         primaryChildFragment.childExpirationTime_opaque = getRemainingWorkInPrimaryTree(
           current,
           workInProgress,
           renderExpirationTime,
         );
         // Skip the primary children, and continue working on the
         // fallback children.
-        workInProgress.memoizedState = mountSuspenseState(renderExpirationTime);
+        workInProgress.memoizedState = SUSPENDED_MARKER;
         return fallbackChildFragment;
       } else {
         // Still haven't timed out. Continue rendering the children, like we
@@ -3384,6 +3379,10 @@ function beginWork(
             return null;
           }
         }
+        case OffscreenComponent: {
+          pushRenderExpirationTime(workInProgress, renderExpirationTime);
+          break;
+        }
       }
       return bailoutOnAlreadyFinishedWork(
         current,

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

@@ -125,6 +125,7 @@ import {
   renderDidSuspend,
   renderDidSuspendDelayIfPossible,
   renderHasNotSuspendedYet,
+  popRenderExpirationTime,
 } from './ReactFiberWorkLoop.new';
 import {createFundamentalStateInstance} from './ReactFiberFundamental.new';
 import {
@@ -1291,6 +1292,7 @@ function completeWork(
       }
       break;
     case OffscreenComponent: {
+      popRenderExpirationTime(workInProgress);
       if (current !== null) {
         const nextState: OffscreenState | null = workInProgress.memoizedState;
         const prevState: OffscreenState | null = current.memoizedState;

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

@@ -55,7 +55,7 @@ import {
   didNotFindHydratableSuspenseInstance,
 } from './ReactFiberHostConfig';
 import {enableSuspenseServerRenderer} from 'shared/ReactFeatureFlags';
-import {Never, NoWork} from './ReactFiberExpirationTime.new';
+import {Never} from './ReactFiberExpirationTime.new';
 
 // The deepest Fiber on the stack involved in a hydration context.
 // This may have been an insertion or a hydration.
@@ -231,7 +231,6 @@ function tryHydrate(fiber, nextInstance) {
         if (suspenseInstance !== null) {
           const suspenseState: SuspenseState = {
             dehydrated: suspenseInstance,
-            baseTime: NoWork,
             retryTime: Never,
           };
           fiber.memoizedState = suspenseState;

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

@@ -29,10 +29,6 @@ export type SuspenseState = {|
   // here to indicate that it is dehydrated (flag) and for quick access
   // to check things like isSuspenseInstancePending.
   dehydrated: null | SuspenseInstance,
-  // Represents the work that was deprioritized when we committed the fallback.
-  // The work outside the boundary already committed at this level, so we cannot
-  // unhide the content without including it.
-  baseTime: ExpirationTimeOpaque,
   // Represents the earliest expiration time we should attempt to hydrate
   // a dehydrated boundary at.
   // Never is the default for dehydrated boundaries.

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

@@ -20,6 +20,7 @@ import {
   ContextProvider,
   SuspenseComponent,
   SuspenseListComponent,
+  OffscreenComponent,
 } from './ReactWorkTags';
 import {DidCapture, NoEffect, ShouldCapture} from './ReactSideEffectTags';
 import {enableSuspenseServerRenderer} from 'shared/ReactFeatureFlags';
@@ -33,6 +34,7 @@ import {
   popTopLevelContextObject as popTopLevelLegacyContextObject,
 } from './ReactFiberContext.new';
 import {popProvider} from './ReactFiberNewContext.new';
+import {popRenderExpirationTime} from './ReactFiberWorkLoop.new';
 
 import invariant from 'shared/invariant';
 
@@ -105,6 +107,9 @@ function unwindWork(
     case ContextProvider:
       popProvider(workInProgress);
       return null;
+    case OffscreenComponent:
+      popRenderExpirationTime(workInProgress);
+      return null;
     default:
       return null;
   }
@@ -141,6 +146,9 @@ function unwindInterruptedWork(interruptedWork: Fiber) {
     case ContextProvider:
       popProvider(interruptedWork);
       break;
+    case OffscreenComponent:
+      popRenderExpirationTime(interruptedWork);
+      break;
     default:
       break;
   }

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

@@ -15,6 +15,7 @@ import type {Interaction} from 'scheduler/src/Tracing';
 import type {SuspenseConfig} from './ReactFiberSuspenseConfig';
 import type {SuspenseState} from './ReactFiberSuspenseComponent.new';
 import type {Effect as HookEffect} from './ReactFiberHooks.new';
+import type {StackCursor} from './ReactFiberStack.new';
 
 import {
   warnAboutDeprecatedLifecycles,
@@ -168,6 +169,11 @@ import {
   getIsUpdatingOpaqueValueInRenderPhaseInDEV,
 } from './ReactFiberHooks.new';
 import {createCapturedValue} from './ReactCapturedValue';
+import {
+  push as pushToStack,
+  pop as popFromStack,
+  createCursor,
+} from './ReactFiberStack.new';
 
 import {
   recordCommitTime,
@@ -231,6 +237,12 @@ let workInProgressRoot: FiberRoot | null = null;
 let workInProgress: Fiber | null = null;
 // The expiration time we're rendering
 let renderExpirationTime: ExpirationTimeOpaque = NoWork;
+
+// Stack that allows components to channge renderExpirationTime for its subtree
+const renderExpirationTimeCursor: StackCursor<ExpirationTimeOpaque> = createCursor(
+  NoWork,
+);
+
 // Whether to root completed, errored, suspended, etc.
 let workInProgressRootExitStatus: RootExitStatus = RootIncomplete;
 // A fatal error, if one is thrown
@@ -1262,6 +1274,19 @@ export function flushControlled(fn: () => mixed): void {
   }
 }
 
+export function pushRenderExpirationTime(
+  fiber: Fiber,
+  subtreeRenderTime: ExpirationTimeOpaque,
+) {
+  pushToStack(renderExpirationTimeCursor, renderExpirationTime, fiber);
+  renderExpirationTime = subtreeRenderTime;
+}
+
+export function popRenderExpirationTime(fiber: Fiber) {
+  renderExpirationTime = renderExpirationTimeCursor.current;
+  popFromStack(renderExpirationTimeCursor, fiber);
+}
+
 function prepareFreshStack(root, expirationTime) {
   root.finishedWork = null;
   root.finishedExpirationTime_opaque = NoWork;