Incremental hydration

acdlite · acdlite · commit 31eefc63eeb6 · 2021-10-30T21:03:11.000-04:00
Stores the tree context on the dehydrated Suspense boundary's state
object so it resume where it left off.
diff --git a/packages/react-dom/src/__tests__/ReactDOMUseId-test.js b/packages/react-dom/src/__tests__/ReactDOMUseId-test.js
@@ -10,9 +10,11 @@
 let JSDOM;
 let React;
 let ReactDOM;
+let Scheduler;
 let clientAct;
 let ReactDOMFizzServer;
 let Stream;
+let Suspense;
 let useId;
 let document;
 let writable;
@@ -27,9 +29,11 @@ describe('useId', () => {
     JSDOM = require('jsdom').JSDOM;
     React = require('react');
     ReactDOM = require('react-dom');
+    Scheduler = require('scheduler');
     clientAct = require('jest-react').act;
     ReactDOMFizzServer = require('react-dom/server');
     Stream = require('stream');
+    Suspense = React.Suspense;
     useId = React.unstable_useId;
 
     // Test Environment
@@ -86,6 +90,11 @@ describe('useId', () => {
     }
   }
 
+  function Text({text}) {
+    Scheduler.unstable_yieldValue(text);
+    return text;
+  }
+
   function normalizeTreeIdForTesting(id) {
     const [serverClientPrefix, base32, hookIndex] = id.split(':');
     if (serverClientPrefix === 'r') {
@@ -308,4 +317,141 @@ describe('useId', () => {
       </div>
     `);
   });
+
+  test('basic incremental hydration', async () => {
+    function App() {
+      return (
+        <div>
+          <Suspense fallback="Loading...">
+            <DivWithId label="A" />
+            <DivWithId label="B" />
+          </Suspense>
+          <DivWithId label="C" />
+        </div>
+      );
+    }
+
+    await serverAct(async () => {
+      const {pipe} = ReactDOMFizzServer.renderToPipeableStream(<App />);
+      pipe(writable);
+    });
+    await clientAct(async () => {
+      ReactDOM.hydrateRoot(container, <App />);
+    });
+    expect(container).toMatchInlineSnapshot(`
+      <div
+        id="container"
+      >
+        <div>
+          <!--$-->
+          <div
+            id="101"
+          />
+          <div
+            id="1001"
+          />
+          <!--/$-->
+          <div
+            id="10"
+          />
+        </div>
+      </div>
+    `);
+  });
+
+  test('inserting a sibling before a dehydrated Suspense boundary', async () => {
+    const span = React.createRef(null);
+    function App({showMore}) {
+      // Note: Using a dynamic array so this is treated as an insertion instead
+      // of an update, because Fiber currently allocates a node even for
+      // empty children.
+      const children = [<Text key="A" text="A" />];
+      if (showMore) {
+        // These are client-only nodes. They aren't not included in the initial
+        // server render.
+        children.push(<Text key="B" text="B" />, <DivWithId key="C" />);
+      }
+      children.push(
+        <Suspense key="boundary" fallback="Loading...">
+          <DivWithId />
+          <DivWithId />
+          <span ref={span} />
+        </Suspense>,
+        <DivWithId key="after" />,
+      );
+
+      return children;
+    }
+
+    await serverAct(async () => {
+      const {pipe} = ReactDOMFizzServer.renderToPipeableStream(<App />);
+      pipe(writable);
+    });
+    expect(Scheduler).toHaveYielded(['A']);
+    const dehydratedSpan = container.getElementsByTagName('span')[0];
+    await clientAct(async () => {
+      const root = ReactDOM.hydrateRoot(container, <App />);
+      expect(Scheduler).toFlushUntilNextPaint(['A']);
+      expect(container).toMatchInlineSnapshot(`
+        <div
+          id="container"
+        >
+          A
+          <!-- -->
+          <!--$-->
+          <div
+            id="110"
+          />
+          <div
+            id="1010"
+          />
+          <span />
+          <!--/$-->
+          <div
+            id="11"
+          />
+        </div>
+      `);
+
+      // The inner boundary hasn't hydrated yet
+      expect(span.current).toBe(null);
+
+      // Insert another sibling before the Suspense boundary
+      root.render(<App showMore={true} />);
+    });
+    expect(Scheduler).toHaveYielded([
+      'A',
+      'B',
+      // The update triggers selective hydration so we render again
+      'A',
+      'B',
+    ]);
+    // The insertions should not cause a mismatch.
+    expect(container).toMatchInlineSnapshot(`
+      <div
+        id="container"
+      >
+        A
+        <!-- -->
+        <!--$-->
+        B
+        <div
+          id="CLIENT_GENERATED_ID"
+        />
+        <div
+          id="110"
+        />
+        <div
+          id="1010"
+        />
+        <span />
+        <!--/$-->
+        <div
+          id="11"
+        />
+      </div>
+    `);
+    // Should have hydrated successfully
+    expect(span.current).toBe(dehydratedSpan);
+  });
 });
diff --git a/packages/react-reconciler/src/ReactFiberBeginWork.new.js b/packages/react-reconciler/src/ReactFiberBeginWork.new.js
@@ -1852,6 +1852,7 @@ function validateFunctionComponentInDev(workInProgress: Fiber, Component: any) {
 
 const SUSPENDED_MARKER: SuspenseState = {
   dehydrated: null,
+  treeContext: null,
   retryLane: NoLane,
 };
 
@@ -2700,6 +2701,7 @@ function updateDehydratedSuspenseComponent(
     reenterHydrationStateFromDehydratedSuspenseInstance(
       workInProgress,
       suspenseInstance,
+      suspenseState.treeContext,
     );
     const nextProps = workInProgress.pendingProps;
     const primaryChildren = nextProps.children;
diff --git a/packages/react-reconciler/src/ReactFiberBeginWork.old.js b/packages/react-reconciler/src/ReactFiberBeginWork.old.js
@@ -1852,6 +1852,7 @@ function validateFunctionComponentInDev(workInProgress: Fiber, Component: any) {
 
 const SUSPENDED_MARKER: SuspenseState = {
   dehydrated: null,
+  treeContext: null,
   retryLane: NoLane,
 };
 
@@ -2700,6 +2701,7 @@ function updateDehydratedSuspenseComponent(
     reenterHydrationStateFromDehydratedSuspenseInstance(
       workInProgress,
       suspenseInstance,
+      suspenseState.treeContext,
     );
     const nextProps = workInProgress.pendingProps;
     const primaryChildren = nextProps.children;
diff --git a/packages/react-reconciler/src/ReactFiberHydrationContext.new.js b/packages/react-reconciler/src/ReactFiberHydrationContext.new.js
@@ -17,6 +17,7 @@ import type {
   HostContext,
 } from './ReactFiberHostConfig';
 import type {SuspenseState} from './ReactFiberSuspenseComponent.new';
+import type {TreeContext} from './ReactFiberTreeContext.new';
 
 import {
   HostComponent,
@@ -62,6 +63,10 @@ import {
 } from './ReactFiberHostConfig';
 import {enableSuspenseServerRenderer} from 'shared/ReactFeatureFlags';
 import {OffscreenLane} from './ReactFiberLane.new';
+import {
+  getSuspendedTreeContext,
+  restoreSuspendedTreeContext,
+} from './ReactFiberTreeContext.new';
 
 // The deepest Fiber on the stack involved in a hydration context.
 // This may have been an insertion or a hydration.
@@ -96,6 +101,7 @@ function enterHydrationState(fiber: Fiber): boolean {
 function reenterHydrationStateFromDehydratedSuspenseInstance(
   fiber: Fiber,
   suspenseInstance: SuspenseInstance,
+  treeContext: TreeContext | null,
 ): boolean {
   if (!supportsHydration) {
     return false;
@@ -105,6 +111,9 @@ function reenterHydrationStateFromDehydratedSuspenseInstance(
   );
   hydrationParentFiber = fiber;
   isHydrating = true;
+  if (treeContext !== null) {
+    restoreSuspendedTreeContext(fiber, treeContext);
+  }
   return true;
 }
 
@@ -287,6 +296,7 @@ function tryHydrate(fiber, nextInstance) {
         if (suspenseInstance !== null) {
           const suspenseState: SuspenseState = {
             dehydrated: suspenseInstance,
+            treeContext: getSuspendedTreeContext(),
             retryLane: OffscreenLane,
           };
           fiber.memoizedState = suspenseState;
diff --git a/packages/react-reconciler/src/ReactFiberHydrationContext.old.js b/packages/react-reconciler/src/ReactFiberHydrationContext.old.js
@@ -17,6 +17,7 @@ import type {
   HostContext,
 } from './ReactFiberHostConfig';
 import type {SuspenseState} from './ReactFiberSuspenseComponent.old';
+import type {TreeContext} from './ReactFiberTreeContext.old';
 
 import {
   HostComponent,
@@ -62,6 +63,10 @@ import {
 } from './ReactFiberHostConfig';
 import {enableSuspenseServerRenderer} from 'shared/ReactFeatureFlags';
 import {OffscreenLane} from './ReactFiberLane.old';
+import {
+  getSuspendedTreeContext,
+  restoreSuspendedTreeContext,
+} from './ReactFiberTreeContext.old';
 
 // The deepest Fiber on the stack involved in a hydration context.
 // This may have been an insertion or a hydration.
@@ -96,6 +101,7 @@ function enterHydrationState(fiber: Fiber): boolean {
 function reenterHydrationStateFromDehydratedSuspenseInstance(
   fiber: Fiber,
   suspenseInstance: SuspenseInstance,
+  treeContext: TreeContext | null,
 ): boolean {
   if (!supportsHydration) {
     return false;
@@ -105,6 +111,9 @@ function reenterHydrationStateFromDehydratedSuspenseInstance(
   );
   hydrationParentFiber = fiber;
   isHydrating = true;
+  if (treeContext !== null) {
+    restoreSuspendedTreeContext(fiber, treeContext);
+  }
   return true;
 }
 
@@ -287,6 +296,7 @@ function tryHydrate(fiber, nextInstance) {
         if (suspenseInstance !== null) {
           const suspenseState: SuspenseState = {
             dehydrated: suspenseInstance,
+            treeContext: getSuspendedTreeContext(),
             retryLane: OffscreenLane,
           };
           fiber.memoizedState = suspenseState;
diff --git a/packages/react-reconciler/src/ReactFiberSuspenseComponent.new.js b/packages/react-reconciler/src/ReactFiberSuspenseComponent.new.js
@@ -11,6 +11,8 @@ import type {ReactNodeList, Wakeable} from 'shared/ReactTypes';
 import type {Fiber} from './ReactInternalTypes';
 import type {SuspenseInstance} from './ReactFiberHostConfig';
 import type {Lane} from './ReactFiberLane.new';
+import type {TreeContext} from './ReactFiberTreeContext.new';
+
 import {SuspenseComponent, SuspenseListComponent} from './ReactWorkTags';
 import {NoFlags, DidCapture} from './ReactFiberFlags';
 import {
@@ -40,6 +42,7 @@ export type SuspenseState = {|
   // here to indicate that it is dehydrated (flag) and for quick access
   // to check things like isSuspenseInstancePending.
   dehydrated: null | SuspenseInstance,
+  treeContext: null | TreeContext,
   // Represents the lane we should attempt to hydrate a dehydrated boundary at.
   // OffscreenLane is the default for dehydrated boundaries.
   // NoLane is the default for normal boundaries, which turns into "normal" pri.
diff --git a/packages/react-reconciler/src/ReactFiberSuspenseComponent.old.js b/packages/react-reconciler/src/ReactFiberSuspenseComponent.old.js
@@ -11,6 +11,8 @@ import type {ReactNodeList, Wakeable} from 'shared/ReactTypes';
 import type {Fiber} from './ReactInternalTypes';
 import type {SuspenseInstance} from './ReactFiberHostConfig';
 import type {Lane} from './ReactFiberLane.old';
+import type {TreeContext} from './ReactFiberTreeContext.old';
+
 import {SuspenseComponent, SuspenseListComponent} from './ReactWorkTags';
 import {NoFlags, DidCapture} from './ReactFiberFlags';
 import {
@@ -40,6 +42,7 @@ export type SuspenseState = {|
   // here to indicate that it is dehydrated (flag) and for quick access
   // to check things like isSuspenseInstancePending.
   dehydrated: null | SuspenseInstance,
+  treeContext: null | TreeContext,
   // Represents the lane we should attempt to hydrate a dehydrated boundary at.
   // OffscreenLane is the default for dehydrated boundaries.
   // NoLane is the default for normal boundaries, which turns into "normal" pri.
diff --git a/packages/react-reconciler/src/ReactFiberTreeContext.new.js b/packages/react-reconciler/src/ReactFiberTreeContext.new.js
@@ -63,6 +63,12 @@ import {getIsHydrating} from './ReactFiberHydrationContext.new';
 import {clz32} from './clz32';
 import {Forked, NoFlags} from './ReactFiberFlags';
 
+export type TreeContext = {
+  id: number,
+  length: number,
+  overflow: string,
+};
+
 // TODO: Use the unified fiber stack module instead of this local one?
 // Intentionally not using it yet to derisk the initial implementation, because
 // the way we push/pop these values is a bit unusual. If there's a mistake, I'd
@@ -228,6 +234,33 @@ export function popTreeContext(workInProgress: Fiber) {
   }
 }
 
+export function getSuspendedTreeContext(): TreeContext | null {
+  warnIfNotHydrating();
+  if (treeContextProvider !== null) {
+    return {
+      id: treeContextId,
+      overflow: treeContextOverflow,
+    };
+  } else {
+    return null;
+  }
+}
+
+export function restoreSuspendedTreeContext(
+  workInProgress: Fiber,
+  suspendedContext: TreeContext,
+) {
+  warnIfNotHydrating();
+
+  idStack[idStackIndex++] = treeContextId;
+  idStack[idStackIndex++] = treeContextOverflow;
+  idStack[idStackIndex++] = treeContextProvider;
+
+  treeContextId = suspendedContext.id;
+  treeContextOverflow = suspendedContext.overflow;
+  treeContextProvider = workInProgress;
+}
+
 function warnIfNotHydrating() {
   if (__DEV__) {
     if (!getIsHydrating()) {
diff --git a/packages/react-reconciler/src/ReactFiberTreeContext.old.js b/packages/react-reconciler/src/ReactFiberTreeContext.old.js
