Fix _document and getInitialProps with React 18 (#35736)

* wip * update test * fix _document logic for edge runtime and rsc * revert deleted file * fix lint error * fix * remove doc gip test * Revert "remove doc gip test" This reverts commit a5fd1d7. * fix test Co-authored-by: Jiachi Liu <inbox@huozhi.im>
vercel · Mar 30, 2022 · 600e0b3 · 600e0b3
1 parent cae9506
commit 600e0b3
Show file tree

Hide file tree

Showing 5 changed files with 159 additions and 87 deletions.
diff --git a/packages/next/server/render.tsx b/packages/next/server/render.tsx
@@ -1268,12 +1268,17 @@ export async function renderToHTML(
       }
     }
 
-    // We make it a function component to enable streaming.
-    if (hasConcurrentFeatures && builtinDocument) {
-      Document = builtinDocument
+    if ((isServerComponent || process.browser) && Document.getInitialProps) {
+      if (builtinDocument) {
+        Document = builtinDocument
+      } else {
+        throw new Error(
+          '`getInitialProps` in Document component is not supported with React Server Components.'
+        )
+      }
     }
 
-    if (!hasConcurrentFeatures && Document.getInitialProps) {
+    async function documentInitialProps() {
       const renderPage: RenderPage = (
         options: ComponentsEnhancer = {}
       ): RenderPageResult | Promise<RenderPageResult> => {
@@ -1323,96 +1328,130 @@ export async function renderToHTML(
         throw new Error(message)
       }
 
-      return {
-        bodyResult: (suffix: string) =>
-          streamFromArray([docProps.html, suffix]),
-        documentElement: (htmlProps: HtmlProps) => (
-          <Document {...htmlProps} {...docProps} />
-        ),
-        head: docProps.head,
-        headTags: await headTags(documentCtx),
-        styles: docProps.styles,
-      }
-    } else {
-      let bodyResult
-
-      const renderContent = () => {
-        return ctx.err && ErrorDebug ? (
-          <Body>
-            <ErrorDebug error={ctx.err} />
-          </Body>
-        ) : (
-          <Body>
-            <AppContainerWithIsomorphicFiberStructure>
-              {isServerComponent && AppMod.__next_rsc__ ? (
-                // _app.server.js is used.
-                <Component {...props.pageProps} router={router} />
-              ) : (
-                <App {...props} Component={Component} router={router} />
-              )}
-            </AppContainerWithIsomorphicFiberStructure>
-          </Body>
-        )
-      }
-
-      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,
-        })
+      return { docProps, documentCtx }
+    }
 
-        bodyResult = async (suffix: string) => {
-          // this must be called inside bodyResult so appWrappers is
-          // up to date when getWrappedApp is called
-
-          const flushEffectHandler = async () => {
-            const allFlushEffects = [
-              styledJsxFlushEffect,
-              ...(flushEffects || []),
-            ]
-            const flushEffectStream = await renderToStream({
-              ReactDOMServer,
-              element: (
-                <>
-                  {allFlushEffects.map((flushEffect, i) => (
-                    <React.Fragment key={i}>{flushEffect()}</React.Fragment>
-                  ))}
-                </>
-              ),
-              generateStaticHTML: true,
-            })
-            const flushed = await streamToString(flushEffectStream)
-            return flushed
-          }
+    const renderContent = () => {
+      return ctx.err && ErrorDebug ? (
+        <Body>
+          <ErrorDebug error={ctx.err} />
+        </Body>
+      ) : (
+        <Body>
+          <AppContainerWithIsomorphicFiberStructure>
+            {isServerComponent && AppMod.__next_rsc__ ? (
+              // _app.server.js is used.
+              <Component {...props.pageProps} router={router} />
+            ) : (
+              <App {...props} Component={Component} router={router} />
+            )}
+          </AppContainerWithIsomorphicFiberStructure>
+        </Body>
+      )
+    }
 
-          return await continueFromInitialStream({
-            renderStream,
-            suffix,
-            dataStream: serverComponentsInlinedTransformStream?.readable,
-            generateStaticHTML: generateStaticHTML || !hasConcurrentFeatures,
-            flushEffectHandler,
-          })
+    if (!hasConcurrentFeatures) {
+      if (Document.getInitialProps) {
+        const documentInitialPropsRes = await documentInitialProps()
+        if (documentInitialPropsRes === null) return null
+        const { docProps, documentCtx } = documentInitialPropsRes
+
+        return {
+          bodyResult: (suffix: string) =>
+            streamFromArray([docProps.html, suffix]),
+          documentElement: (htmlProps: HtmlProps) => (
+            <Document {...htmlProps} {...docProps} />
+          ),
+          head: docProps.head,
+          headTags: await headTags(documentCtx),
+          styles: docProps.styles,
         }
       } else {
         const content = renderContent()
         // for non-concurrent rendering we need to ensure App is rendered
         // before _document so that updateHead is called/collected before
         // rendering _document's head
         const result = ReactDOMServer.renderToString(content)
-        bodyResult = (suffix: string) => streamFromArray([result, suffix])
+        const bodyResult = (suffix: string) => streamFromArray([result, suffix])
+
+        const styles = jsxStyleRegistry.styles()
+        jsxStyleRegistry.flush()
+
+        return {
+          bodyResult,
+          documentElement: () => (Document as any)(),
+          head,
+          headTags: [],
+          styles,
+        }
+      }
+    } else {
+      let bodyResult
+
+      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 flushEffectHandler = async () => {
+          const allFlushEffects = [
+            styledJsxFlushEffect,
+            ...(flushEffects || []),
+          ]
+          const flushEffectStream = await renderToStream({
+            ReactDOMServer,
+            element: (
+              <>
+                {allFlushEffects.map((flushEffect, i) => (
+                  <React.Fragment key={i}>{flushEffect()}</React.Fragment>
+                ))}
+              </>
+            ),
+            generateStaticHTML: true,
+          })
+          const flushed = await streamToString(flushEffectStream)
+          return flushed
+        }
+
+        return await continueFromInitialStream({
+          renderStream,
+          suffix,
+          dataStream: serverComponentsInlinedTransformStream?.readable,
+          generateStaticHTML: generateStaticHTML || !hasConcurrentFeatures,
+          flushEffectHandler,
+        })
       }
 
       const styles = jsxStyleRegistry.styles()
       jsxStyleRegistry.flush()
 
+      const documentInitialPropsRes =
+        isServerComponent || process.browser || !Document.getInitialProps
+          ? {}
+          : await documentInitialProps()
+      if (documentInitialPropsRes === null) return null
+
+      const documentElement = () => {
+        if (isServerComponent || process.browser) {
+          return (Document as any)()
+        }
+
+        const { docProps } = (documentInitialPropsRes as any) || {}
+        return <Document {...htmlProps} {...docProps} />
+      }
+
       return {
         bodyResult,
-        documentElement: () => (Document as any)(),
+        documentElement,
         head,
         headTags: [],
         styles,

diff --git a/test/development/basic/styled-components.test.ts b/test/development/basic/styled-components.test.ts
@@ -17,6 +17,8 @@ describe('styled-components SWC transform', () => {
       },
       dependencies: {
         'styled-components': '5.3.3',
+        react: 'latest',
+        'react-dom': 'latest',
       },
     })
   })
@@ -34,6 +36,7 @@ describe('styled-components SWC transform', () => {
     })
     return foundLog
   }
+
   it('should not have hydration mismatch with styled-components transform enabled', async () => {
     let browser
     try {
@@ -56,4 +59,14 @@ describe('styled-components SWC transform', () => {
       }
     }
   })
+
+  it('should render the page with correct styles', async () => {
+    const browser = await webdriver(next.appPort, '/')
+
+    expect(
+      await browser.eval(
+        `window.getComputedStyle(document.querySelector('#btn')).color`
+      )
+    ).toBe('rgb(255, 255, 255)')
+  })
 })
diff --git a/test/development/basic/styled-components/next.config.js b/test/development/basic/styled-components/next.config.js
@@ -1,5 +1,17 @@
-module.exports = {
+const path = require('path')
+
+let withReact18 = (config) => config
+
+try {
+  // only used when running inside of the monorepo not when isolated
+  withReact18 = require(path.join(
+    __dirname,
+    '../../../integration/react-18/test/with-react-18'
+  ))
+} catch (_) {}
+
+module.exports = withReact18({
   compiler: {
     styledComponents: true,
   },
-}
+})
diff --git a/test/development/basic/styled-components/pages/index.js b/test/development/basic/styled-components/pages/index.js
@@ -32,7 +32,9 @@ export default function Home() {
         GitHub
       </Button>
 
-      <Button href="/docs">Documentation</Button>
+      <Button id="btn" href="/docs">
+        Documentation
+      </Button>
     </div>
   )
 }
diff --git a/test/integration/react-streaming-and-server-components/test/index.test.js b/test/integration/react-streaming-and-server-components/test/index.test.js
@@ -241,11 +241,17 @@ const cssSuite = {
 }
 
 const documentSuite = {
-  runTests: (context) => {
-    it('should error when custom _document has getInitialProps method', async () => {
-      const res = await fetchViaHTTP(context.appPort, '/')
-      expect(res.status).toBe(500)
-    })
+  runTests: (context, env) => {
+    if (env === 'dev') {
+      it('should error when custom _document has getInitialProps method', async () => {
+        const res = await fetchViaHTTP(context.appPort, '/')
+        expect(res.status).toBe(500)
+      })
+    } else {
+      it('should failed building', async () => {
+        expect(context.code).toBe(1)
+      })
+    }
   },
   beforeAll: () => documentPage.write(documentWithGip),
   afterAll: () => documentPage.delete(),
@@ -270,7 +276,7 @@ function runSuite(suiteName, env, options) {
       options.beforeAll?.()
       if (env === 'prod') {
         context.appPort = await findPort()
-        await nextBuild(context.appDir)
+        context.code = (await nextBuild(context.appDir)).code
         context.server = await nextStart(context.appDir, context.appPort)
       }
       if (env === 'dev') {