Try not. Do... or do not. Hydrate Suspense Boundaries. (#32851)

Assertively claim a SuspenseInstance. We already know we're hydrating.

If there's no match, it throws anyway. So there's no other code path.

Author: sebmarkbage

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -203,7 +203,6 @@ import {
   pushFallbackTreeSuspenseHandler,
   pushOffscreenSuspenseHandler,
   reuseSuspenseHandlerOnStack,
-  popSuspenseHandler,
 } from './ReactFiberSuspenseContext';
 import {
   pushHiddenContext,
@@ -245,7 +244,7 @@ import {
   claimHydratableSingleton,
   tryToClaimNextHydratableInstance,
   tryToClaimNextHydratableTextInstance,
-  tryToClaimNextHydratableSuspenseInstance,
+  claimNextHydratableSuspenseInstance,
   warnIfHydrating,
   queueHydrationError,
 } from './ReactFiberHydrationContext';
@@ -2151,24 +2150,14 @@ function updateSuspenseComponent(
       } else {
         pushFallbackTreeSuspenseHandler(workInProgress);
       }
-      tryToClaimNextHydratableSuspenseInstance(workInProgress);
-      // This could've been a dehydrated suspense component.
-      const suspenseState: null | SuspenseState = workInProgress.memoizedState;
-      if (suspenseState !== null) {
-        const dehydrated = suspenseState.dehydrated;
-        if (dehydrated !== null) {
-          return mountDehydratedSuspenseComponent(
-            workInProgress,
-            dehydrated,
-            renderLanes,
-          );
-        }
-      }
-      // If hydration didn't succeed, fall through to the normal Suspense path.
-      // To avoid a stack mismatch we need to pop the Suspense handler that we
-      // pushed above. This will become less awkward when move the hydration
-      // logic to its own fiber.
-      popSuspenseHandler(workInProgress);
+      // This throws if we fail to hydrate.
+      const dehydrated: SuspenseInstance =
+        claimNextHydratableSuspenseInstance(workInProgress);
+      return mountDehydratedSuspenseComponent(
+        workInProgress,
+        dehydrated,
+        renderLanes,
+      );
     }
 
     const nextPrimaryChildren = nextProps.children;

packages/react-reconciler/src/ReactFiberHydrationContext.js

@@ -272,7 +272,10 @@ function tryHydrateText(fiber: Fiber, nextInstance: any) {
   return false;
 }
 
-function tryHydrateSuspense(fiber: Fiber, nextInstance: any) {
+function tryHydrateSuspense(
+  fiber: Fiber,
+  nextInstance: any,
+): null | SuspenseInstance {
   // fiber is a SuspenseComponent Fiber
   const suspenseInstance = canHydrateSuspenseInstance(
     nextInstance,
@@ -298,9 +301,8 @@ function tryHydrateSuspense(fiber: Fiber, nextInstance: any) {
     // While a Suspense Instance does have children, we won't step into
     // it during the first pass. Instead, we'll reenter it later.
     nextHydratableInstance = null;
-    return true;
   }
-  return false;
+  return suspenseInstance;
 }
 
 export const HydrationMismatchException: mixed = new Error(
@@ -423,15 +425,16 @@ function tryToClaimNextHydratableTextInstance(fiber: Fiber): void {
   }
 }
 
-function tryToClaimNextHydratableSuspenseInstance(fiber: Fiber): void {
-  if (!isHydrating) {
-    return;
-  }
+function claimNextHydratableSuspenseInstance(fiber: Fiber): SuspenseInstance {
   const nextInstance = nextHydratableInstance;
-  if (!nextInstance || !tryHydrateSuspense(fiber, nextInstance)) {
+  const suspenseInstance = nextInstance
+    ? tryHydrateSuspense(fiber, nextInstance)
+    : null;
+  if (suspenseInstance === null) {
     warnNonHydratedInstance(fiber, nextInstance);
-    throwOnHydrationMismatch(fiber);
+    throw throwOnHydrationMismatch(fiber);
   }
+  return suspenseInstance;
 }
 
 export function tryToClaimNextHydratableFormMarkerInstance(
@@ -790,7 +793,7 @@ export {
   claimHydratableSingleton,
   tryToClaimNextHydratableInstance,
   tryToClaimNextHydratableTextInstance,
-  tryToClaimNextHydratableSuspenseInstance,
+  claimNextHydratableSuspenseInstance,
   prepareToHydrateHostInstance,
   prepareToHydrateHostTextInstance,
   prepareToHydrateHostSuspenseInstance,