Refactor server/render for SSR streaming

packages/next/build/webpack/loaders/next-middleware-ssr-loader/index.ts

@@ -57,47 +57,7 @@ export default async function middlewareRSCLoader(this: any) {
           throw new Error('Your page must export a \`default\` component')
         }
 
-        function wrapReadable(readable) {
-          const encoder = new TextEncoder()
-          const transformStream = new TransformStream()
-          const writer = transformStream.writable.getWriter()
-          const reader = readable.getReader()
-          const process = () => {
-            reader.read().then(({ done, value }) => {
-              if (!done) {
-                writer.write(typeof value === 'string' ? encoder.encode(value) : value)
-                process()
-              } else {
-                writer.close()
-              }
-            })
-          }
-          process()
-          return transformStream.readable
-        }
-
-        ${
-          isServerComponent
-            ? `
-        const renderFlight = props => renderToReadableStream(createElement(Page, props), rscManifest)
-
-        let responseCache
-        const FlightWrapper = props => {
-          let response = responseCache
-          if (!response) {
-            responseCache = response = createFromReadableStream(renderFlight(props))
-          }
-          return response.readRoot()
-        }
-        const Component = props => {
-          return createElement(
-            React.Suspense,
-            { fallback: null },
-            createElement(FlightWrapper, props)
-          )
-        }`
-            : `const Component = Page`
-        }
+        const Component = Page
 
         async function render(request) {
           const url = request.nextUrl
@@ -115,23 +75,12 @@ export default async function middlewareRSCLoader(this: any) {
             isServerComponent
               ? `
           // Flight data request
-          const isFlightDataRequest = query.__flight__ !== undefined
-          if (isFlightDataRequest) {
+          const renderServerComponentData = query.__flight__ !== undefined
+          if (renderServerComponentData) {
             delete query.__flight__
-            return new Response(
-              wrapReadable(
-                renderFlight({
-                  router: {
-                    route: pathname,
-                    asPath: pathname,
-                    pathname: pathname,
-                    query,
-                  }
-                })
-              )
-            )
           }`
-              : ''
+              : `
+          const renderServerComponentData = false`
           }
 
           const renderOpts = {
@@ -156,7 +105,14 @@ export default async function middlewareRSCLoader(this: any) {
             basePath: ${JSON.stringify(basePath || '')},
             supportsDynamicHTML: true,
             concurrentFeatures: true,
-            renderServerComponent: ${isServerComponent ? 'true' : 'false'},
+            renderServerComponentData,
+            renderServerComponent: ${
+              !isServerComponent
+                ? 'null'
+                : `
+              props => renderToReadableStream(createElement(Page, props), rscManifest)
+            `
+            },
           }
 
           const transformStream = new TransformStream()
@@ -173,7 +129,8 @@ export default async function middlewareRSCLoader(this: any) {
             )
             result.pipe({
               write: str => writer.write(encoder.encode(str)),
-              end: () => writer.close()
+              end: () => writer.close(),
+              // Not implemented: cork/uncork/on/removeListener
             })
           } catch (err) {
             return new Response(

packages/next/server/render.tsx

@@ -3,6 +3,7 @@ import { ParsedUrlQuery } from 'querystring'
 import type { Writable as WritableType } from 'stream'
 import React from 'react'
 import ReactDOMServer from 'react-dom/server'
+import { createFromReadableStream } from 'next/dist/compiled/react-server-dom-webpack'
 import { StyleRegistry, createStyleRegistry } from 'styled-jsx'
 import { UnwrapPromise } from '../lib/coalesced-function'
 import {
@@ -203,7 +204,8 @@ export type RenderOptsPartial = {
   devOnlyCacheBusterQueryString?: string
   resolvedUrl?: string
   resolvedAsPath?: string
-  renderServerComponent?: null | (() => Promise<string>)
+  renderServerComponent?: null | ((props: any) => ReadableStream)
+  renderServerComponentData?: boolean
   distDir?: string
   locale?: string
   locales?: string[]
@@ -274,6 +276,46 @@ function checkRedirectValues(
   }
 }
 
+// Create the wrapper component for a Flight stream.
+function createServerComponentRenderer(
+  OriginalComponent: React.ComponentType,
+  renderServerComponent: NonNullable<RenderOpts['renderServerComponent']>
+) {
+  let responseCache: any
+  const ServerComponentWrapper = (props: any) => {
+    let response = responseCache
+    if (!response) {
+      responseCache = response = createFromReadableStream(
+        renderServerComponent(props)
+      )
+    }
+    return response.readRoot()
+  }
+  const Component = (props: any) => {
+    return (
+      <React.Suspense fallback={null}>
+        <ServerComponentWrapper {...props} />
+      </React.Suspense>
+    )
+  }
+
+  // Although it's not allowed to attach some static methods to Component,
+  // we still re-assign all the component APIs to keep the behavior unchanged.
+  for (const methodName of [
+    'getInitialProps',
+    'getStaticProps',
+    'getServerSideProps',
+    'getStaticPaths',
+  ]) {
+    const method = (OriginalComponent as any)[methodName]
+    if (method) {
+      ;(Component as any)[methodName] = method
+    }
+  }
+
+  return Component
+}
+
 export async function renderToHTML(
   req: IncomingMessage,
   res: ServerResponse,
@@ -298,7 +340,6 @@ export async function renderToHTML(
     App,
     Document,
     pageConfig = {},
-    Component,
     buildManifest,
     fontManifest,
     reactLoadableManifest,
@@ -307,6 +348,7 @@ export async function renderToHTML(
     getStaticPaths,
     getServerSideProps,
     renderServerComponent,
+    renderServerComponentData,
     isDataReq,
     params,
     previewProps,
@@ -316,6 +358,11 @@ export async function renderToHTML(
     concurrentFeatures,
   } = renderOpts
 
+  const isServerComponent = !!renderServerComponent
+  const Component = isServerComponent
+    ? createServerComponentRenderer(renderOpts.Component, renderServerComponent)
+    : renderOpts.Component
+
   const getFontDefinition = (url: string): string => {
     if (fontManifest) {
       return getFontDefinitionFromManifest(url, fontManifest)
@@ -359,8 +406,6 @@ export async function renderToHTML(
 
   const hasPageGetInitialProps = !!(Component as any).getInitialProps
 
-  const isRSC = !!renderServerComponent
-
   const pageIsDynamic = isDynamicRoute(pathname)
 
   const isAutoExport =
@@ -775,7 +820,9 @@ export async function renderToHTML(
     props.pageProps = Object.assign(
       {},
       props.pageProps,
-      'props' in data ? data.props : undefined
+      'props' in data ? data.props : undefined,
+      // Pass router to the Server Component as a temporary workaround.
+      isServerComponent ? { router } : undefined
     )
 
     // pass up revalidate and props for export
@@ -943,6 +990,16 @@ export async function renderToHTML(
   // the response might be finished on the getInitialProps call
   if (isResSent(res) && !isSSG) return null
 
+  if (renderServerComponentData) {
+    return new RenderResult((res, next) => {
+      const { startWriting } = connectReactServerReadableStreamToPiper(
+        res.write,
+        next
+      )
+      startWriting(renderServerComponent!(props))
+    })
+  }
+
   // we preload the buildManifest for auto-export dynamic pages
   // to speed up hydrating query values
   let filteredBuildManifest = buildManifest
@@ -1081,7 +1138,7 @@ export async function renderToHTML(
 
         return concurrentFeatures
           ? process.browser
-            ? await renderToReadableStream(content)
+            ? await renderToWebStream(content)
             : await renderToNodeStream(content, generateStaticHTML)
           : piperFromArray([ReactDOMServer.renderToString(content)])
       }
@@ -1154,7 +1211,7 @@ export async function renderToHTML(
       err: renderOpts.err ? serializeError(dev, renderOpts.err) : undefined, // Error if one happened, otherwise don't sent in the resulting HTML
       gsp: !!getStaticProps ? true : undefined, // whether the page is getStaticProps
       gssp: !!getServerSideProps ? true : undefined, // whether the page is getServerSideProps
-      rsc: isRSC ? true : undefined, // whether the page is a server components page
+      rsc: isServerComponent ? true : undefined, // whether the page is a server components page
       customServer, // whether the user is using a custom server
       gip: hasPageGetInitialProps ? true : undefined, // whether the page has getInitialProps
       appGip: !defaultAppGetInitialProps ? true : undefined, // whether the _app has getInitialProps
@@ -1467,41 +1524,71 @@ function renderToNodeStream(
   })
 }
 
-function renderToReadableStream(
-  element: React.ReactElement
-): NodeWritablePiper {
-  return (res, next) => {
-    let bufferedString = ''
-    let shellCompleted = false
+function connectReactServerReadableStreamToPiper(
+  write: (s: string) => boolean,
+  next: (err?: Error) => void
+) {
+  let bufferedString = ''
+
+  function flushBuffer() {
+    // Intentionally delayed writing when using ReadableStream due to the lack
+    // of cork/uncork APIs.
+    setTimeout(() => {
+      if (!bufferedString) return
+      if (write(bufferedString)) {
+        bufferedString = ''
+      }
+    }, 0)
+  }
 
-    const readable = (ReactDOMServer as any).renderToReadableStream(element, {
-      onCompleteShell() {
-        shellCompleted = true
-        if (bufferedString) {
-          res.write(bufferedString)
-          bufferedString = ''
-        }
-      },
-    })
+  function startWriting(readable: ReadableStream) {
     const reader = readable.getReader()
     const decoder = new TextDecoder()
     const process = () => {
       reader.read().then(({ done, value }: any) => {
         if (!done) {
           const s = typeof value === 'string' ? value : decoder.decode(value)
-          if (shellCompleted) {
-            res.write(s)
-          } else {
-            bufferedString += s
-          }
+          bufferedString += s
+          flushBuffer()
           process()
         } else {
-          next()
+          // Make sure it's scheduled after the current flushing.
+          setTimeout(() => next(), 0)
         }
       })
     }
     process()
   }
+
+  return {
+    flushBuffer,
+    startWriting,
+  }
+}
+
+function renderToWebStream(element: React.ReactElement): NodeWritablePiper {
+  return (res, next) => {
+    let shellCompleted = false
+
+    const { flushBuffer, startWriting } =
+      connectReactServerReadableStreamToPiper((s: string) => {
+        // Buffer result until the shell is completed.
+        if (shellCompleted) {
+          res.write(s)
+          return true
+        }
+        return false
+      }, next)
+
+    startWriting(
+      (ReactDOMServer as any).renderToReadableStream(element, {
+        onCompleteShell() {
+          shellCompleted = true
+          flushBuffer()
+        },
+      })
+    )
+  }
 }
 
 function chainPipers(pipers: NodeWritablePiper[]): NodeWritablePiper {

packages/next/types/misc.d.ts

@@ -1,6 +1,7 @@
 /* eslint-disable import/no-extraneous-dependencies */
 declare module 'next/dist/compiled/babel/plugin-transform-modules-commonjs'
 declare module 'next/dist/compiled/babel/plugin-syntax-jsx'
+declare module 'next/dist/compiled/react-server-dom-webpack'
 declare module 'browserslist'
 
 declare module 'cssnano-simple' {