Deal with fallback content in Partial Hydration (#14884)

* Replace SSR fallback content with new suspense content

* The three states of a Dehydrated Suspense

This introduces three states for dehydrated suspense boundaries

Pending - This means that the tree is still waiting for additional data or
to be populated with its final content.

Fallback - This means that the tree has entered a permanent fallback state
and no more data from the server is to be expected. This means that the
client should take over and try to render instead. The fallback nodes will
be deleted.

Normal - The node has entered its final content and is now ready to be
hydrated.

* Rename retryTimedOutBoundary to resolveRetryThenable

This doesn't just retry. It assumes that resolving a thenable means that
it is ok to clear it from the thenable cache.

We'll reuse the retryTimedOutBoundary logic separately.

* Register a callback to be fired when a boundary changes away from pending

It's now possible to switch from a pending state to either hydrating
or replacing the content.

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

@@ -634,4 +634,228 @@ describe('ReactDOMServerPartialHydration', () => {
     expect(ref.current).toBe(span);
     expect(container.textContent).toBe('Hi');
   });
+
+  it('replaces the fallback with client content if it is not rendered by the server', async () => {
+    let suspend = false;
+    let promise = new Promise(resolvePromise => {});
+    let ref = React.createRef();
+
+    function Child() {
+      if (suspend) {
+        throw promise;
+      } else {
+        return 'Hello';
+      }
+    }
+
+    function App() {
+      return (
+        <div>
+          <Suspense fallback="Loading...">
+            <span ref={ref}>
+              <Child />
+            </span>
+          </Suspense>
+        </div>
+      );
+    }
+
+    // First we render the final HTML. With the streaming renderer
+    // this may have suspense points on the server but here we want
+    // to test the completed HTML. Don't suspend on the server.
+    suspend = true;
+    let finalHTML = ReactDOMServer.renderToString(<App />);
+    let container = document.createElement('div');
+    container.innerHTML = finalHTML;
+
+    expect(container.getElementsByTagName('span').length).toBe(0);
+
+    // On the client we have the data available quickly for some reason.
+    suspend = false;
+    let root = ReactDOM.unstable_createRoot(container, {hydrate: true});
+    root.render(<App />);
+    jest.runAllTimers();
+
+    expect(container.textContent).toBe('Hello');
+
+    let span = container.getElementsByTagName('span')[0];
+    expect(ref.current).toBe(span);
+  });
+
+  it('waits for pending content to come in from the server and then hydrates it', async () => {
+    let suspend = false;
+    let promise = new Promise(resolvePromise => {});
+    let ref = React.createRef();
+
+    function Child() {
+      if (suspend) {
+        throw promise;
+      } else {
+        return 'Hello';
+      }
+    }
+
+    function App() {
+      return (
+        <div>
+          <Suspense fallback="Loading...">
+            <span ref={ref}>
+              <Child />
+            </span>
+          </Suspense>
+        </div>
+      );
+    }
+
+    // We're going to simulate what Fizz will do during streaming rendering.
+
+    // First we generate the HTML of the loading state.
+    suspend = true;
+    let loadingHTML = ReactDOMServer.renderToString(<App />);
+    // Then we generate the HTML of the final content.
+    suspend = false;
+    let finalHTML = ReactDOMServer.renderToString(<App />);
+
+    let container = document.createElement('div');
+    container.innerHTML = loadingHTML;
+
+    let suspenseNode = container.firstChild.firstChild;
+    expect(suspenseNode.nodeType).toBe(8);
+    // Put the suspense node in hydration state.
+    suspenseNode.data = '$?';
+
+    // This will simulates new content streaming into the document and
+    // replacing the fallback with final content.
+    function streamInContent() {
+      let temp = document.createElement('div');
+      temp.innerHTML = finalHTML;
+      let finalSuspenseNode = temp.firstChild.firstChild;
+      let fallbackContent = suspenseNode.nextSibling;
+      let finalContent = finalSuspenseNode.nextSibling;
+      suspenseNode.parentNode.replaceChild(finalContent, fallbackContent);
+      suspenseNode.data = '$';
+      if (suspenseNode._reactRetry) {
+        suspenseNode._reactRetry();
+      }
+    }
+
+    // We're still showing a fallback.
+    expect(container.getElementsByTagName('span').length).toBe(0);
+
+    // Attempt to hydrate the content.
+    suspend = false;
+    let root = ReactDOM.unstable_createRoot(container, {hydrate: true});
+    root.render(<App />);
+    jest.runAllTimers();
+
+    // We're still loading because we're waiting for the server to stream more content.
+    expect(container.textContent).toBe('Loading...');
+
+    // The server now updates the content in place in the fallback.
+    streamInContent();
+
+    // The final HTML is now in place.
+    expect(container.textContent).toBe('Hello');
+    let span = container.getElementsByTagName('span')[0];
+
+    // But it is not yet hydrated.
+    expect(ref.current).toBe(null);
+
+    jest.runAllTimers();
+
+    // Now it's hydrated.
+    expect(ref.current).toBe(span);
+  });
+
+  it('handles an error on the client if the server ends up erroring', async () => {
+    let suspend = false;
+    let promise = new Promise(resolvePromise => {});
+    let ref = React.createRef();
+
+    function Child() {
+      if (suspend) {
+        throw promise;
+      } else {
+        throw new Error('Error Message');
+      }
+    }
+
+    class ErrorBoundary extends React.Component {
+      state = {error: null};
+      static getDerivedStateFromError(error) {
+        return {error};
+      }
+      render() {
+        if (this.state.error) {
+          return <div ref={ref}>{this.state.error.message}</div>;
+        }
+        return this.props.children;
+      }
+    }
+
+    function App() {
+      return (
+        <ErrorBoundary>
+          <div>
+            <Suspense fallback="Loading...">
+              <span ref={ref}>
+                <Child />
+              </span>
+            </Suspense>
+          </div>
+        </ErrorBoundary>
+      );
+    }
+
+    // We're going to simulate what Fizz will do during streaming rendering.
+
+    // First we generate the HTML of the loading state.
+    suspend = true;
+    let loadingHTML = ReactDOMServer.renderToString(<App />);
+
+    let container = document.createElement('div');
+    container.innerHTML = loadingHTML;
+
+    let suspenseNode = container.firstChild.firstChild;
+    expect(suspenseNode.nodeType).toBe(8);
+    // Put the suspense node in hydration state.
+    suspenseNode.data = '$?';
+
+    // This will simulates the server erroring and putting the fallback
+    // as the final state.
+    function streamInError() {
+      suspenseNode.data = '$!';
+      if (suspenseNode._reactRetry) {
+        suspenseNode._reactRetry();
+      }
+    }
+
+    // We're still showing a fallback.
+    expect(container.getElementsByTagName('span').length).toBe(0);
+
+    // Attempt to hydrate the content.
+    suspend = false;
+    let root = ReactDOM.unstable_createRoot(container, {hydrate: true});
+    root.render(<App />);
+    jest.runAllTimers();
+
+    // We're still loading because we're waiting for the server to stream more content.
+    expect(container.textContent).toBe('Loading...');
+
+    // The server now updates the content in place in the fallback.
+    streamInError();
+
+    // The server errored, but we still haven't hydrated. We don't know if the
+    // client will succeed yet, so we still show the loading state.
+    expect(container.textContent).toBe('Loading...');
+    expect(ref.current).toBe(null);
+
+    jest.runAllTimers();
+
+    // Hydrating should've generated an error and replaced the suspense boundary.
+    expect(container.textContent).toBe('Error Message');
+
+    let div = container.getElementsByTagName('div')[0];
+    expect(ref.current).toBe(div);
+  });
 });

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

@@ -94,6 +94,8 @@ describe('ReactDOMServerSuspense', () => {
     );
     const e = c.children[0];
 
-    expect(e.innerHTML).toBe('<div>Children</div><div>Fallback</div>');
+    expect(e.innerHTML).toBe(
+      '<div>Children</div><!--$!--><div>Fallback</div><!--/$-->',
+    );
   });
 });

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

@@ -56,7 +56,7 @@ export type Props = {
 export type Container = Element | Document;
 export type Instance = Element;
 export type TextInstance = Text;
-export type SuspenseInstance = Comment;
+export type SuspenseInstance = Comment & {_reactRetry?: () => void};
 export type HydratableInstance = Instance | TextInstance | SuspenseInstance;
 export type PublicInstance = Element | Text;
 type HostContextDev = {
@@ -89,6 +89,8 @@ if (__DEV__) {
 
 const SUSPENSE_START_DATA = '$';
 const SUSPENSE_END_DATA = '/$';
+const SUSPENSE_PENDING_START_DATA = '$?';
+const SUSPENSE_FALLBACK_START_DATA = '$!';
 
 const STYLE = 'style';
 
@@ -458,7 +460,11 @@ export function clearSuspenseBoundary(
         } else {
           depth--;
         }
-      } else if (data === SUSPENSE_START_DATA) {
+      } else if (
+        data === SUSPENSE_START_DATA ||
+        data === SUSPENSE_PENDING_START_DATA ||
+        data === SUSPENSE_FALLBACK_START_DATA
+      ) {
         depth++;
       }
     }
@@ -554,6 +560,21 @@ export function canHydrateSuspenseInstance(
   return ((instance: any): SuspenseInstance);
 }
 
+export function isSuspenseInstancePending(instance: SuspenseInstance) {
+  return instance.data === SUSPENSE_PENDING_START_DATA;
+}
+
+export function isSuspenseInstanceFallback(instance: SuspenseInstance) {
+  return instance.data === SUSPENSE_FALLBACK_START_DATA;
+}
+
+export function registerSuspenseInstanceRetry(
+  instance: SuspenseInstance,
+  callback: () => void,
+) {
+  instance._reactRetry = callback;
+}
+
 export function getNextHydratableSibling(
   instance: HydratableInstance,
 ): null | HydratableInstance {
@@ -565,7 +586,9 @@ export function getNextHydratableSibling(
     node.nodeType !== TEXT_NODE &&
     (!enableSuspenseServerRenderer ||
       node.nodeType !== COMMENT_NODE ||
-      (node: any).data !== SUSPENSE_START_DATA)
+      ((node: any).data !== SUSPENSE_START_DATA &&
+        (node: any).data !== SUSPENSE_PENDING_START_DATA &&
+        (node: any).data !== SUSPENSE_FALLBACK_START_DATA))
   ) {
     node = node.nextSibling;
   }
@@ -583,7 +606,9 @@ export function getFirstHydratableChild(
     next.nodeType !== TEXT_NODE &&
     (!enableSuspenseServerRenderer ||
       next.nodeType !== COMMENT_NODE ||
-      (next: any).data !== SUSPENSE_START_DATA)
+      ((next: any).data !== SUSPENSE_START_DATA &&
+        (next: any).data !== SUSPENSE_FALLBACK_START_DATA &&
+        (next: any).data !== SUSPENSE_PENDING_START_DATA))
   ) {
     next = next.nextSibling;
   }

packages/react-dom/src/server/ReactPartialRenderer.js

@@ -835,6 +835,7 @@ class ReactDOMServerRenderer {
                 'suspense fallback not found, something is broken',
               );
               this.stack.push(fallbackFrame);
+              out[this.suspenseDepth] += '<!--$!-->';
               // Skip flushing output since we're switching to the fallback
               continue;
             } else {
@@ -996,8 +997,7 @@ class ReactDOMServerRenderer {
               children: fallbackChildren,
               childIndex: 0,
               context: context,
-              footer: '',
-              out: '',
+              footer: '<!--/$-->',
             };
             const frame: Frame = {
               fallbackFrame,

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -84,7 +84,11 @@ import {
 import {
   shouldSetTextContent,
   shouldDeprioritizeSubtree,
+  isSuspenseInstancePending,
+  isSuspenseInstanceFallback,
+  registerSuspenseInstanceRetry,
 } from './ReactFiberHostConfig';
+import type {SuspenseInstance} from './ReactFiberHostConfig';
 import {pushHostContext, pushHostContainer} from './ReactFiberHostContext';
 import {
   pushProvider,
@@ -129,6 +133,7 @@ import {
   createWorkInProgress,
   isSimpleFunctionComponent,
 } from './ReactFiber';
+import {retryTimedOutBoundary} from './ReactFiberScheduler';
 
 const ReactCurrentOwner = ReactSharedInternals.ReactCurrentOwner;
 
@@ -1637,10 +1642,21 @@ function updateDehydratedSuspenseComponent(
     workInProgress.expirationTime = Never;
     return null;
   }
+  if ((workInProgress.effectTag & DidCapture) !== NoEffect) {
+    // 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;
+  }
   // We use childExpirationTime to indicate that a child might depend on context, so if
   // any context has changed, we need to treat is as if the input might have changed.
   const hasContextChanged = current.childExpirationTime >= renderExpirationTime;
-  if (didReceiveUpdate || hasContextChanged) {
+  const suspenseInstance = (current.stateNode: SuspenseInstance);
+  if (
+    didReceiveUpdate ||
+    hasContextChanged ||
+    isSuspenseInstanceFallback(suspenseInstance)
+  ) {
     // This boundary has changed since the first render. This means that we are now unable to
     // hydrate it. We might still be able to hydrate it using an earlier expiration time but
     // during this render we can't. Instead, we're going to delete the whole subtree and
@@ -1648,6 +1664,10 @@ function updateDehydratedSuspenseComponent(
     // or fallback. The real Suspense boundary will suspend for a while so we have some time
     // to ensure it can produce real content, but all state and pending events will be lost.
 
+    // Alternatively, this boundary is in a permanent fallback state. In this case, we'll never
+    // get an update and we'll never be able to hydrate the final content. Let's just try the
+    // client side render instead.
+
     // Detach from the current dehydrated boundary.
     current.alternate = null;
     workInProgress.alternate = null;
@@ -1677,8 +1697,26 @@ function updateDehydratedSuspenseComponent(
     workInProgress.effectTag |= Placement;
     // Retry as a real Suspense component.
     return updateSuspenseComponent(null, workInProgress, renderExpirationTime);
-  }
-  if ((workInProgress.effectTag & DidCapture) === NoEffect) {
+  } else if (isSuspenseInstancePending(suspenseInstance)) {
+    // This component is still pending more data from the server, so we can't hydrate its
+    // content. We treat it as if this component suspended itself. It might seem as if
+    // we could just try to render it client-side instead. However, this will perform a
+    // lot of unnecessary work and is unlikely to complete since it often will suspend
+    // on missing data anyway. Additionally, the server might be able to render more
+    // than we can on the client yet. In that case we'd end up with more fallback states
+    // on the client than if we just leave it alone. If the server times out or errors
+    // these should update this boundary to the permanent Fallback state instead.
+    // Mark it as having captured (i.e. suspended).
+    workInProgress.effectTag |= DidCapture;
+    // Leave the children in place. I.e. empty.
+    workInProgress.child = null;
+    // Register a callback to retry this boundary once the server has sent the result.
+    registerSuspenseInstanceRetry(
+      suspenseInstance,
+      retryTimedOutBoundary.bind(null, current),
+    );
+    return null;
+  } else {
     // This is the first attempt.
     reenterHydrationStateFromDehydratedSuspenseInstance(workInProgress);
     const nextProps = workInProgress.pendingProps;
@@ -1690,11 +1728,6 @@ function updateDehydratedSuspenseComponent(
       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;
   }
 }
 

packages/react-reconciler/src/ReactFiberCommitWork.js

@@ -94,7 +94,7 @@ import {
 import {
   captureCommitPhaseError,
   requestCurrentTime,
-  retryTimedOutBoundary,
+  resolveRetryThenable,
 } from './ReactFiberScheduler';
 import {
   NoEffect as NoHookEffect,
@@ -1232,7 +1232,7 @@ function commitWork(current: Fiber | null, finishedWork: Fiber): void {
         }
         thenables.forEach(thenable => {
           // Memoize using the boundary fiber to prevent redundant listeners.
-          let retry = retryTimedOutBoundary.bind(null, finishedWork, thenable);
+          let retry = resolveRetryThenable.bind(null, finishedWork, thenable);
           if (enableSchedulerTracing) {
             retry = Schedule_tracing_wrap(retry);
           }

packages/react-reconciler/src/ReactFiberCompleteWork.js

@@ -710,7 +710,12 @@ function completeWork(
       const nextDidTimeout = nextState !== null;
       const prevDidTimeout = current !== null && current.memoizedState !== null;
 
-      if (current !== null && !nextDidTimeout && prevDidTimeout) {
+      if (current === null) {
+        // In cases where we didn't find a suitable hydration boundary we never
+        // downgraded this to a DehydratedSuspenseComponent, but we still need to
+        // pop the hydration state since we might be inside the insertion tree.
+        popHydrationState(workInProgress);
+      } else if (!nextDidTimeout && prevDidTimeout) {
         // We just switched from the fallback to the normal children. Delete
         // the fallback.
         // TODO: Would it be better to store the fallback fragment on

packages/react-reconciler/src/ReactFiberScheduler.js

@@ -1698,7 +1698,20 @@ function pingSuspendedRoot(
   }
 }
 
-function retryTimedOutBoundary(boundaryFiber: Fiber, thenable: Thenable) {
+function retryTimedOutBoundary(boundaryFiber: Fiber) {
+  const currentTime = requestCurrentTime();
+  const retryTime = computeExpirationForFiber(currentTime, boundaryFiber);
+  const root = scheduleWorkToRoot(boundaryFiber, retryTime);
+  if (root !== null) {
+    markPendingPriorityLevel(root, retryTime);
+    const rootExpirationTime = root.expirationTime;
+    if (rootExpirationTime !== NoWork) {
+      requestWork(root, rootExpirationTime);
+    }
+  }
+}
+
+function resolveRetryThenable(boundaryFiber: Fiber, thenable: Thenable) {
   // The boundary fiber (a Suspense component) previously timed out and was
   // rendered in its fallback state. One of the promises that suspended it has
   // resolved, which means at least part of the tree was likely unblocked. Try
@@ -1729,16 +1742,7 @@ function retryTimedOutBoundary(boundaryFiber: Fiber, thenable: Thenable) {
     retryCache.delete(thenable);
   }
 
-  const currentTime = requestCurrentTime();
-  const retryTime = computeExpirationForFiber(currentTime, boundaryFiber);
-  const root = scheduleWorkToRoot(boundaryFiber, retryTime);
-  if (root !== null) {
-    markPendingPriorityLevel(root, retryTime);
-    const rootExpirationTime = root.expirationTime;
-    if (rootExpirationTime !== NoWork) {
-      requestWork(root, rootExpirationTime);
-    }
-  }
+  retryTimedOutBoundary(boundaryFiber);
 }
 
 function scheduleWorkToRoot(fiber: Fiber, expirationTime): FiberRoot | null {
@@ -2589,6 +2593,7 @@ export {
   renderDidError,
   pingSuspendedRoot,
   retryTimedOutBoundary,
+  resolveRetryThenable,
   markLegacyErrorBoundaryAsFailed,
   isAlreadyFailedLegacyErrorBoundary,
   scheduleWork,

packages/react-reconciler/src/ReactFiberUnwindWork.js

@@ -66,7 +66,7 @@ import {
   markLegacyErrorBoundaryAsFailed,
   isAlreadyFailedLegacyErrorBoundary,
   pingSuspendedRoot,
-  retryTimedOutBoundary,
+  resolveRetryThenable,
 } from './ReactFiberScheduler';
 
 import invariant from 'shared/invariant';
@@ -372,11 +372,7 @@ function throwException(
         // Memoize using the boundary fiber to prevent redundant listeners.
         if (!retryCache.has(thenable)) {
           retryCache.add(thenable);
-          let retry = retryTimedOutBoundary.bind(
-            null,
-            workInProgress,
-            thenable,
-          );
+          let retry = resolveRetryThenable.bind(null, workInProgress, thenable);
           if (enableSchedulerTracing) {
             retry = Schedule_tracing_wrap(retry);
           }

packages/react-reconciler/src/forks/ReactFiberHostConfig.custom.js

@@ -107,6 +107,12 @@ export const canHydrateInstance = $$$hostConfig.canHydrateInstance;
 export const canHydrateTextInstance = $$$hostConfig.canHydrateTextInstance;
 export const canHydrateSuspenseInstance =
   $$$hostConfig.canHydrateSuspenseInstance;
+export const isSuspenseInstancePending =
+  $$$hostConfig.isSuspenseInstancePending;
+export const isSuspenseInstanceFallback =
+  $$$hostConfig.isSuspenseInstanceFallback;
+export const registerSuspenseInstanceRetry =
+  $$$hostConfig.registerSuspenseInstanceRetry;
 export const getNextHydratableSibling = $$$hostConfig.getNextHydratableSibling;
 export const getFirstHydratableChild = $$$hostConfig.getFirstHydratableChild;
 export const hydrateInstance = $$$hostConfig.hydrateInstance;

packages/shared/HostConfigWithNoHydration.js

@@ -27,6 +27,9 @@ export const supportsHydration = false;
 export const canHydrateInstance = shim;
 export const canHydrateTextInstance = shim;
 export const canHydrateSuspenseInstance = shim;
+export const isSuspenseInstancePending = shim;
+export const isSuspenseInstanceFallback = shim;
+export const registerSuspenseInstanceRetry = shim;
 export const getNextHydratableSibling = shim;
 export const getFirstHydratableChild = shim;
 export const hydrateInstance = shim;