implement preamble and postambl for react-dom/server

gnoff · gnoff · commit 88402c6bd57b · 2022-08-09T08:47:23.000-07:00
diff --git a/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js b/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js
@@ -128,6 +128,8 @@ describe('ReactDOMFizzServer', () => {
     buffer = '';
     const fakeBody = document.createElement('body');
     fakeBody.innerHTML = bufferedContent;
+    const parent =
+      container.nodeName === '#document' ? container.body : container;
     while (fakeBody.firstChild) {
       const node = fakeBody.firstChild;
       if (
@@ -137,13 +139,37 @@ describe('ReactDOMFizzServer', () => {
         const script = document.createElement('script');
         script.textContent = node.textContent;
         fakeBody.removeChild(node);
-        container.appendChild(script);
+        parent.appendChild(script);
       } else {
-        container.appendChild(node);
+        parent.appendChild(node);
       }
     }
   }
 
+  async function actIntoEmptyDocument(callback) {
+    await callback();
+    // Await one turn around the event loop.
+    // This assumes that we'll flush everything we have so far.
+    await new Promise(resolve => {
+      setImmediate(resolve);
+    });
+    if (hasErrored) {
+      throw fatalError;
+    }
+    // JSDOM doesn't support stream HTML parser so we need to give it a proper fragment.
+    // We also want to execute any scripts that are embedded.
+    // We assume that we have now received a proper fragment of HTML.
+    const bufferedContent = buffer;
+    // Test Environment
+    const jsdom = new JSDOM(bufferedContent, {
+      runScripts: 'dangerously',
+    });
+    window = jsdom.window;
+    document = jsdom.window.document;
+    container = document;
+    buffer = '';
+  }
+
   function getVisibleChildren(element) {
     const children = [];
     let node = element.firstChild;
@@ -4194,6 +4220,113 @@ describe('ReactDOMFizzServer', () => {
     );
   });
 
+  it('emits html and head start tags (the preamble) before other content if rendered in the shell', async () => {
+    await actIntoEmptyDocument(() => {
+      const {pipe} = ReactDOMFizzServer.renderToPipeableStream(
+        <>
+          <title data-baz="baz">a title</title>
+          <html data-foo="foo">
+            <head data-bar="bar" />
+            <body>a body</body>
+          </html>
+        </>,
+      );
+      pipe(writable);
+    });
+    expect(getVisibleChildren(document)).toEqual(
+      <html data-foo="foo">
+        <head data-bar="bar">
+          <title data-baz="baz">a title</title>
+        </head>
+        <body>a body</body>
+      </html>,
+    );
+
+    // Hydrate the same thing on the client. We expect this to still fail because <title> is not a Resource
+    // and is unmatched on hydration
+    const errors = [];
+    const root = ReactDOMClient.hydrateRoot(
+      document,
+      <>
+        <title data-baz="baz">a title</title>
+        <html data-foo="foo">
+          <head data-bar="bar" />
+          <body>a body</body>
+        </html>
+      </>,
+      {
+        onRecoverableError: (err, errInfo) => {
+          errors.push(err.message);
+        },
+      },
+    );
+    expect(() => {
+      try {
+        expect(() => {
+          expect(Scheduler).toFlushWithoutYielding();
+        }).toThrow('Invalid insertion of HTML node in #document node.');
+      } catch (e) {
+        console.log('e', e);
+      }
+    }).toErrorDev(
+      [
+        'Warning: Expected server HTML to contain a matching <title> in <#document>.',
+        'Warning: An error occurred during hydration. The server HTML was replaced with client content in <#document>.',
+        'Warning: validateDOMNesting(...): <title> cannot appear as a child of <#document>',
+      ],
+      {withoutStack: 1},
+    );
+    expect(errors).toEqual([
+      'Hydration failed because the initial UI does not match what was rendered on the server.',
+      'There was an error while hydrating. Because the error happened outside of a Suspense boundary, the entire root will switch to client rendering.',
+    ]);
+  });
+
+  it('holds back body and html closing tags (the postamble) until all pending tasks are completed', async () => {
+    const chunks = [];
+    writable.on('data', chunk => {
+      chunks.push(chunk);
+    });
+
+    await actIntoEmptyDocument(() => {
+      const {pipe} = ReactDOMFizzServer.renderToPipeableStream(
+        <html>
+          <head />
+          <body>
+            first
+            <Suspense>
+              <AsyncText text="second" />
+            </Suspense>
+          </body>
+        </html>,
+      );
+      pipe(writable);
+    });
+
+    expect(getVisibleChildren(document)).toEqual(
+      <html>
+        <head />
+        <body>{'first'}</body>
+      </html>,
+    );
+
+    await act(() => {
+      resolveText('second');
+    });
+
+    expect(getVisibleChildren(document)).toEqual(
+      <html>
+        <head />
+        <body>
+          {'first'}
+          {'second'}
+        </body>
+      </html>,
+    );
+
+    expect(chunks.pop()).toEqual('</body></html>');
+  });
+
   describe('text separators', () => {
     // To force performWork to start before resolving AsyncText but before piping we need to wait until
     // after scheduleWork which currently uses setImmediate to delay performWork
diff --git a/packages/react-dom/src/server/ReactDOMServerFormatConfig.js b/packages/react-dom/src/server/ReactDOMServerFormatConfig.js
@@ -242,6 +242,26 @@ export function getChildFormatContext(
   return parentContext;
 }
 
+export function isPreambleInsertion(type: string): boolean {
+  switch (type) {
+    case 'html':
+    case 'head': {
+      return true;
+    }
+  }
+  return false;
+}
+
+export function isPostambleInsertion(type: string): boolean {
+  switch (type) {
+    case 'body':
+    case 'html': {
+      return true;
+    }
+  }
+  return false;
+}
+
 export type SuspenseBoundaryID = null | PrecomputedChunk;
 
 export const UNINITIALIZED_SUSPENSE_BOUNDARY_ID: SuspenseBoundaryID = null;
@@ -1405,11 +1425,13 @@ const DOCTYPE: PrecomputedChunk = stringToPrecomputedChunk('<!DOCTYPE html>');
 
 export function pushStartInstance(
   target: Array<Chunk | PrecomputedChunk>,
+  preamble: Array<Chunk | PrecomputedChunk>,
   type: string,
   props: Object,
   responseState: ResponseState,
   formatContext: FormatContext,
 ): ReactNodeList {
+  target = isPreambleInsertion(type) ? preamble : target;
   if (__DEV__) {
     validateARIAProperties(type, props);
     validateInputProperties(type, props);
@@ -1521,9 +1543,11 @@ const endTag2 = stringToPrecomputedChunk('>');
 
 export function pushEndInstance(
   target: Array<Chunk | PrecomputedChunk>,
+  postamble: Array<Chunk | PrecomputedChunk>,
   type: string,
   props: Object,
 ): void {
+  target = isPostambleInsertion(type) ? postamble : target;
   switch (type) {
     // Omitted close tags
     // TODO: Instead of repeating this switch we could try to pass a flag from above.
diff --git a/packages/react-noop-renderer/src/ReactNoopServer.js b/packages/react-noop-renderer/src/ReactNoopServer.js
@@ -113,6 +113,7 @@ const ReactNoopServer = ReactFizzServer({
   },
   pushStartInstance(
     target: Array<Uint8Array>,
+    preamble: Array<Uint8Array>,
     type: string,
     props: Object,
   ): ReactNodeList {
@@ -128,6 +129,7 @@ const ReactNoopServer = ReactFizzServer({
 
   pushEndInstance(
     target: Array<Uint8Array>,
+    postamble: Array<Uint8Array>,
     type: string,
     props: Object,
   ): void {
diff --git a/packages/react-server/src/ReactFizzServer.js b/packages/react-server/src/ReactFizzServer.js
@@ -200,6 +200,8 @@ export opaque type Request = {
   clientRenderedBoundaries: Array<SuspenseBoundary>, // Errored or client rendered but not yet flushed.
   completedBoundaries: Array<SuspenseBoundary>, // Completed but not yet fully flushed boundaries to show.
   partialBoundaries: Array<SuspenseBoundary>, // Partially completed boundaries that can flush its segments early.
+  +preamble: Array<Chunk | PrecomputedChunk>, // Chunks that need to be emitted before any segment chunks.
+  +postamble: Array<Chunk | PrecomputedChunk>, // Chunks that need to be emitted after segments, waiting for all pending root tasks to finish
   // onError is called when an error happens anywhere in the tree. It might recover.
   // The return string is used in production  primarily to avoid leaking internals, secondarily to save bytes.
   // Returning null/undefined will cause a defualt error message in production
@@ -272,6 +274,8 @@ export function createRequest(
     clientRenderedBoundaries: [],
     completedBoundaries: [],
     partialBoundaries: [],
+    preamble: [],
+    postamble: [],
     onError: onError === undefined ? defaultErrorHandler : onError,
     onAllReady: onAllReady === undefined ? noop : onAllReady,
     onShellReady: onShellReady === undefined ? noop : onShellReady,
@@ -632,6 +636,7 @@ function renderHostElement(
   const segment = task.blockedSegment;
   const children = pushStartInstance(
     segment.chunks,
+    request.preamble,
     type,
     props,
     request.responseState,
@@ -647,7 +652,7 @@ function renderHostElement(
   // We expect that errors will fatal the whole task and that we don't need
   // the correct context. Therefore this is not in a finally.
   segment.formatContext = prevContext;
-  pushEndInstance(segment.chunks, type, props);
+  pushEndInstance(segment.chunks, request.postamble, type, props);
   segment.lastPushedText = false;
   popComponentStackInDEV(task);
 }
@@ -2054,6 +2059,7 @@ function flushCompletedQueues(
   request: Request,
   destination: Destination,
 ): void {
+  let allComplete = false;
   beginWriting(destination);
   try {
     // The structure of this is to go through each queue one by one and write
@@ -2063,20 +2069,30 @@ function flushCompletedQueues(
 
     // TODO: Emit preloading.
 
-    // TODO: It's kind of unfortunate to keep checking this array after we've already
-    // emitted the root.
+    let i;
     const completedRootSegment = request.completedRootSegment;
-    if (completedRootSegment !== null && request.pendingRootTasks === 0) {
-      flushSegment(request, destination, completedRootSegment);
-      request.completedRootSegment = null;
-      writeCompletedRoot(destination, request.responseState);
+    if (completedRootSegment !== null) {
+      if (request.pendingRootTasks === 0) {
+        const preamble = request.preamble;
+        for (i = 0; i < preamble.length; i++) {
+          // we expect the preamble to be tiny and will ignore backpressure
+          writeChunk(destination, preamble[i]);
+        }
+        preamble.length = 0;
+
+        flushSegment(request, destination, completedRootSegment);
+        request.completedRootSegment = null;
+        writeCompletedRoot(destination, request.responseState);
+      } else {
+        // We haven't flushed the root yet so we don't need to check boundaries further down
+        return;
+      }
     }
 
     // We emit client rendering instructions for already emitted boundaries first.
     // This is so that we can signal to the client to start client rendering them as
     // soon as possible.
     const clientRenderedBoundaries = request.clientRenderedBoundaries;
-    let i;
     for (i = 0; i < clientRenderedBoundaries.length; i++) {
       const boundary = clientRenderedBoundaries[i];
       if (!flushClientRenderedBoundary(request, destination, boundary)) {
@@ -2138,9 +2154,7 @@ function flushCompletedQueues(
       }
     }
     largeBoundaries.splice(0, i);
-  } finally {
-    completeWriting(destination);
-    flushBuffered(destination);
+
     if (
       request.allPendingTasks === 0 &&
       request.pingedTasks.length === 0 &&
@@ -2149,6 +2163,17 @@ function flushCompletedQueues(
       // We don't need to check any partially completed segments because
       // either they have pending task or they're complete.
     ) {
+      allComplete = true;
+      const postamble = request.postamble;
+      for (i = 0; i < postamble.length; i++) {
+        writeChunk(destination, postamble[i]);
+      }
+      postamble.length = 0;
+    }
+  } finally {
+    completeWriting(destination);
+    flushBuffered(destination);
+    if (allComplete) {
       if (__DEV__) {
         if (request.abortableTasks.size !== 0) {
           console.error(
