Switch to client rendering if root receives update

If a hydration root receives an update before the outermost shell has
finished hydrating, we should give up hydrating and switch to
client rendering.

Since the shell is expected to commit quickly, this doesn't happen that
often. The most common sequence is something in the shell suspends, and
then the user quickly navigates to a different screen, triggering a
top-level update.

Instead of immediately switching to client rendering, we could first
attempt to hydration at higher priority, like we do for updates that
occur inside nested dehydrated trees.

But since this case is expected to be rare, and mainly only happens when
the shell is suspended, an attempt at higher priority would likely end
up suspending again anyway, so it would be wasted effort. Implementing
it this way would also require us to add a new lane especially for root
hydration. For simplicity's sake, we'll immediately switch to client
rendering. In the future, if we find another use case for a root
hydration lane, we'll reconsider.

packages/jest-react/src/internalAct.js

@@ -22,7 +22,7 @@ import enqueueTask from 'shared/enqueueTask';
 
 let actingUpdatesScopeDepth = 0;
 
-export function act(scope: () => Thenable<mixed> | void) {
+export function act<T>(scope: () => Thenable<T> | T): Thenable<T> {
   if (Scheduler.unstable_flushAllWithoutAsserting === undefined) {
     throw Error(
       'This version of `act` requires a special mock build of Scheduler.',
@@ -66,20 +66,21 @@ export function act(scope: () => Thenable<mixed> | void) {
   // returned and 2) we could use async/await. Since it's only our used in
   // our test suite, we should be able to.
   try {
-    const thenable = scope();
+    const result = scope();
     if (
-      typeof thenable === 'object' &&
-      thenable !== null &&
-      typeof thenable.then === 'function'
+      typeof result === 'object' &&
+      result !== null &&
+      typeof result.then === 'function'
     ) {
+      const thenableResult: Thenable<T> = (result: any);
       return {
-        then(resolve: () => void, reject: (error: mixed) => void) {
-          thenable.then(
-            () => {
+        then(resolve, reject) {
+          thenableResult.then(
+            returnValue => {
               flushActWork(
                 () => {
                   unwind();
-                  resolve();
+                  resolve(returnValue);
                 },
                 error => {
                   unwind();
@@ -95,13 +96,19 @@ export function act(scope: () => Thenable<mixed> | void) {
         },
       };
     } else {
+      const returnValue: T = (result: any);
       try {
         // TODO: Let's not support non-async scopes at all in our tests. Need to
         // migrate existing tests.
         let didFlushWork;
         do {
           didFlushWork = Scheduler.unstable_flushAllWithoutAsserting();
         } while (didFlushWork);
+        return {
+          then(resolve, reject) {
+            resolve(returnValue);
+          },
+        };
       } finally {
         unwind();
       }

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

@@ -140,10 +140,10 @@ describe('ReactDOMFizzShellHydration', () => {
     }
   }
 
-  // function Text({text}) {
-  //   Scheduler.unstable_yieldValue(text);
-  //   return text;
-  // }
+  function Text({text}) {
+    Scheduler.unstable_yieldValue(text);
+    return text;
+  }
 
   function AsyncText({text}) {
     readText(text);
@@ -213,4 +213,34 @@ describe('ReactDOMFizzShellHydration', () => {
     expect(Scheduler).toHaveYielded(['Shell']);
     expect(container.textContent).toBe('Shell');
   });
+
+  test('updating the root before the shell hydrates forces a client render', async () => {
+    function App() {
+      return <AsyncText text="Shell" />;
+    }
+
+    // Server render
+    await resolveText('Shell');
+    await serverAct(async () => {
+      const {pipe} = ReactDOMFizzServer.renderToPipeableStream(<App />);
+      pipe(writable);
+    });
+    expect(Scheduler).toHaveYielded(['Shell']);
+
+    // Clear the cache and start rendering on the client
+    resetTextCache();
+
+    // Hydration suspends because the data for the shell hasn't loaded yet
+    const root = await clientAct(async () => {
+      return ReactDOM.hydrateRoot(container, <App />);
+    });
+    expect(Scheduler).toHaveYielded(['Suspend! [Shell]']);
+    expect(container.textContent).toBe('Shell');
+
+    await clientAct(async () => {
+      root.render(<Text text="New screen" />);
+    });
+    expect(Scheduler).toHaveYielded(['New screen']);
+    expect(container.textContent).toBe('New screen');
+  });
 });

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

@@ -1966,21 +1966,18 @@ describe('ReactDOMServerPartialHydration', () => {
     expect(b.textContent).toBe('B');
 
     const root = ReactDOM.hydrateRoot(container, <App />);
+
+    // Commit the shell
+
     // Increase hydration priority to higher than "offscreen".
     root.unstable_scheduleHydration(b);
 
     suspend = true;
 
     await act(async () => {
       if (gate(flags => flags.enableSyncDefaultUpdates)) {
-        React.startTransition(() => {
-          root.render(<App />);
-        });
-
         expect(Scheduler).toFlushAndYieldThrough(['Before', 'After']);
       } else {
-        root.render(<App />);
-
         expect(Scheduler).toFlushAndYieldThrough(['Before']);
         // This took a long time to render.
         Scheduler.unstable_advanceTime(1000);

packages/react-dom/src/client/ReactDOMRoot.js

@@ -60,6 +60,7 @@ import {
 
 import {
   createContainer,
+  createHydrationContainer,
   updateContainer,
   findHostInstanceWithNoPortals,
   registerMutableSourceForHydration,
@@ -261,10 +262,10 @@ export function hydrateRoot(
     }
   }
 
-  const root = createContainer(
+  const root = createHydrationContainer(
+    initialChildren,
     container,
     ConcurrentRoot,
-    true, // hydrate
     hydrationCallbacks,
     isStrictMode,
     concurrentUpdatesByDefaultOverride,
@@ -284,9 +285,6 @@ export function hydrateRoot(
     }
   }
 
-  // Render the initial children
-  updateContainer(initialChildren, root, null, null);
-
   return new ReactDOMHydrationRoot(root);
 }
 

packages/react-reconciler/src/ReactFiberReconciler.js

@@ -16,6 +16,7 @@ import {enableNewReconciler} from 'shared/ReactFeatureFlags';
 
 import {
   createContainer as createContainer_old,
+  createHydrationContainer as createHydrationContainer_old,
   updateContainer as updateContainer_old,
   batchedUpdates as batchedUpdates_old,
   deferredUpdates as deferredUpdates_old,
@@ -53,6 +54,7 @@ import {
 
 import {
   createContainer as createContainer_new,
+  createHydrationContainer as createHydrationContainer_new,
   updateContainer as updateContainer_new,
   batchedUpdates as batchedUpdates_new,
   deferredUpdates as deferredUpdates_new,
@@ -91,6 +93,9 @@ import {
 export const createContainer = enableNewReconciler
   ? createContainer_new
   : createContainer_old;
+export const createHydrationContainer = enableNewReconciler
+  ? createHydrationContainer_new
+  : createHydrationContainer_old;
 export const updateContainer = enableNewReconciler
   ? updateContainer_new
   : updateContainer_old;

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

@@ -57,6 +57,7 @@ import {
   requestEventTime,
   requestUpdateLane,
   scheduleUpdateOnFiber,
+  scheduleInitialHydrationOnRoot,
   flushRoot,
   batchedUpdates,
   flushSync,
@@ -244,6 +245,8 @@ function findHostInstanceWithWarning(
 export function createContainer(
   containerInfo: Container,
   tag: RootTag,
+  // TODO: We can remove hydration-specific stuff from createContainer once
+  // we delete legacy mode. The new root API uses createDehydratedContainer.
   hydrate: boolean,
   hydrationCallbacks: null | SuspenseHydrationCallbacks,
   isStrictMode: boolean,
@@ -265,6 +268,49 @@ export function createContainer(
   );
 }
 
+export function createHydrationContainer(
+  initialChildren: ReactNodeList,
+  containerInfo: Container,
+  tag: RootTag,
+  hydrationCallbacks: null | SuspenseHydrationCallbacks,
+  isStrictMode: boolean,
+  concurrentUpdatesByDefaultOverride: null | boolean,
+  identifierPrefix: string,
+  onRecoverableError: (error: mixed) => void,
+  transitionCallbacks: null | TransitionTracingCallbacks,
+): OpaqueRoot {
+  const hydrate = true;
+  const root = createFiberRoot(
+    containerInfo,
+    tag,
+    hydrate,
+    hydrationCallbacks,
+    isStrictMode,
+    concurrentUpdatesByDefaultOverride,
+    identifierPrefix,
+    onRecoverableError,
+    transitionCallbacks,
+  );
+
+  // TODO: Move this to FiberRoot constructor
+  root.context = getContextForSubtree(null);
+
+  // Schedule the initial render. In a hydration root, this is different from
+  // a regular update because the initial render must match was was rendered
+  // on the server.
+  const current = root.current;
+  const eventTime = requestEventTime();
+  const lane = requestUpdateLane(current);
+  const update = createUpdate(eventTime, lane);
+  // Caution: React DevTools currently depends on this property
+  // being called "element".
+  update.payload = {element: initialChildren};
+  enqueueUpdate(current, update, lane);
+  scheduleInitialHydrationOnRoot(root, lane, eventTime);
+
+  return root;
+}
+
 export function updateContainer(
   element: ReactNodeList,
   container: OpaqueRoot,

packages/react-reconciler/src/ReactFiberReconciler.old.js

@@ -57,6 +57,7 @@ import {
   requestEventTime,
   requestUpdateLane,
   scheduleUpdateOnFiber,
+  scheduleInitialHydrationOnRoot,
   flushRoot,
   batchedUpdates,
   flushSync,
@@ -244,6 +245,8 @@ function findHostInstanceWithWarning(
 export function createContainer(
   containerInfo: Container,
   tag: RootTag,
+  // TODO: We can remove hydration-specific stuff from createContainer once
+  // we delete legacy mode. The new root API uses createDehydratedContainer.
   hydrate: boolean,
   hydrationCallbacks: null | SuspenseHydrationCallbacks,
   isStrictMode: boolean,
@@ -265,6 +268,49 @@ export function createContainer(
   );
 }
 
+export function createHydrationContainer(
+  initialChildren: ReactNodeList,
+  containerInfo: Container,
+  tag: RootTag,
+  hydrationCallbacks: null | SuspenseHydrationCallbacks,
+  isStrictMode: boolean,
+  concurrentUpdatesByDefaultOverride: null | boolean,
+  identifierPrefix: string,
+  onRecoverableError: (error: mixed) => void,
+  transitionCallbacks: null | TransitionTracingCallbacks,
+): OpaqueRoot {
+  const hydrate = true;
+  const root = createFiberRoot(
+    containerInfo,
+    tag,
+    hydrate,
+    hydrationCallbacks,
+    isStrictMode,
+    concurrentUpdatesByDefaultOverride,
+    identifierPrefix,
+    onRecoverableError,
+    transitionCallbacks,
+  );
+
+  // TODO: Move this to FiberRoot constructor
+  root.context = getContextForSubtree(null);
+
+  // Schedule the initial render. In a hydration root, this is different from
+  // a regular update because the initial render must match was was rendered
+  // on the server.
+  const current = root.current;
+  const eventTime = requestEventTime();
+  const lane = requestUpdateLane(current);
+  const update = createUpdate(eventTime, lane);
+  // Caution: React DevTools currently depends on this property
+  // being called "element".
+  update.payload = {element: initialChildren};
+  enqueueUpdate(current, update, lane);
+  scheduleInitialHydrationOnRoot(root, lane, eventTime);
+
+  return root;
+}
+
 export function updateContainer(
   element: ReactNodeList,
   container: OpaqueRoot,

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

@@ -517,8 +517,31 @@ export function scheduleUpdateOnFiber(
       }
     }
 
-    // TODO: Consolidate with `isInterleavedUpdate` check
-    if (root === workInProgressRoot) {
+    if (root.isDehydrated && root.tag !== LegacyRoot) {
+      // This root's shell hasn't hydrated yet. Revert to client rendering.
+      // TODO: Log a recoverable error
+      if (workInProgressRoot === root) {
+        // If this happened during an interleaved event, interrupt the
+        // in-progress hydration. Theoretically, we could attempt to force a
+        // synchronous hydration before switching to client rendering, but the
+        // most common reason the shell hasn't hydrated yet is because it
+        // suspended. So it's very likely to suspend again anyway. For
+        // simplicity, we'll skip that atttempt and go straight to
+        // client rendering.
+        //
+        // Another way to model this would be to give the initial hydration its
+        // own special lane. However, it may not be worth adding a lane solely
+        // for this purpose, so we'll wait until we find another use case before
+        // adding it.
+        //
+        // TODO: Consider only interrupting hydration if the priority of the
+        // update is higher than default.
+        prepareFreshStack(root, NoLanes);
+      }
+      root.isDehydrated = false;
+    } else if (root === workInProgressRoot) {
+      // TODO: Consolidate with `isInterleavedUpdate` check
+
       // Received an update to a tree that's in the middle of rendering. Mark
       // that there was an interleaved update work on this root. Unless the
       // `deferRenderPhaseUpdateToNextBatch` flag is off and this is a render
@@ -564,6 +587,26 @@ export function scheduleUpdateOnFiber(
   return root;
 }
 
+export function scheduleInitialHydrationOnRoot(
+  root: FiberRoot,
+  lane: Lane,
+  eventTime: number,
+) {
+  // This is a special fork of scheduleUpdateOnFiber that is only used to
+  // schedule the initial hydration of a root that has just been created. Most
+  // of the stuff in scheduleUpdateOnFiber can be skipped.
+  //
+  // The main reason for this separate path, though, is to distinguish the
+  // initial children from subsequent updates. In fully client-rendered roots
+  // (createRoot instead of hydrateRoot), all top-level renders are modeled as
+  // updates, but hydration roots are special because the initial render must
+  // match what was rendered on the server.
+  const current = root.current;
+  current.lanes = lane;
+  markRootUpdated(root, lane, eventTime);
+  ensureRootIsScheduled(root, eventTime);
+}
+
 // This is split into a separate function so we can mark a fiber with pending
 // work without treating it as a typical update that originates from an event;
 // e.g. retrying a Suspense boundary isn't an update, but it does schedule work