Deal with fallback content in Partial Hydration

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,