Wait for shell resolve with gIP is customized in react 18 (#36792)

When getInitialProps is customized with react 18, since gIP requires to return `html` as doc property which could be used by  user-land customization, we do blocking-rendering there and passdown the `html` to document

Fixes #36675
Closes #36419

packages/next/server/node-web-streams-helper.ts

@@ -1,5 +1,9 @@
 import { nonNullable } from '../lib/non-nullable'
 
+export type ReactReadableStream = ReadableStream<Uint8Array> & {
+  allReady?: Promise<void> | undefined
+}
+
 export function readableStreamTee<T = any>(
   readable: ReadableStream<T>
 ): [ReadableStream<T>, ReadableStream<T>] {
@@ -138,29 +142,24 @@ export function renderToInitialStream({
 }: {
   ReactDOMServer: any
   element: React.ReactElement
-}): Promise<
-  ReadableStream<Uint8Array> & {
-    allReady?: Promise<void>
-  }
-> {
+}): Promise<ReactReadableStream> {
   return ReactDOMServer.renderToReadableStream(element)
 }
 
-export async function continueFromInitialStream({
-  suffix,
-  dataStream,
-  generateStaticHTML,
-  flushEffectHandler,
-  renderStream,
-}: {
-  suffix?: string
-  dataStream?: ReadableStream<Uint8Array>
-  generateStaticHTML: boolean
-  flushEffectHandler?: () => string
-  renderStream: ReadableStream<Uint8Array> & {
-    allReady?: Promise<void>
+export async function continueFromInitialStream(
+  renderStream: ReactReadableStream,
+  {
+    suffix,
+    dataStream,
+    generateStaticHTML,
+    flushEffectHandler,
+  }: {
+    suffix?: string
+    dataStream?: ReadableStream<Uint8Array>
+    generateStaticHTML: boolean
+    flushEffectHandler?: () => string
   }
-}): Promise<ReadableStream<Uint8Array>> {
+): Promise<ReadableStream<Uint8Array>> {
   const closeTag = '</body></html>'
   const suffixUnclosed = suffix ? suffix.split(closeTag)[0] : null
 
@@ -198,12 +197,11 @@ export async function renderToStream({
   flushEffectHandler?: () => string
 }): Promise<ReadableStream<Uint8Array>> {
   const renderStream = await renderToInitialStream({ ReactDOMServer, element })
-  return continueFromInitialStream({
+  return continueFromInitialStream(renderStream, {
     suffix,
     dataStream,
     generateStaticHTML,
     flushEffectHandler,
-    renderStream,
   })
 }
 

packages/next/server/render.tsx

@@ -20,6 +20,7 @@ import type { FontManifest } from './font-utils'
 import type { LoadComponentsReturnType, ManifestItem } from './load-components'
 import type { GetServerSideProps, GetStaticProps, PreviewData } from '../types'
 import type { UnwrapPromise } from '../lib/coalesced-function'
+import type { ReactReadableStream } from './node-web-streams-helper'
 
 import React from 'react'
 import { createFromReadableStream } from 'next/dist/compiled/react-server-dom-webpack'
@@ -1340,11 +1341,11 @@ export async function renderToHTML(
       }
     }
 
-    async function documentInitialProps(
+    async function loadDocumentInitialProps(
       renderShell?: (
         _App: AppType,
         _Component: NextComponentType
-      ) => Promise<void>
+      ) => Promise<ReactReadableStream>
     ) {
       const renderPage: RenderPage = (
         options: ComponentsEnhancer = {}
@@ -1373,11 +1374,13 @@ export async function renderToHTML(
           enhanceComponents(options, App, Component)
 
         if (renderShell) {
-          return renderShell(EnhancedApp, EnhancedComponent).then(() => {
-            // When using concurrent features, we don't have or need the full
-            // html so it's fine to return nothing here.
-            return { html: '', head }
-          })
+          return renderShell(EnhancedApp, EnhancedComponent).then(
+            async (stream) => {
+              const forwardStream = readableStreamTee(stream)[1]
+              const html = await streamToString(forwardStream)
+              return { html, head }
+            }
+          )
         }
 
         const html = ReactDOMServer.renderToString(
@@ -1438,9 +1441,9 @@ export async function renderToHTML(
     if (!process.env.__NEXT_REACT_ROOT) {
       // Enabling react legacy rendering mode: __NEXT_REACT_ROOT = false
       if (Document.getInitialProps) {
-        const documentInitialPropsRes = await documentInitialProps()
-        if (documentInitialPropsRes === null) return null
-        const { docProps, documentCtx } = documentInitialPropsRes
+        const documentInitialProps = await loadDocumentInitialProps()
+        if (documentInitialProps === null) return null
+        const { docProps, documentCtx } = documentInitialProps
 
         return {
           bodyResult: (suffix: string) =>
@@ -1473,97 +1476,97 @@ export async function renderToHTML(
       }
     } else {
       // Enabling react concurrent rendering mode: __NEXT_REACT_ROOT = true
-      let renderStream: ReadableStream<Uint8Array> & {
-        allReady?: Promise<void> | undefined
-      }
-
       const renderShell = async (
         EnhancedApp: AppType,
         EnhancedComponent: NextComponentType
       ) => {
         const content = renderContent(EnhancedApp, EnhancedComponent)
-        renderStream = await renderToInitialStream({
+        return await renderToInitialStream({
           ReactDOMServer,
           element: content,
         })
       }
 
-      const bodyResult = async (suffix: string) => {
-        // this must be called inside bodyResult so appWrappers is
-        // up to date when `wrapApp` is called
-
-        const flushEffectHandler = (): string => {
-          const allFlushEffects = [
-            styledJsxFlushEffect,
-            ...(flushEffects || []),
-          ]
-          const flushed = ReactDOMServer.renderToString(
-            <>
-              {allFlushEffects.map((flushEffect, i) => (
-                <React.Fragment key={i}>{flushEffect()}</React.Fragment>
-              ))}
-            </>
-          )
-          return flushed
-        }
+      const createBodyResult =
+        (initialStream: ReactReadableStream) => (suffix: string) => {
+          // this must be called inside bodyResult so appWrappers is
+          // up to date when `wrapApp` is called
+          const flushEffectHandler = (): string => {
+            const allFlushEffects = [
+              styledJsxFlushEffect,
+              ...(flushEffects || []),
+            ]
+            const flushed = ReactDOMServer.renderToString(
+              <>
+                {allFlushEffects.map((flushEffect, i) => (
+                  <React.Fragment key={i}>{flushEffect()}</React.Fragment>
+                ))}
+              </>
+            )
+            return flushed
+          }
 
-        // Handle static data for server components.
-        async function generateStaticFlightDataIfNeeded() {
-          if (serverComponentsPageDataTransformStream) {
-            // If it's a server component with the Node.js runtime, we also
-            // statically generate the page data.
-            let data = ''
-
-            const readable = serverComponentsPageDataTransformStream.readable
-            const reader = readable.getReader()
-            const textDecoder = new TextDecoder()
-
-            while (true) {
-              const { done, value } = await reader.read()
-              if (done) {
-                break
+          // Handle static data for server components.
+          async function generateStaticFlightDataIfNeeded() {
+            if (serverComponentsPageDataTransformStream) {
+              // If it's a server component with the Node.js runtime, we also
+              // statically generate the page data.
+              let data = ''
+
+              const readable = serverComponentsPageDataTransformStream.readable
+              const reader = readable.getReader()
+              const textDecoder = new TextDecoder()
+
+              while (true) {
+                const { done, value } = await reader.read()
+                if (done) {
+                  break
+                }
+                data += decodeText(value, textDecoder)
               }
-              data += decodeText(value, textDecoder)
-            }
 
-            ;(renderOpts as any).pageData = {
-              ...(renderOpts as any).pageData,
-              __flight__: data,
+              ;(renderOpts as any).pageData = {
+                ...(renderOpts as any).pageData,
+                __flight__: data,
+              }
+              return data
             }
-            return data
           }
-        }
 
-        // @TODO: A potential improvement would be to reuse the inlined
-        // data stream, or pass a callback inside as this doesn't need to
-        // be streamed.
-        // Do not use `await` here.
-        generateStaticFlightDataIfNeeded()
-        return await continueFromInitialStream({
-          renderStream,
-          suffix,
-          dataStream: serverComponentsInlinedTransformStream?.readable,
-          generateStaticHTML,
-          flushEffectHandler,
-        })
-      }
+          // @TODO: A potential improvement would be to reuse the inlined
+          // data stream, or pass a callback inside as this doesn't need to
+          // be streamed.
+          // Do not use `await` here.
+          generateStaticFlightDataIfNeeded()
+          return continueFromInitialStream(initialStream, {
+            suffix,
+            dataStream: serverComponentsInlinedTransformStream?.readable,
+            generateStaticHTML,
+            flushEffectHandler,
+          })
+        }
 
       const hasDocumentGetInitialProps = !(
         isServerComponent ||
         process.env.NEXT_RUNTIME === 'edge' ||
         !Document.getInitialProps
       )
 
+      let bodyResult: (s: string) => Promise<ReadableStream<Uint8Array>>
+
       // If it has getInitialProps, we will render the shell in `renderPage`.
       // Otherwise we do it right now.
       let documentInitialPropsRes:
         | {}
-        | Awaited<ReturnType<typeof documentInitialProps>>
+        | Awaited<ReturnType<typeof loadDocumentInitialProps>>
       if (hasDocumentGetInitialProps) {
-        documentInitialPropsRes = await documentInitialProps(renderShell)
+        documentInitialPropsRes = await loadDocumentInitialProps(renderShell)
         if (documentInitialPropsRes === null) return null
+        const { docProps } = documentInitialPropsRes as any
+        bodyResult = createBodyResult(streamFromArray([docProps.html]))
       } else {
-        await renderShell(App, Component)
+        const stream = await renderShell(App, Component)
+        bodyResult = createBodyResult(stream)
         documentInitialPropsRes = {}
       }
 

packages/next/server/view-render.tsx

@@ -506,8 +506,7 @@ export async function renderToHTML(
     // Do not use `await` here.
     // generateStaticFlightDataIfNeeded()
 
-    return await continueFromInitialStream({
-      renderStream,
+    return await continueFromInitialStream(renderStream, {
       suffix: '',
       dataStream: serverComponentsInlinedTransformStream?.readable,
       generateStaticHTML: generateStaticHTML || !hasConcurrentFeatures,

test/integration/amphtml-fragment-style/test/index.test.js

@@ -12,19 +12,14 @@ import {
 } from 'next-test-utils'
 
 const appDir = join(__dirname, '../')
-const nodeArgs = ['-r', join(appDir, '../../lib/react-17-require-hook.js')]
 let appPort
 let app
 
 describe('AMP Fragment Styles', () => {
   beforeAll(async () => {
-    await nextBuild(appDir, [], {
-      nodeArgs,
-    })
+    await nextBuild(appDir, [])
     appPort = await findPort()
-    app = await nextStart(appDir, appPort, {
-      nodeArgs,
-    })
+    app = await nextStart(appDir, appPort)
   })
   afterAll(() => killApp(app))
 

test/integration/app-document/pages/_document.js

@@ -42,6 +42,7 @@ export default class MyDocument extends Document {
 
     return {
       ...result,
+      cssInJsCount: (result.html.match(/css-in-js-class/g) || []).length,
       customProperty: 'Hello Document',
       withCSP: ctx.query.withCSP,
     }
@@ -74,6 +75,7 @@ export default class MyDocument extends Document {
           <p id="document-hmr">Hello Document HMR</p>
           <Main />
           <NextScript nonce="test-nonce" />
+          <div id="css-in-cjs-count">{this.props.cssInJsCount}</div>
         </body>
       </Html>
     )

test/integration/app-document/pages/index.js

@@ -2,8 +2,10 @@ import Link from 'next/link'
 export default () => (
   <div>
     <div className="page-index">index</div>
+    <span className="css-in-js-class" />
     <Link href="/about">
       <a id="about-link">about</a>
     </Link>
+    <span className="css-in-js-class" />
   </div>
 )

test/integration/app-document/test/rendering.js

@@ -22,6 +22,14 @@ export default function ({ app }, suiteName, render, fetch) {
         expect($('body').hasClass('custom_class')).toBe(true)
       })
 
+      it('Document.getInitialProps returns html prop representing app shell', async () => {
+        // Extract css-in-js-class from the rendered HTML, which is returned by Document.getInitialProps
+        const $index = await get$('/')
+        const $about = await get$('/about')
+        expect($index('#css-in-cjs-count').text()).toBe('2')
+        expect($about('#css-in-cjs-count').text()).toBe('0')
+      })
+
       test('It injects custom head tags', async () => {
         const $ = await get$('/')
         expect($('head').text()).toMatch('body { margin: 0 }')

test/integration/styled-jsx-module/test/index.test.js

@@ -12,7 +12,6 @@ import {
 } from 'next-test-utils'
 
 const appDir = join(__dirname, '../app')
-const nodeArgs = ['-r', join(appDir, '../../../lib/react-17-require-hook.js')]
 let appPort
 let app
 
@@ -49,13 +48,9 @@ function runTests() {
 describe('styled-jsx using in node_modules', () => {
   describe('Production', () => {
     beforeAll(async () => {
-      await nextBuild(appDir, undefined, {
-        nodeArgs,
-      })
+      await nextBuild(appDir)
       appPort = await findPort()
-      app = await nextStart(appDir, appPort, {
-        nodeArgs,
-      })
+      app = await nextStart(appDir, appPort)
     })
     afterAll(() => killApp(app))
 
@@ -65,9 +60,7 @@ describe('styled-jsx using in node_modules', () => {
   describe('Development', () => {
     beforeAll(async () => {
       appPort = await findPort()
-      app = await launchApp(appDir, appPort, {
-        nodeArgs,
-      })
+      app = await launchApp(appDir, appPort)
     })
     afterAll(() => killApp(app))