Add useFlushEffect hook

errors/functional-document-rsc.md

@@ -0,0 +1,9 @@
+# Functional Next.js Document Components don't currently support React Hooks or Suspense
+
+#### Why This Error Occurred
+
+You tried to use [Suspense](https://reactjs.org/docs/react-api.html#suspense) or a [React hook](https://reactjs.org/docs/hooks-reference.html) like `useState` inside of a functional custom Next.js `Document` component. While these components are conceptually similar to [React Server Components](https://reactjs.org/blog/2020/12/21/data-fetching-with-react-server-components.html), they are very limited in functionality while React Server Components are still experimental.
+
+#### Possible Ways to Fix It
+
+Move any relevant suspense or hooks to the [`App` Component](https://nextjs.org/docs/advanced-features/custom-app) instead.

errors/manifest.json

@@ -514,6 +514,10 @@
         {
           "title": "experimental-jest-transformer",
           "path": "/errors/experimental-jest-transformer.md"
+        },
+        {
+          "title": "functional-document-rsc",
+          "path": "/errors/functional-document-rsc.md"
         }
       ]
     }

packages/next/pages/_document.tsx

@@ -12,8 +12,10 @@ import { cleanAmpPath } from '../server/utils'
 import { htmlEscapeJsonString } from '../server/htmlescape'
 import Script, { ScriptProps } from '../client/script'
 import isError from '../lib/is-error'
+import { useFlushEffect } from '../server/functional-document'
 
 export { DocumentContext, DocumentInitialProps, DocumentProps }
+export { useFlushEffect }
 
 export type OriginProps = {
   nonce?: string

packages/next/server/functional-document.ts

@@ -0,0 +1,60 @@
+import { FunctionComponent, ReactElement } from 'react'
+
+type FlushEffect = () => ReactElement
+type FlushEffectHook = (fn: FlushEffect) => void
+let CURRENT_HOOK_IMPL: FlushEffectHook | null = null
+
+export async function renderFunctionalDocument(
+  Document: FunctionComponent
+): Promise<[ReactElement | null, Array<FlushEffect>]> {
+  const flushEffects: Array<FlushEffect> = []
+  const nextHookImpl: FlushEffectHook = (fn) => {
+    flushEffects.push(fn)
+  }
+  const prevHookImpl = CURRENT_HOOK_IMPL
+
+  try {
+    flushEffects.length = 0
+    CURRENT_HOOK_IMPL = nextHookImpl
+    // Note: we intentionally do not pass props to functional `Document` components. Since this
+    // component is only used on the very first render, we want to prevent handing applications
+    // a footgun where page behavior can unexpectedly differ. Instead, applications should
+    // move such logic to `pages/_app`.
+    const elem = Document({})
+    return [elem, flushEffects]
+  } catch (err: unknown) {
+    if (
+      err &&
+      typeof err === 'object' &&
+      typeof (err as any).then === 'function'
+    ) {
+      throw new Error(
+        'Functional Next.js Document components do not currently support Suspense.\n' +
+          'Read more: https://nextjs.org/docs/messages/functional-document-rsc'
+      )
+    }
+
+    if (
+      err instanceof Error &&
+      /Invalid hook call|Minified React error #321/.test(err.message)
+    ) {
+      throw new Error(
+        'Functional Next.js Document components do not currently support React hooks.\n' +
+          'Read more: https://nextjs.org/docs/messages/functional-document-rsc'
+      )
+    }
+    throw err
+  } finally {
+    CURRENT_HOOK_IMPL = prevHookImpl
+  }
+}
+
+export function useFlushEffect(fn: FlushEffect): void {
+  const currHookImpl = CURRENT_HOOK_IMPL
+  if (!currHookImpl) {
+    throw new Error(
+      'The useFlushEffect hook can only be used by the `Document` component in `pages/_document`'
+    )
+  }
+  return currHookImpl(fn)
+}

packages/next/server/render.tsx

@@ -62,6 +62,7 @@ import {
 import { DomainLocale } from './config'
 import RenderResult, { NodeWritablePiper } from './render-result'
 import isError from '../lib/is-error'
+import { renderFunctionalDocument } from './functional-document'
 
 let Writable: typeof import('stream').Writable
 let Buffer: typeof import('buffer').Buffer
@@ -1007,7 +1008,11 @@ export async function renderToHTML(
     )
     const reader = stream.getReader()
     return new RenderResult((innerRes, next) => {
-      bufferedReadFromReadableStream(reader, (val) => innerRes.write(val)).then(
+      bufferedReadFromReadableStream(
+        reader,
+        (val) => innerRes.write(val),
+        () => null
+      ).then(
         () => next(),
         (innerErr) => next(innerErr)
       )
@@ -1148,6 +1153,26 @@ export async function renderToHTML(
         styles: docProps.styles,
       }
     } else {
+      const styledJsxFlushEffect = () => {
+        const styles = jsxStyleRegistry.styles() as any as React.ReactElement[]
+        jsxStyleRegistry.flush()
+        return styles.length > 0 ? <>{styles}</> : null
+      }
+      const [document, flushEffects] = await renderFunctionalDocument(
+        Document as any
+      )
+      const getFlushPrefix = () => {
+        const elements = [styledJsxFlushEffect, ...flushEffects]
+        .map((fn) => fn())
+          .filter(Boolean)
+          .map((elem, i) => React.cloneElement(elem!, { key: i }))
+        if (elements.length > 0) {
+          return renderToStaticString(<>{elements}</>)
+        } else {
+          return null
+        }
+      }
+
       const bodyResult = async () => {
         const content = (
           <Body>
@@ -1161,16 +1186,16 @@ export async function renderToHTML(
           </Body>
         )
 
-        return concurrentFeatures
+        return await (concurrentFeatures
           ? process.browser
-            ? await renderToWebStream(content)
-            : await renderToNodeStream(content, generateStaticHTML)
-          : piperFromArray([ReactDOMServer.renderToString(content)])
+            ? renderToWebStream(content, getFlushPrefix)
+            : renderToNodeStream(content, generateStaticHTML, getFlushPrefix)
+          : renderToStringStream(content, getFlushPrefix))
       }
 
       return {
         bodyResult,
-        documentElement: () => (Document as any)(),
+        documentElement: () => document,
         useMainContent: (fn?: (content: JSX.Element) => JSX.Element) => {
           if (fn) {
             appWrappers.push(fn)
@@ -1180,7 +1205,7 @@ export async function renderToHTML(
         },
         head,
         headTags: [],
-        styles: jsxStyleRegistry.styles(),
+        styles: [],
       }
     }
   }
@@ -1284,30 +1309,7 @@ export async function renderToHTML(
     </AmpStateContext.Provider>
   )
 
-  let documentHTML: string
-  if (process.browser) {
-    // There is no `renderToStaticMarkup` exposed in the web environment, use
-    // blocking `renderToReadableStream` to get the similar result.
-    let result = ''
-    const readable = (ReactDOMServer as any).renderToReadableStream(document, {
-      onError: (e: any) => {
-        throw e
-      },
-    })
-    const reader = readable.getReader()
-    const decoder = new TextDecoder()
-    while (true) {
-      const { done, value } = await reader.read()
-      if (done) {
-        break
-      }
-      result += typeof value === 'string' ? value : decoder.decode(value)
-    }
-    documentHTML = result
-  } else {
-    documentHTML = ReactDOMServer.renderToStaticMarkup(document)
-  }
-
+  const documentHTML = await renderToStaticString(document)
   const nonRenderedComponents = []
   const expectedDocComponents = ['Main', 'Head', 'NextScript', 'Html']
 
@@ -1411,6 +1413,31 @@ export async function renderToHTML(
   return new RenderResult(chainPipers(pipers))
 }
 
+async function renderToStaticString(element: JSX.Element): Promise<string> {
+  if (process.browser) {
+    // There is no `renderToStaticMarkup` exposed in the web environment, use
+    // blocking `renderToReadableStream` to get the similar result.
+    let result = ''
+    const readable = (ReactDOMServer as any).renderToReadableStream(element, {
+      onError: (e: any) => {
+        throw e
+      },
+    })
+    const reader = readable.getReader()
+    const decoder = new TextDecoder()
+    while (true) {
+      const { done, value } = await reader.read()
+      if (done) {
+        break
+      }
+      result += typeof value === 'string' ? value : decoder.decode(value)
+    }
+    return result
+  } else {
+    return ReactDOMServer.renderToStaticMarkup(element)
+  }
+}
+
 function errorToJSON(err: Error) {
   return {
     name: err.name,
@@ -1435,9 +1462,24 @@ function serializeError(
   }
 }
 
+async function renderToStringStream(
+  element: React.ReactElement,
+  getFlushPrefix: () => Promise<string> | null
+): Promise<NodeWritablePiper> {
+  const pipers = [ReactDOMServer.renderToString(element)]
+  const flushPrefix = getFlushPrefix()
+
+  if (flushPrefix) {
+    pipers.unshift(await flushPrefix)
+  }
+
+  return piperFromArray(pipers)
+}
+
 function renderToNodeStream(
   element: React.ReactElement,
-  generateStaticHTML: boolean
+  generateStaticHTML: boolean,
+  getFlushPrefix: () => Promise<string> | null
 ): Promise<NodeWritablePiper> {
   return new Promise((resolve, reject) => {
     let underlyingStream: {
@@ -1449,22 +1491,27 @@ function renderToNodeStream(
     const stream = new Writable({
       // Use the buffer from the underlying stream
       highWaterMark: 0,
-      writev(chunks, callback) {
-        let str = ''
-        for (let { chunk } of chunks) {
-          str += chunk.toString()
-        }
+      async writev(chunks, callback) {
+        try {
+          if (!underlyingStream) {
+            throw new Error(
+              'invariant: write called without an underlying stream. This is a bug in Next.js'
+            )
+          }
 
-        if (!underlyingStream) {
-          throw new Error(
-            'invariant: write called without an underlying stream. This is a bug in Next.js'
-          )
-        }
+          const flushPrefix = getFlushPrefix()
+          let str = flushPrefix ? await flushPrefix : ''
+          for (let { chunk } of chunks) {
+            str += chunk.toString()
+          }
 
-        if (!underlyingStream.writable.write(str)) {
-          underlyingStream.queuedCallbacks.push(() => callback())
-        } else {
-          callback()
+          if (!underlyingStream.writable.write(str)) {
+            underlyingStream.queuedCallbacks.push(() => callback())
+          } else {
+            callback()
+          }
+        } catch (err: unknown) {
+          callback(err as any)
         }
       },
     })
@@ -1551,7 +1598,8 @@ function renderToNodeStream(
 
 async function bufferedReadFromReadableStream(
   reader: ReadableStreamDefaultReader,
-  writeFn: (val: string) => void
+  writeFn: (val: string) => void,
+  getFlushPrefix: () => Promise<string> | null
 ): Promise<void> {
   const decoder = new TextDecoder()
   let bufferedString = ''
@@ -1576,6 +1624,11 @@ async function bufferedReadFromReadableStream(
       break
     }
 
+    const flushPrefix = getFlushPrefix()
+    if (flushPrefix) {
+      bufferedString += await flushPrefix
+    }
+
     bufferedString += typeof value === 'string' ? value : decoder.decode(value)
     flushBuffer()
   }
@@ -1585,7 +1638,8 @@ async function bufferedReadFromReadableStream(
 }
 
 function renderToWebStream(
-  element: React.ReactElement
+  element: React.ReactElement,
+  getFlushPrefix: () => Promise<string> | null
 ): Promise<NodeWritablePiper> {
   return new Promise((resolve, reject) => {
     let resolved = false
@@ -1602,8 +1656,10 @@ function renderToWebStream(
         if (!resolved) {
           resolved = true
           resolve((res, next) => {
-            bufferedReadFromReadableStream(reader, (val) =>
-              res.write(val)
+            bufferedReadFromReadableStream(
+              reader,
+              (val) => res.write(val),
+              getFlushPrefix
             ).then(
               () => next(),
               (err) => next(err)

test/integration/document-functional-flush-effect/pages/_document.js

@@ -0,0 +1,16 @@
+import React from 'react'
+import { useFlushEffect, Html, Head, Main, NextScript } from 'next/document'
+
+export default function Document() {
+  useFlushEffect(() => <foo />)
+  useFlushEffect(() => <bar />)
+  return (
+    <Html>
+      <Head />
+      <body>
+        <Main />
+        <NextScript />
+      </body>
+    </Html>
+  )
+}

test/integration/document-functional-flush-effect/pages/index.js

@@ -0,0 +1,3 @@
+export default function Index() {
+  return <span>Hello World!</span>
+}

test/integration/document-functional-flush-effect/test/index.test.js

@@ -0,0 +1,24 @@
+/* eslint-env jest */
+
+import { join } from 'path'
+import { findPort, launchApp, killApp, renderViaHTTP } from 'next-test-utils'
+
+const appDir = join(__dirname, '..')
+let appPort
+let app
+
+describe('Functional Custom Document', () => {
+  describe('development mode', () => {
+    beforeAll(async () => {
+      appPort = await findPort()
+      app = await launchApp(appDir, appPort)
+    })
+
+    afterAll(() => killApp(app))
+
+    it('supports custom flush handlers', async () => {
+      const html = await renderViaHTTP(appPort, '/')
+      expect(html).toMatch(/<foo><\/foo><bar><\/bar><span>Hello World!<\/span>/)
+    })
+  })
+})