Switch to client rendering if root receives update

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,16 @@ describe('ReactDOMServerPartialHydration', () => {
     expect(b.textContent).toBe('B');
 
     const root = ReactDOM.hydrateRoot(container, <App />);
+
     // 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 createHydrationContainer.
   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 createHydrationContainer.
   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