Retry hydrating at offscreen pri or after ping if suspended

packages/react-dom/src/__tests__/ReactDOMServerPartialHydration-test.js

@@ -28,7 +28,7 @@ describe('ReactDOMServerPartialHydration', () => {
     Suspense = React.Suspense;
   });
 
-  it('hydrates a parent even if a child Suspense boundary is blocked', () => {
+  it('hydrates a parent even if a child Suspense boundary is blocked', async () => {
     let suspend = false;
     let resolve;
     let promise = new Promise(resolvePromise => (resolve = resolvePromise));
@@ -65,31 +65,35 @@ describe('ReactDOMServerPartialHydration', () => {
     // On the client we don't have all data yet but we want to start
     // hydrating anyway.
     suspend = true;
-    ReactDOM.hydrate(<App />, container);
+    const root = ReactDOM.unstable_createRoot(container, {hydrate: true});
+    root.render(<App />);
+    jest.runAllTimers();
 
     // Resolving the promise should continue hydration
+    suspend = false;
     resolve();
+    await promise;
+    jest.runAllTimers();
   });
 
   it('can insert siblings before the dehydrated boundary', () => {
     let suspend = false;
-    let resolve;
-    let promise = new Promise(resolvePromise => (resolve = resolvePromise));
+    let promise = new Promise(() => {});
     let showSibling;
 
     function Child() {
       if (suspend) {
         throw promise;
       } else {
-        return 'Hello';
+        return 'Second';
       }
     }
 
     function Sibling() {
       let [visible, setVisibilty] = React.useState(false);
       showSibling = () => setVisibilty(true);
       if (visible) {
-        return <div>Hello</div>;
+        return <div>First</div>;
       }
       return null;
     }
@@ -123,9 +127,6 @@ describe('ReactDOMServerPartialHydration', () => {
     showSibling();
 
     expect(container.firstChild.firstChild.tagName).toBe('DIV');
-    expect(container.firstChild.firstChild.textContent).toBe('Hello');
-
-    // Resolving the promise should continue hydration
-    resolve();
+    expect(container.firstChild.firstChild.textContent).toBe('First');
   });
 });

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -1397,7 +1397,11 @@ function updateSuspenseComponent(
     tryToClaimNextHydratableInstance(workInProgress);
     // This could've changed the tag if this was a dehydrated suspense component.
     if (workInProgress.tag === DehydratedSuspenseComponent) {
-      return updateDehydratedSuspenseComponent(null, workInProgress);
+      return updateDehydratedSuspenseComponent(
+        null,
+        workInProgress,
+        renderExpirationTime,
+      );
     }
 
     // This is the initial mount. This branch is pretty simple because there's
@@ -1609,8 +1613,37 @@ function updateSuspenseComponent(
 function updateDehydratedSuspenseComponent(
   current: Fiber | null,
   workInProgress: Fiber,
+  renderExpirationTime: ExpirationTime,
 ) {
-  return null;
+  if (current === null) {
+    // During the first pass, we'll bail out and not drill into the children.
+    // Instead, we'll leave the content in place and try to hydrate it later.
+    workInProgress.expirationTime = workInProgress.childExpirationTime = Never;
+    return null;
+  }
+  if ((workInProgress.effectTag & DidCapture) === NoEffect) {
+    // This is the first attempt.
+    const prevProps = current.memoizedProps;
+    const nextProps = workInProgress.pendingProps;
+    if (prevProps !== nextProps) {
+      // TODO: Delete children and upgrade to a regular suspense component without
+      // hydrating.
+    }
+    // TODO: Restore hydration state
+    const nextChildren = nextProps.children;
+    workInProgress.child = mountChildFibers(
+      workInProgress,
+      null,
+      nextChildren,
+      renderExpirationTime,
+    );
+    return workInProgress.child;
+  } else {
+    // Something suspended. Leave the existing children in place.
+    // TODO: In non-concurrent mode, should we commit the nodes we have hydrated so far?
+    workInProgress.child = null;
+    return null;
+  }
 }
 
 function updatePortalComponent(
@@ -1897,6 +1930,13 @@ function beginWork(
           }
           break;
         }
+        case DehydratedSuspenseComponent: {
+          // 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.
+          workInProgress.effectTag |= DidCapture;
+          break;
+        }
       }
       return bailoutOnAlreadyFinishedWork(
         current,
@@ -2067,7 +2107,11 @@ function beginWork(
       );
     }
     case DehydratedSuspenseComponent: {
-      return updateDehydratedSuspenseComponent(current, workInProgress);
+      return updateDehydratedSuspenseComponent(
+        current,
+        workInProgress,
+        renderExpirationTime,
+      );
     }
     default:
       invariant(

packages/react-reconciler/src/ReactFiberCompleteWork.js

@@ -773,6 +773,15 @@ function completeWork(
             'This is probably a bug in React.',
         );
         skipPastDehydratedSuspenseInstance(workInProgress);
+      } else if ((workInProgress.effectTag & DidCapture) === NoEffect) {
+        // This boundary did not suspend so it's now hydrated.
+        // To handle any future suspense cases, we're going to now upgrade it
+        // to a Suspense component. We detach it from the existing current fiber.
+        current.alternate = null;
+        workInProgress.alternate = null;
+        workInProgress.tag = SuspenseComponent;
+        workInProgress.memoizedState = null;
+        workInProgress.stateNode = null;
       }
       break;
     }

packages/react-reconciler/src/ReactFiberScheduler.js

@@ -60,6 +60,8 @@ import {
   HostRoot,
   MemoComponent,
   SimpleMemoComponent,
+  SuspenseComponent,
+  DehydratedSuspenseComponent,
 } from 'shared/ReactWorkTags';
 import {
   enableSchedulerTracing,
@@ -1697,8 +1699,21 @@ function retryTimedOutBoundary(boundaryFiber: Fiber, thenable: Thenable) {
   // resolved, which means at least part of the tree was likely unblocked. Try
   // rendering again, at a new expiration time.
 
-  const retryCache: WeakSet<Thenable> | Set<Thenable> | null =
-    boundaryFiber.stateNode;
+  let retryCache: WeakSet<Thenable> | Set<Thenable> | null;
+  switch (boundaryFiber.tag) {
+    case SuspenseComponent:
+      retryCache = boundaryFiber.stateNode;
+      break;
+    case DehydratedSuspenseComponent:
+      retryCache = boundaryFiber.memoizedState;
+      break;
+    default:
+      invariant(
+        false,
+        'Pinged unknown suspense boundary type. ' +
+          'This is probably a bug in React.',
+      );
+  }
   if (retryCache !== null) {
     // The thenable resolved, so we no longer need to memoize, because it will
     // never be thrown again.

packages/react-reconciler/src/ReactFiberUnwindWork.js

@@ -66,6 +66,7 @@ import {
   markLegacyErrorBoundaryAsFailed,
   isAlreadyFailedLegacyErrorBoundary,
   pingSuspendedRoot,
+  retryTimedOutBoundary,
 } from './ReactFiberScheduler';
 
 import invariant from 'shared/invariant';
@@ -77,6 +78,7 @@ import {
 } from './ReactFiberExpirationTime';
 import {findEarliestOutstandingPriorityLevel} from './ReactFiberPendingPriority';
 
+const PossiblyWeakSet = typeof WeakSet === 'function' ? WeakSet : Set;
 const PossiblyWeakMap = typeof WeakMap === 'function' ? WeakMap : Map;
 
 function createRootErrorUpdate(
@@ -148,6 +150,43 @@ function createClassErrorUpdate(
   return update;
 }
 
+function attachPingListener(
+  root: FiberRoot,
+  renderExpirationTime: ExpirationTime,
+  thenable: Thenable,
+) {
+  // Attach a listener to the promise to "ping" the root and retry. But
+  // only if one does not already exist for the current render expiration
+  // time (which acts like a "thread ID" here).
+  let pingCache = root.pingCache;
+  let threadIDs;
+  if (pingCache === null) {
+    pingCache = root.pingCache = new PossiblyWeakMap();
+    threadIDs = new Set();
+    pingCache.set(thenable, threadIDs);
+  } else {
+    threadIDs = pingCache.get(thenable);
+    if (threadIDs === undefined) {
+      threadIDs = new Set();
+      pingCache.set(thenable, threadIDs);
+    }
+  }
+  if (!threadIDs.has(renderExpirationTime)) {
+    // Memoize using the thread ID to prevent redundant listeners.
+    threadIDs.add(renderExpirationTime);
+    let ping = pingSuspendedRoot.bind(
+      null,
+      root,
+      thenable,
+      renderExpirationTime,
+    );
+    if (enableSchedulerTracing) {
+      ping = Schedule_tracing_wrap(ping);
+    }
+    thenable.then(ping, ping);
+  }
+}
+
 function throwException(
   root: FiberRoot,
   returnFiber: Fiber,
@@ -271,36 +310,7 @@ function throwException(
         // Confirmed that the boundary is in a concurrent mode tree. Continue
         // with the normal suspend path.
 
-        // Attach a listener to the promise to "ping" the root and retry. But
-        // only if one does not already exist for the current render expiration
-        // time (which acts like a "thread ID" here).
-        let pingCache = root.pingCache;
-        let threadIDs;
-        if (pingCache === null) {
-          pingCache = root.pingCache = new PossiblyWeakMap();
-          threadIDs = new Set();
-          pingCache.set(thenable, threadIDs);
-        } else {
-          threadIDs = pingCache.get(thenable);
-          if (threadIDs === undefined) {
-            threadIDs = new Set();
-            pingCache.set(thenable, threadIDs);
-          }
-        }
-        if (!threadIDs.has(renderExpirationTime)) {
-          // Memoize using the thread ID to prevent redundant listeners.
-          threadIDs.add(renderExpirationTime);
-          let ping = pingSuspendedRoot.bind(
-            null,
-            root,
-            thenable,
-            renderExpirationTime,
-          );
-          if (enableSchedulerTracing) {
-            ping = Schedule_tracing_wrap(ping);
-          }
-          thenable.then(ping, ping);
-        }
+        attachPingListener(root, renderExpirationTime, thenable);
 
         let absoluteTimeoutMs;
         if (earliestTimeoutMs === -1) {
@@ -344,7 +354,35 @@ function throwException(
         workInProgress.tag === DehydratedSuspenseComponent &&
         shouldCaptureDehydratedSuspense(workInProgress)
       ) {
-        // TODO
+        attachPingListener(root, renderExpirationTime, thenable);
+
+        // Since we already have a current fiber, we can eagerly add a ping listener.
+        let retryCache = workInProgress.memoizedState;
+        if (retryCache === null) {
+          retryCache = workInProgress.memoizedState = new PossiblyWeakSet();
+          const current = workInProgress.alternate;
+          invariant(
+            current,
+            'A dehydrated suspense boundary must commit before trying to render. ' +
+              'This is probably a bug in React.',
+          );
+          current.memoizedState = retryCache;
+        }
+        // Memoize using the boundary fiber to prevent redundant listeners.
+        if (!retryCache.has(thenable)) {
+          retryCache.add(thenable);
+          let retry = retryTimedOutBoundary.bind(
+            null,
+            workInProgress,
+            thenable,
+          );
+          if (enableSchedulerTracing) {
+            retry = Schedule_tracing_wrap(retry);
+          }
+          thenable.then(retry, retry);
+        }
+        workInProgress.effectTag |= ShouldCapture;
+        workInProgress.expirationTime = renderExpirationTime;
         return;
       }
       // This boundary already captured during this render. Continue to the next
@@ -459,7 +497,12 @@ function unwindWork(
     }
     case DehydratedSuspenseComponent: {
       // TODO: popHydrationState
-      // TODO: Maybe re-render if it captured?
+      const effectTag = workInProgress.effectTag;
+      if (effectTag & ShouldCapture) {
+        workInProgress.effectTag = (effectTag & ~ShouldCapture) | DidCapture;
+        // Captured a suspense effect. Re-render the boundary.
+        return workInProgress;
+      }
       return null;
     }
     case HostPortal: