Flush initial styled-jsx in gIP first in concurrent rendering (#36594)

huozhi · kodiakhq[bot] · shuding · web-flow · commit fcec7587791e · 2022-05-02T22:52:46.000+02:00
* Use flushed effects to generate styled-jsx styles insted of gIP by default

* ensure styles are flushed inside the default getInitialProps

Co-authored-by: kodiakhq[bot] &lt;49736102+kodiakhq[bot]@users.noreply.github.com&gt;
Co-authored-by: Shu Ding &lt;g@shud.in&gt;
diff --git a/packages/next/pages/_document.tsx b/packages/next/pages/_document.tsx
@@ -1,9 +1,13 @@
 import React, { Component, ReactElement, ReactNode, useContext } from 'react'
-import { OPTIMIZED_FONT_PROVIDERS } from '../shared/lib/constants'
+import {
+  OPTIMIZED_FONT_PROVIDERS,
+  NEXT_BUILTIN_DOCUMENT,
+} from '../shared/lib/constants'
 import type {
   DocumentContext,
   DocumentInitialProps,
   DocumentProps,
+  DocumentType,
 } from '../shared/lib/utils'
 import { BuildManifest, getPageFiles } from '../server/get-page-files'
 import { cleanAmpPath } from '../server/utils'
@@ -309,7 +313,7 @@ export default class Document<P = {}> extends Component<DocumentProps & P> {
 
 // Add a special property to the built-in `Document` component so later we can
 // identify if a user customized `Document` is used or not.
-;(Document as any).__next_internal_document =
+const InternalFunctionDocument: DocumentType =
   function InternalFunctionDocument() {
     return (
       <Html>
@@ -321,6 +325,7 @@ export default class Document<P = {}> extends Component<DocumentProps & P> {
       </Html>
     )
   }
+;(Document as any)[NEXT_BUILTIN_DOCUMENT] = InternalFunctionDocument
 
 export function Html(
   props: React.DetailedHTMLProps<
diff --git a/packages/next/server/base-server.ts b/packages/next/server/base-server.ts
@@ -27,6 +27,7 @@ import { parse as parseQs } from 'querystring'
 import { format as formatUrl, parse as parseUrl } from 'url'
 import { getRedirectStatus } from '../lib/load-custom-routes'
 import {
+  NEXT_BUILTIN_DOCUMENT,
   SERVERLESS_DIRECTORY,
   SERVER_DIRECTORY,
   STATIC_STATUS_PAGES,
@@ -1221,7 +1222,7 @@ export default abstract class Server<ServerOptions extends Options = Options> {
         // When concurrent features is enabled, the built-in `Document`
         // component also supports dynamic HTML.
         (this.renderOpts.reactRoot &&
-          !!(components.Document as any)?.__next_internal_document)
+          NEXT_BUILTIN_DOCUMENT in components.Document)
 
       // Disable dynamic HTML in cases that we know it won't be generated,
       // so that we can continue generating a cache key when possible.
diff --git a/packages/next/server/render.tsx b/packages/next/server/render.tsx
@@ -6,6 +6,7 @@ import type { DomainLocale } from './config'
 import type {
   AppType,
   DocumentInitialProps,
+  DocumentType,
   DocumentProps,
   DocumentContext,
   NextComponentType,
@@ -35,6 +36,7 @@ import {
   UNSTABLE_REVALIDATE_RENAME_ERROR,
 } from '../lib/constants'
 import {
+  NEXT_BUILTIN_DOCUMENT,
   SERVER_PROPS_ID,
   STATIC_PROPS_ID,
   STATIC_STATUS_PAGES,
@@ -769,6 +771,7 @@ export async function renderToHTML(
 
       const { html, head } = await docCtx.renderPage({ enhanceApp })
       const styles = jsxStyleRegistry.styles({ nonce: options.nonce })
+      jsxStyleRegistry.flush()
       return { html, head, styles }
     },
   }
@@ -1311,14 +1314,14 @@ export async function renderToHTML(
     // 1. Using `Document.getInitialProps` in the Edge runtime.
     // 2. Using the class component `Document` with concurrent features.
 
-    const builtinDocument = (Document as any).__next_internal_document as
-      | typeof Document
-      | undefined
+    const BuiltinFunctionalDocument: DocumentType | undefined = (
+      Document as any
+    )[NEXT_BUILTIN_DOCUMENT]
 
     if (process.env.NEXT_RUNTIME === 'edge' && Document.getInitialProps) {
       // In the Edge runtime, `Document.getInitialProps` isn't supported.
       // We throw an error here if it's customized.
-      if (!builtinDocument) {
+      if (!BuiltinFunctionalDocument) {
         throw new Error(
           '`getInitialProps` in Document component is not supported with the Edge Runtime.'
         )
@@ -1329,8 +1332,8 @@ export async function renderToHTML(
       (isServerComponent || process.env.NEXT_RUNTIME === 'edge') &&
       Document.getInitialProps
     ) {
-      if (builtinDocument) {
-        Document = builtinDocument
+      if (BuiltinFunctionalDocument) {
+        Document = BuiltinFunctionalDocument
       } else {
         throw new Error(
           '`getInitialProps` in Document component is not supported with React Server Components.'
@@ -1433,6 +1436,7 @@ 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
@@ -1468,6 +1472,7 @@ export async function renderToHTML(
         }
       }
     } else {
+      // Enabling react concurrent rendering mode: __NEXT_REACT_ROOT = true
       let renderStream: ReadableStream<Uint8Array> & {
         allReady?: Promise<void> | undefined
       }
@@ -1534,13 +1539,11 @@ export async function renderToHTML(
         // be streamed.
         // Do not use `await` here.
         generateStaticFlightDataIfNeeded()
-
         return await continueFromInitialStream({
           renderStream,
           suffix,
           dataStream: serverComponentsInlinedTransformStream?.readable,
-          generateStaticHTML:
-            generateStaticHTML || !process.env.__NEXT_REACT_ROOT,
+          generateStaticHTML,
           flushEffectHandler,
         })
       }
@@ -1572,8 +1575,8 @@ export async function renderToHTML(
 
         return <Document {...htmlProps} {...docProps} />
       }
-      let styles
 
+      let styles
       if (hasDocumentGetInitialProps) {
         styles = docProps.styles
       } else {
diff --git a/packages/next/shared/lib/constants.ts b/packages/next/shared/lib/constants.ts
@@ -25,6 +25,7 @@ export const CLIENT_PUBLIC_FILES_PATH = 'public'
 export const CLIENT_STATIC_FILES_PATH = 'static'
 export const CLIENT_STATIC_FILES_RUNTIME = 'runtime'
 export const STRING_LITERAL_DROP_BUNDLE = '__NEXT_DROP_CLIENT_FILE__'
+export const NEXT_BUILTIN_DOCUMENT = '__NEXT_BUILTIN_DOCUMENT__'
 
 // server/middleware-flight-manifest.js
 export const MIDDLEWARE_FLIGHT_MANIFEST = 'middleware-flight-manifest'
diff --git a/test/e2e/styled-jsx/app/pages/index.js b/test/e2e/styled-jsx/app/pages/index.js
@@ -0,0 +1,8 @@
+export default function Page() {
+  return (
+    <div>
+      <style jsx>{'p { color: red; }'}</style>
+      <p>hello world</p>
+    </div>
+  )
+}
diff --git a/test/e2e/styled-jsx/index.test.ts b/test/e2e/styled-jsx/index.test.ts
@@ -1,25 +1,19 @@
-import { createNext } from 'e2e-utils'
+import path from 'path'
+import { createNext, FileRef } from 'e2e-utils'
 import { NextInstance } from 'test/lib/next-modes/base'
 import { renderViaHTTP } from 'next-test-utils'
 import cheerio from 'cheerio'
 
+const appDir = path.join(__dirname, 'app')
+
 function runTest(packageManager?: string) {
   describe(`styled-jsx${packageManager ? ' ' + packageManager : ''}`, () => {
     let next: NextInstance
 
     beforeAll(async () => {
       next = await createNext({
         files: {
-          'pages/index.js': `
-            export default function Page() {
-              return (
-                <div>
-                  <style jsx>{'p { color: red; }'}</style>
-                  <p>hello world</p>
-                </div>
-              )
-            }
-          `,
+          pages: new FileRef(path.join(appDir, 'pages')),
         },
         ...(packageManager
           ? {
@@ -33,7 +27,7 @@ function runTest(packageManager?: string) {
     it('should contain styled-jsx styles in html', async () => {
       const html = await renderViaHTTP(next.url, '/')
       const $ = cheerio.load(html)
-      expect($('head').text()).toContain('color:red')
+      expect($('html').text()).toMatch(/color:(\s)*red/)
     })
   })
 }
diff --git a/test/integration/react-18/test/concurrent.js b/test/integration/react-18/test/concurrent.js
@@ -32,7 +32,14 @@ export default (context, _render) => {
     })
   })
 
-  it('flushes styles as the page renders', async () => {
+  it('flushes styled-jsx styles as the page renders', async () => {
+    const html = await renderViaHTTP(
+      context.appPort,
+      '/use-flush-effect/styled-jsx'
+    )
+    const stylesOccurrence = html.match(/color:(\s)*blue/g) || []
+    expect(stylesOccurrence.length).toBe(1)
+
     await withBrowser('/use-flush-effect/styled-jsx', async (browser) => {
       await check(
         () => browser.waitForElementByCss('#__jsx-900f996af369fc74').text(),
