Flush styles effects (#39268)

Use flush effects to custom apply css-in-js solution to app. Re-introduce flush effects to app-render, and remove default support of styled-jsx in `app/`. So that users will choose their own css-in-js solution if they need any customization. styled-jsx won't appear in client bundle if you didn't use it.

For now we have to inject the initial styles before `</head>` to avoid hydration errors. Later on we can remove this once react can handle it.

- [x] inject styles before end of head element
- [x] add tests

package.json

@@ -191,6 +191,7 @@
     "semver": "7.3.7",
     "shell-quote": "1.7.3",
     "styled-components": "5.3.3",
+    "styled-jsx": "link:packages/next/node_modules/styled-jsx",
     "styled-jsx-plugin-postcss": "3.0.2",
     "tailwindcss": "1.1.3",
     "taskr": "1.1.0",

packages/next/build/webpack-config.ts

@@ -974,7 +974,7 @@ export default async function getBaseWebpackConfig(
   }
 
   const rscSharedRegex =
-    /(node_modules\/react\/|\/shared\/lib\/(head-manager-context|router-context)\.js|node_modules\/styled-jsx\/)/
+    /(node_modules\/react\/|\/shared\/lib\/(head-manager-context|router-context|flush-effects)\.js|node_modules\/styled-jsx\/)/
 
   let webpackConfig: webpack.Configuration = {
     parallelism: Number(process.env.NEXT_WEBPACK_PARALLELISM) || undefined,

packages/next/client/components/hooks-client-context.ts

@@ -1,5 +1,5 @@
 import { createContext } from 'react'
-import { NextParsedUrlQuery } from '../../server/request-meta'
+import type { NextParsedUrlQuery } from '../../server/request-meta'
 
 export const SearchParamsContext = createContext<NextParsedUrlQuery>(
   null as any

packages/next/client/components/hooks-client.ts

@@ -12,6 +12,11 @@ import {
   LayoutRouterContext,
 } from '../../shared/lib/app-router-context'
 
+export {
+  FlushEffectsContext,
+  useFlushEffects,
+} from '../../shared/lib/flush-effects'
+
 /**
  * Get the current search params. For example useSearchParams() would return {"foo": "bar"} when ?foo=bar
  */

packages/next/client/components/match-segments.ts

@@ -1,4 +1,4 @@
-import { Segment } from '../../server/app-render'
+import type { Segment } from '../../server/app-render'
 
 export const matchSegment = (
   existingSegment: Segment,

packages/next/server/app-render.tsx

@@ -6,7 +6,6 @@ import React from 'react'
 import { ParsedUrlQuery, stringify as stringifyQuery } from 'querystring'
 import { createFromReadableStream } from 'next/dist/compiled/react-server-dom-webpack'
 import { renderToReadableStream } from 'next/dist/compiled/react-server-dom-webpack/writer.browser.server'
-import { StyleRegistry, createStyleRegistry } from 'styled-jsx'
 import { NextParsedUrlQuery } from './request-meta'
 import RenderResult from './render-result'
 import {
@@ -23,6 +22,7 @@ import { htmlEscapeJsonString } from './htmlescape'
 import { shouldUseReactRoot, stripInternalQueries } from './utils'
 import { NextApiRequestCookies } from './api-utils'
 import { matchSegment } from '../client/components/match-segments'
+import { FlushEffectsContext } from '../client/components/hooks-client'
 
 // this needs to be required lazily so that `next-server` can set
 // the env before we require
@@ -985,18 +985,26 @@ export async function renderToHTMLOrFlight(
     }
   )
 
-  /**
-   * Style registry for styled-jsx
-   */
-  const jsxStyleRegistry = createStyleRegistry()
+  let flushEffectsHandler: (() => React.ReactNode) | null = null
+  function FlushEffects({ children }: { children: JSX.Element }) {
+    // Reset flushEffectsHandler on each render
+    flushEffectsHandler = null
+    const setFlushEffectsHandler = React.useCallback(
+      (handler: () => React.ReactNode) => {
+        if (flushEffectsHandler)
+          throw new Error(
+            'The `useFlushEffects` hook cannot be used more than once.'
+          )
+        flushEffectsHandler = handler
+      },
+      []
+    )
 
-  /**
-   * styled-jsx styles as React Component
-   */
-  const styledJsxFlushEffect = (): React.ReactNode => {
-    const styles = jsxStyleRegistry.styles()
-    jsxStyleRegistry.flush()
-    return <>{styles}</>
+    return (
+      <FlushEffectsContext.Provider value={setFlushEffectsHandler}>
+        {children}
+      </FlushEffectsContext.Provider>
+    )
   }
 
   /**
@@ -1015,11 +1023,18 @@ export async function renderToHTMLOrFlight(
   const generateStaticHTML = supportsDynamicHTML !== true
   const bodyResult = async () => {
     const content = (
-      <StyleRegistry registry={jsxStyleRegistry}>
+      <FlushEffects>
         <ServerComponentsRenderer />
-      </StyleRegistry>
+      </FlushEffects>
     )
 
+    const flushEffectHandler = (): string => {
+      const flushed = ReactDOMServer.renderToString(
+        <>{flushEffectsHandler && flushEffectsHandler()}</>
+      )
+      return flushed
+    }
+
     const renderStream = await renderToInitialStream({
       ReactDOMServer,
       element: content,
@@ -1031,17 +1046,13 @@ export async function renderToHTMLOrFlight(
       },
     })
 
-    const flushEffectHandler = (): string => {
-      const flushed = ReactDOMServer.renderToString(styledJsxFlushEffect())
-      return flushed
-    }
-
     const hasConcurrentFeatures = !!runtime
 
     return await continueFromInitialStream(renderStream, {
       dataStream: serverComponentsInlinedTransformStream?.readable,
       generateStaticHTML: generateStaticHTML || !hasConcurrentFeatures,
       flushEffectHandler,
+      flushEffectsToHead: true,
       initialStylesheets,
     })
   }

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

@@ -137,7 +137,7 @@ export function createFlushEffectStream(
 }
 
 export function createHeadInjectionTransformStream(
-  inject: string
+  inject: () => string
 ): TransformStream<Uint8Array, Uint8Array> {
   let injected = false
   return new TransformStream({
@@ -147,7 +147,7 @@ export function createHeadInjectionTransformStream(
       if (!injected && (index = content.indexOf('</head')) !== -1) {
         injected = true
         const injectedContent =
-          content.slice(0, index) + inject + content.slice(index)
+          content.slice(0, index) + inject() + content.slice(index)
         controller.enqueue(encodeText(injectedContent))
       } else {
         controller.enqueue(chunk)
@@ -175,12 +175,14 @@ export async function continueFromInitialStream(
     dataStream,
     generateStaticHTML,
     flushEffectHandler,
+    flushEffectsToHead,
     initialStylesheets,
   }: {
     suffix?: string
     dataStream?: ReadableStream<Uint8Array>
     generateStaticHTML: boolean
     flushEffectHandler?: () => string
+    flushEffectsToHead: boolean
     initialStylesheets?: string[]
   }
 ): Promise<ReadableStream<Uint8Array>> {
@@ -193,15 +195,22 @@ export async function continueFromInitialStream(
 
   const transforms: Array<TransformStream<Uint8Array, Uint8Array>> = [
     createBufferedTransformStream(),
-    flushEffectHandler ? createFlushEffectStream(flushEffectHandler) : null,
+    flushEffectHandler && !flushEffectsToHead
+      ? createFlushEffectStream(flushEffectHandler)
+      : null,
     suffixUnclosed != null ? createDeferredSuffixStream(suffixUnclosed) : null,
     dataStream ? createInlineDataStream(dataStream) : null,
     suffixUnclosed != null ? createSuffixStream(closeTag) : null,
-    createHeadInjectionTransformStream(
-      (initialStylesheets || [])
+    createHeadInjectionTransformStream(() => {
+      const inlineStyleLinks = (initialStylesheets || [])
         .map((href) => `<link rel="stylesheet" href="/_next/${href}">`)
         .join('')
-    ),
+      // TODO-APP: Inject flush effects to end of head in app layout rendering, to avoid
+      // hydration errors. Remove this once it's ready to be handled by react itself.
+      const flushEffectsContent =
+        flushEffectHandler && flushEffectsToHead ? flushEffectHandler() : ''
+      return inlineStyleLinks + flushEffectsContent
+    }),
   ].filter(nonNullable)
 
   return transforms.reduce(

packages/next/server/render.tsx

@@ -1269,6 +1269,7 @@ export async function renderToHTML(
           dataStream: serverComponentsInlinedTransformStream?.readable,
           generateStaticHTML,
           flushEffectHandler,
+          flushEffectsToHead: false,
         })
       }
 

packages/next/shared/lib/flush-effects.tsx

@@ -0,0 +1,14 @@
+import React, { createContext, useContext } from 'react'
+
+export type FlushEffectsHook = (callbacks: () => React.ReactNode) => void
+
+export const FlushEffectsContext = createContext<FlushEffectsHook | null>(
+  null as any
+)
+
+export function useFlushEffects(callbacks: () => React.ReactNode): void {
+  const flushEffectsImpl = useContext(FlushEffectsContext)
+  // Should have no effects on client where there's no flush effects provider
+  if (!flushEffectsImpl) return
+  return flushEffectsImpl(callbacks)
+}

test/e2e/app-dir/rsc-basic.test.ts

@@ -41,6 +41,8 @@ describe('app dir - react server components', () => {
         'next.config.js': new FileRef(path.join(appDir, 'next.config.js')),
       },
       dependencies: {
+        'styled-jsx': 'latest',
+        'styled-components': '6.0.0-alpha.5',
         react: 'experimental',
         'react-dom': 'experimental',
       },
@@ -320,11 +322,15 @@ describe('app dir - react server components', () => {
     expect(content).toContain('bar.server.js:')
   })
 
-  it.skip('should SSR styled-jsx correctly', async () => {
-    const html = await renderViaHTTP(next.url, '/styled-jsx')
-    const styledJsxClass = getNodeBySelector(html, 'h1').attr('class')
+  it('should render initial styles of css-in-js in SSR correctly', async () => {
+    const html = await renderViaHTTP(next.url, '/css-in-js')
+    const head = getNodeBySelector(html, 'head').html()
 
-    expect(html).toContain(`h1.${styledJsxClass}{color:red}`)
+    // from styled-jsx
+    expect(head).toMatch(/{color:(\s*)purple;?}/)
+
+    // from styled-components
+    expect(head).toMatch(/{color:(\s*)blue;?}/)
   })
 
   it('should support streaming for flight response', async () => {

test/e2e/app-dir/rsc-basic/app/css-in-js/page.server.js

@@ -0,0 +1,11 @@
+import Comp from './styled-jsx.client'
+import StyledComp from './styled-components.client'
+
+export default function Page() {
+  return (
+    <div>
+      <Comp />
+      <StyledComp />
+    </div>
+  )
+}

test/e2e/app-dir/rsc-basic/app/css-in-js/styled-components.client.js

@@ -0,0 +1,30 @@
+import styled from 'styled-components'
+
+const Button = styled.button`
+  display: inline-block;
+  border-radius: 3px;
+  padding: 0.5rem 0;
+  margin: 0.5rem 1rem;
+  width: 11rem;
+  color: blue;
+  border: 2px solid blue;
+`
+
+const Box = styled.div`
+  border: 1px solid blue;
+  padding: 8px;
+  margin: 8px 0;
+`
+
+const Title = styled.h3`
+  color: blue;
+`
+
+export default () => {
+  return (
+    <Box>
+      <Title>styled-components</Title>
+      <Button>{`💅 This area belongs to styled-components`}</Button>
+    </Box>
+  )
+}

test/e2e/app-dir/rsc-basic/app/css-in-js/styled-jsx.client.js

@@ -0,0 +1,19 @@
+export default function Comp() {
+  return (
+    <div>
+      <style jsx>{`
+        h3 {
+          color: purple;
+        }
+        .box {
+          padding: 8px;
+          border: 2px solid purple;
+        }
+      `}</style>
+      <div className="box">
+        <h3>styled-jsx</h3>
+        <p>This area is rendered by styled-jsx</p>
+      </div>
+    </div>
+  )
+}

test/e2e/app-dir/rsc-basic/app/layout.server.js

@@ -1,12 +1,15 @@
 import React from 'react'
+import RootStyleRegistry from './root-style-registry.client'
 
 export default function AppLayout({ children }) {
   return (
     <html>
       <head>
         <title>RSC</title>
       </head>
-      <body>{children}</body>
+      <body>
+        <RootStyleRegistry>{children}</RootStyleRegistry>
+      </body>
     </html>
   )
 }

test/e2e/app-dir/rsc-basic/app/root-style-registry.client.js

@@ -0,0 +1,43 @@
+import React from 'react'
+import { StyleRegistry, createStyleRegistry } from 'styled-jsx'
+import { ServerStyleSheet, StyleSheetManager } from 'styled-components'
+import { useFlushEffects } from 'next/dist/client/components/hooks-client'
+import { useState } from 'react'
+
+export default function RootStyleRegistry({ children }) {
+  const [jsxStyleRegistry] = useState(() => createStyleRegistry())
+  const [styledComponentsStyleSheet] = useState(() => new ServerStyleSheet())
+  const styledJsxFlushEffect = () => {
+    const styles = jsxStyleRegistry.styles()
+    jsxStyleRegistry.flush()
+    return <>{styles}</>
+  }
+  const styledComponentsFlushEffect = () => {
+    const styles = styledComponentsStyleSheet.getStyleElement()
+    styledComponentsStyleSheet.seal()
+
+    return <>{styles}</>
+  }
+
+  useFlushEffects(() => {
+    const effects = styledComponentsFlushEffect()
+
+    return (
+      <>
+        {styledJsxFlushEffect()}
+        {effects}
+      </>
+    )
+  })
+
+  // Only include style registry on server side for SSR
+  if (typeof window === 'undefined') {
+    return (
+      <StyleSheetManager sheet={styledComponentsStyleSheet.instance}>
+        <StyleRegistry registry={jsxStyleRegistry}>{children}</StyleRegistry>
+      </StyleSheetManager>
+    )
+  }
+
+  return children
+}