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,