From 9fc19048bdda0248f96f8a5971b8f6b2e488e346 Mon Sep 17 00:00:00 2001
From: Shu Ding <g@shud.in>
Date: Wed, 30 Mar 2022 15:15:13 +0200
Subject: [PATCH] Ensure the app shell is rendered before rendering the
 document (#35732)

---
 .../next/server/node-web-streams-helper.ts    | 50 +++++++++++++++++++
 packages/next/server/render.tsx               | 18 +++++--
 2 files changed, 64 insertions(+), 4 deletions(-)

diff --git a/packages/next/server/node-web-streams-helper.ts b/packages/next/server/node-web-streams-helper.ts
index b8b1cc6ff98d..f854b4e675f3 100644
--- a/packages/next/server/node-web-streams-helper.ts
+++ b/packages/next/server/node-web-streams-helper.ts
@@ -233,6 +233,56 @@ export function createFlushEffectStream(
   })
 }
 
+export async function renderToInitialStream({
+  ReactDOMServer,
+  element,
+}: {
+  ReactDOMServer: typeof import('react-dom/server')
+  element: React.ReactElement
+}): Promise<
+  ReadableStream<Uint8Array> & {
+    allReady?: Promise<void>
+  }
+> {
+  return await (ReactDOMServer as any).renderToReadableStream(element)
+}
+
+export async function continueFromInitialStream({
+  suffix,
+  dataStream,
+  generateStaticHTML,
+  flushEffectHandler,
+  renderStream,
+}: {
+  suffix?: string
+  dataStream?: ReadableStream<Uint8Array>
+  generateStaticHTML: boolean
+  flushEffectHandler?: () => Promise<string>
+  renderStream: ReadableStream<Uint8Array> & {
+    allReady?: Promise<void>
+  }
+}): Promise<ReadableStream<Uint8Array>> {
+  const closeTag = '</body></html>'
+  const suffixUnclosed = suffix ? suffix.split(closeTag)[0] : null
+
+  if (generateStaticHTML) {
+    await renderStream.allReady
+  }
+
+  const transforms: Array<TransformStream<Uint8Array, Uint8Array>> = [
+    createBufferedTransformStream(),
+    flushEffectHandler ? createFlushEffectStream(flushEffectHandler) : null,
+    suffixUnclosed != null ? createPrefixStream(suffixUnclosed) : null,
+    dataStream ? createInlineDataStream(dataStream) : null,
+    suffixUnclosed != null ? createSuffixStream(closeTag) : null,
+  ].filter(Boolean) as any
+
+  return transforms.reduce(
+    (readable, transform) => pipeThrough(readable, transform),
+    renderStream
+  )
+}
+
 export async function renderToStream({
   ReactDOMServer,
   element,
diff --git a/packages/next/server/render.tsx b/packages/next/server/render.tsx
index 2b5f36910aef..fd045c7dcb6f 100644
--- a/packages/next/server/render.tsx
+++ b/packages/next/server/render.tsx
@@ -74,6 +74,8 @@ import {
   chainStreams,
   createBufferedTransformStream,
   renderToStream,
+  renderToInitialStream,
+  continueFromInitialStream,
 } from './node-web-streams-helper'
 import { ImageConfigContext } from '../shared/lib/image-config-context'
 import { FlushEffectsContext } from '../shared/lib/flush-effects'
@@ -1354,11 +1356,20 @@ export async function renderToHTML(
       }
 
       if (hasConcurrentFeatures) {
+        let renderStream: any
+
+        // We start rendering the shell earlier, before returning the head tags
+        // to `documentResult`.
+        const content = renderContent()
+        renderStream = await renderToInitialStream({
+          ReactDOMServer,
+          element: content,
+        })
+
         bodyResult = async (suffix: string) => {
           // this must be called inside bodyResult so appWrappers is
           // up to date when getWrappedApp is called
 
-          const content = renderContent()
           const flushEffectHandler = async () => {
             const allFlushEffects = [
               styledJsxFlushEffect,
@@ -1379,9 +1390,8 @@ export async function renderToHTML(
             return flushed
           }
 
-          return await renderToStream({
-            ReactDOMServer,
-            element: content,
+          return await continueFromInitialStream({
+            renderStream,
             suffix,
             dataStream: serverComponentsInlinedTransformStream?.readable,
             generateStaticHTML: generateStaticHTML || !hasConcurrentFeatures,