From 2412a0f406a180f7829078b5a1b4d3696e322c5a Mon Sep 17 00:00:00 2001
From: Ethan Arrowood <ethan@arrowood.dev>
Date: Thu, 23 May 2024 10:44:26 -0600
Subject: [PATCH] convert use-flight-response and try to tee Readable

---
 .../next/src/server/app-render/app-render.tsx |  86 +++++++---
 .../app-render/use-flight-response/index.d.ts |  19 +++
 .../app-render/use-flight-response/index.js   |   8 +
 .../use-flight-response.edge.ts}              |  26 ++-
 .../use-flight-response.node.ts               | 152 ++++++++++++++++++
 .../server/stream-utils/stream-utils.node.ts  |  58 ++-----
 6 files changed, 266 insertions(+), 83 deletions(-)
 create mode 100644 packages/next/src/server/app-render/use-flight-response/index.d.ts
 create mode 100644 packages/next/src/server/app-render/use-flight-response/index.js
 rename packages/next/src/server/app-render/{use-flight-response.tsx => use-flight-response/use-flight-response.edge.ts} (84%)
 create mode 100644 packages/next/src/server/app-render/use-flight-response/use-flight-response.node.ts

diff --git a/packages/next/src/server/app-render/app-render.tsx b/packages/next/src/server/app-render/app-render.tsx
index d794b80acf394..654658e38bb12 100644
--- a/packages/next/src/server/app-render/app-render.tsx
+++ b/packages/next/src/server/app-render/app-render.tsx
@@ -624,7 +624,7 @@ function ReactServerEntrypoint<T>({
   nonce?: string
 }): T {
   preinitScripts()
-  const response = useFlightStream(
+  const response = useFlightStream<T>(
     reactServerStream,
     clientReferenceManifest,
     nonce
@@ -979,27 +979,38 @@ async function renderToHTMLOrFlightImpl(
         }
       )
 
-      let resultStream: ReadableStream<Uint8Array>
+      // let resultStream: ReadableStream<Uint8Array>
+      // if (
+      //   process.env.NEXT_RUNTIME === 'nodejs' &&
+      //   !(serverStream instanceof ReadableStream)
+      // ) {
+      //   const { PassThrough, Readable } =
+      //     require('node:stream') as typeof import('node:stream')
+      //   resultStream = Readable.toWeb(
+      //     serverStream.pipe(new PassThrough())
+      //   ) as ReadableStream<Uint8Array>
+      // } else if (!(serverStream instanceof ReadableStream)) {
+      //   throw new Error(
+      //     'Invariant. Stream is not a ReadableStream in non-Node.js runtime'
+      //   )
+      // } else {
+      //   resultStream = serverStream
+      // }
+
+      let renderStream, dataStream
+
       if (
         process.env.NEXT_RUNTIME === 'nodejs' &&
         !(serverStream instanceof ReadableStream)
       ) {
-        const { PassThrough, Readable } =
-          require('node:stream') as typeof import('node:stream')
-        resultStream = Readable.toWeb(
-          serverStream.pipe(new PassThrough())
-        ) as ReadableStream<Uint8Array>
-      } else if (!(serverStream instanceof ReadableStream)) {
-        throw new Error(
-          'Invariant. Stream is not a ReadableStream in non-Node.js runtime'
-        )
+        const { teeReadable } = require('../stream-utils')
+        ;[renderStream, dataStream] = teeReadable(serverStream)
       } else {
-        resultStream = serverStream
+        // We are going to consume this render both for SSR and for inlining the flight data
+        // @ts-ignore
+        ;[renderStream, dataStream] = serverStream.tee()
       }
 
-      // We are going to consume this render both for SSR and for inlining the flight data
-      let [renderStream, dataStream] = resultStream.tee()
-
       const children = (
         <HeadManagerContext.Provider
           value={{
@@ -1110,7 +1121,20 @@ async function renderToHTMLOrFlightImpl(
           } else {
             // We may still be rendering the RSC stream even though the HTML is finished.
             // We wait for the RSC stream to complete and check again if dynamic was used
-            const [original, flightSpy] = dataStream.tee()
+            let original, flightSpy
+
+            if (
+              process.env.NEXT_RUNTIME === 'nodejs' &&
+              !(dataStream instanceof ReadableStream)
+            ) {
+              const { teeReadable } = require('../stream-utils')
+              ;[original, flightSpy] = teeReadable(dataStream)
+            } else {
+              // We are going to consume this render both for SSR and for inlining the flight data
+              // @ts-ignore
+              ;[original, flightSpy] = dataStream.tee()
+            }
+
             dataStream = original
 
             await flightRenderComplete(flightSpy)
@@ -1198,13 +1222,23 @@ async function renderToHTMLOrFlightImpl(
                 renderedHTMLStream = convertReadable(resultStream2)
               }
 
+              let inlinedDataStream = createInlinedDataReadableStream(
+                dataStream,
+                nonce,
+                formState
+              )
+
+              if (
+                process.env.NEXT_RUNTIME === 'nodejs' &&
+                !(inlinedDataStream instanceof ReadableStream)
+              ) {
+                inlinedDataStream = convertReadable(inlinedDataStream)
+              }
+
               return {
                 stream: await continueStaticPrerender(renderedHTMLStream, {
-                  inlinedDataStream: createInlinedDataReadableStream(
-                    dataStream,
-                    nonce,
-                    formState
-                  ),
+                  // @ts-ignore
+                  inlinedDataStream,
                   getServerInsertedHTML,
                 }),
               }
@@ -1213,15 +1247,22 @@ async function renderToHTMLOrFlightImpl(
         } else if (renderOpts.postponed) {
           stream = convertReadable(stream)
           // This is a continuation of either an Incomplete or Dynamic Data Prerender.
-          const inlinedDataStream = createInlinedDataReadableStream(
+          let inlinedDataStream = createInlinedDataReadableStream(
             dataStream,
             nonce,
             formState
           )
+          if (
+            process.env.NEXT_RUNTIME === 'nodejs' &&
+            !(inlinedDataStream instanceof ReadableStream)
+          ) {
+            inlinedDataStream = convertReadable(inlinedDataStream)
+          }
           if (resumed) {
             // We have new HTML to stream and we also need to include server inserted HTML
             return {
               stream: await continueDynamicHTMLResume(stream, {
+                // @ts-ignore
                 inlinedDataStream,
                 getServerInsertedHTML,
               }),
@@ -1230,6 +1271,7 @@ async function renderToHTMLOrFlightImpl(
             // We are continuing a Dynamic Data Prerender and simply need to append new inlined flight data
             return {
               stream: await continueDynamicDataResume(stream, {
+                // @ts-ignore
                 inlinedDataStream,
               }),
             }
diff --git a/packages/next/src/server/app-render/use-flight-response/index.d.ts b/packages/next/src/server/app-render/use-flight-response/index.d.ts
new file mode 100644
index 0000000000000..0ec228e177b3c
--- /dev/null
+++ b/packages/next/src/server/app-render/use-flight-response/index.d.ts
@@ -0,0 +1,19 @@
+import type { Readable } from 'node:stream'
+import type { DeepReadonly } from '../../../shared/lib/deep-readonly'
+import type { ClientReferenceManifest } from '../../../build/webpack/plugins/flight-manifest-plugin'
+
+export function useFlightStream<T>(
+  flightStream: Readable | ReadableStream<Uint8Array>,
+  clientReferenceManifest: DeepReadonly<ClientReferenceManifest>,
+  nonce?: string
+): Promise<T>
+
+export function flightRenderComplete(
+  flightStream: Readable | ReadableStream<Uint8Array>
+): Promise<void>
+
+export function createInlinedDataReadableStream(
+  flightStream: Readable | ReadableStream<Uint8Array>,
+  nonce: string | undefined,
+  formState: unknown | null
+): Readable | ReadableStream<Uint8Array>
diff --git a/packages/next/src/server/app-render/use-flight-response/index.js b/packages/next/src/server/app-render/use-flight-response/index.js
new file mode 100644
index 0000000000000..2fd6fc6598200
--- /dev/null
+++ b/packages/next/src/server/app-render/use-flight-response/index.js
@@ -0,0 +1,8 @@
+if (
+  process.env.NEXT_RUNTIME === 'nodejs' &&
+  process.env.EXPERIMENTAL_NODE_STREAMS_SUPPORT === '1'
+) {
+  module.exports = require('./use-flight-response.node.js')
+} else {
+  module.exports = require('./use-flight-response.edge.js')
+}
diff --git a/packages/next/src/server/app-render/use-flight-response.tsx b/packages/next/src/server/app-render/use-flight-response/use-flight-response.edge.ts
similarity index 84%
rename from packages/next/src/server/app-render/use-flight-response.tsx
rename to packages/next/src/server/app-render/use-flight-response/use-flight-response.edge.ts
index 5d31ab6b0e52d..5792b247a46a4 100644
--- a/packages/next/src/server/app-render/use-flight-response.tsx
+++ b/packages/next/src/server/app-render/use-flight-response/use-flight-response.edge.ts
@@ -1,8 +1,8 @@
-import type { ClientReferenceManifest } from '../../build/webpack/plugins/flight-manifest-plugin'
-import type { BinaryStreamOf } from './app-render'
+import type { ClientReferenceManifest } from '../../../build/webpack/plugins/flight-manifest-plugin'
+import type { BinaryStreamOf } from '../app-render'
 
-import { htmlEscapeJsonString } from '../htmlescape'
-import type { DeepReadonly } from '../../shared/lib/deep-readonly'
+import { htmlEscapeJsonString } from '../../htmlescape'
+import type { DeepReadonly } from '../../../shared/lib/deep-readonly'
 import type { Readable } from 'node:stream'
 
 const isEdgeRuntime = process.env.NEXT_RUNTIME === 'edge'
@@ -23,7 +23,7 @@ const encoder = new TextEncoder()
  * This is only used for renderToHTML, the Flight response does not need additional wrappers.
  */
 export function useFlightStream<T>(
-  flightStream: Readable | BinaryStreamOf<T>,
+  flightStream: BinaryStreamOf<T>,
   clientReferenceManifest: DeepReadonly<ClientReferenceManifest>,
   nonce?: string
 ): Promise<T> {
@@ -38,20 +38,12 @@ export function useFlightStream<T>(
   // @TODO: investigate why the aliasing for turbopack doesn't pick this up, requiring this runtime check
   if (process.env.TURBOPACK) {
     createFromStream =
-      flightStream instanceof ReadableStream
-        ? // eslint-disable-next-line import/no-extraneous-dependencies
-          require('react-server-dom-turbopack/client.edge')
-            .createFromReadableStream
-        : // eslint-disable-next-line import/no-extraneous-dependencies
-          require('react-server-dom-turbopack/client.node').createFromNodeStream
+      // eslint-disable-next-line import/no-extraneous-dependencies
+      require('react-server-dom-turbopack/client.edge').createFromReadableStream
   } else {
     createFromStream =
-      flightStream instanceof ReadableStream
-        ? // eslint-disable-next-line import/no-extraneous-dependencies
-          require('react-server-dom-webpack/client.edge')
-            .createFromReadableStream
-        : // eslint-disable-next-line import/no-extraneous-dependencies
-          require('react-server-dom-webpack/client.node').createFromNodeStream
+      // eslint-disable-next-line import/no-extraneous-dependencies
+      require('react-server-dom-webpack/client.edge').createFromReadableStream
   }
 
   const newResponse = createFromStream(flightStream, {
diff --git a/packages/next/src/server/app-render/use-flight-response/use-flight-response.node.ts b/packages/next/src/server/app-render/use-flight-response/use-flight-response.node.ts
new file mode 100644
index 0000000000000..07590cc3521d2
--- /dev/null
+++ b/packages/next/src/server/app-render/use-flight-response/use-flight-response.node.ts
@@ -0,0 +1,152 @@
+import { Readable } from 'node:stream'
+import type { ClientReferenceManifest } from '../../../build/webpack/plugins/flight-manifest-plugin'
+import type { DeepReadonly } from '../../../shared/lib/deep-readonly'
+import { htmlEscapeJsonString } from '../../htmlescape'
+import isError from '../../../lib/is-error'
+
+const flightResponses = new WeakMap<Readable, Promise<any>>()
+const encoder = new TextEncoder()
+
+export function useFlightStream<T>(
+  flightStream: Readable,
+  clientReferenceManifest: DeepReadonly<ClientReferenceManifest>,
+  nonce?: string
+): Promise<T> {
+  const response = flightResponses.get(flightStream)
+
+  if (response) {
+    return response
+  }
+
+  // react-server-dom-webpack/client.edge must not be hoisted for require cache clearing to work correctly
+  let createFromStream
+  // @TODO: investigate why the aliasing for turbopack doesn't pick this up, requiring this runtime check
+  if (process.env.TURBOPACK) {
+    createFromStream =
+      // eslint-disable-next-line import/no-extraneous-dependencies
+      require('react-server-dom-turbopack/client.node').createFromNodeStream
+  } else {
+    createFromStream =
+      // eslint-disable-next-line import/no-extraneous-dependencies
+      require('react-server-dom-webpack/client.node').createFromNodeStream
+  }
+
+  const newResponse = createFromStream(flightStream, {
+    ssrManifest: {
+      moduleLoading: clientReferenceManifest.moduleLoading,
+      moduleMap: clientReferenceManifest.ssrModuleMapping,
+    },
+    nonce,
+  })
+
+  flightResponses.set(flightStream, newResponse)
+
+  return newResponse
+}
+
+export function flightRenderComplete(flightStream: Readable): Promise<void> {
+  return new Promise((resolve, reject) => {
+    flightStream
+      .resume()
+      .on('end', () => {
+        resolve()
+      })
+      .on('error', (error) => {
+        reject(error)
+      })
+  })
+}
+
+const INLINE_FLIGHT_PAYLOAD_BOOTSTRAP = 0
+const INLINE_FLIGHT_PAYLOAD_DATA = 1
+const INLINE_FLIGHT_PAYLOAD_FORM_STATE = 2
+const INLINE_FLIGHT_PAYLOAD_BINARY = 3
+
+export function createInlinedDataReadableStream(
+  flightStream: Readable,
+  nonce: string | undefined,
+  formState: unknown | null
+): Readable {
+  const startScriptTag = nonce
+    ? `<script nonce=${JSON.stringify(nonce)}>`
+    : '<script>'
+
+  const decoder = new TextDecoder('utf-8', { fatal: true })
+
+  if (flightStream.readableFlowing) {
+    flightStream.pause()
+  }
+
+  return new Readable({
+    construct(callback) {
+      try {
+        const chunk = encoder.encode(
+          `${startScriptTag}(self.__next_f=self.__next_f||[]).push(${htmlEscapeJsonString(
+            JSON.stringify([INLINE_FLIGHT_PAYLOAD_BOOTSTRAP])
+          )});self.__next_f.push(${htmlEscapeJsonString(
+            JSON.stringify([INLINE_FLIGHT_PAYLOAD_FORM_STATE, formState])
+          )})</script>`
+        )
+        this.push(chunk)
+        return callback(null)
+      } catch (error) {
+        return isError(error) ? callback(error) : callback()
+      }
+    },
+    read() {
+      try {
+        const chunk = flightStream.read()
+        if (chunk) {
+          try {
+            const decodedString = decoder.decode(chunk, { stream: true })
+            writeFlightDataInstruction(this, startScriptTag, decodedString)
+          } catch {
+            writeFlightDataInstruction(this, startScriptTag, chunk)
+          }
+        } else {
+          try {
+            const decodedString = decoder.decode()
+            if (decodedString) {
+              writeFlightDataInstruction(this, startScriptTag, decodedString)
+            }
+          } catch {}
+
+          this.push(null)
+        }
+      } catch (error) {
+        if (isError(error)) {
+          this.destroy(error)
+        }
+      }
+    },
+  })
+}
+
+function writeFlightDataInstruction(
+  readable: Readable,
+  scriptStart: string,
+  chunk: string | Uint8Array
+) {
+  let htmlInlinedData: string
+
+  if (typeof chunk === 'string') {
+    htmlInlinedData = htmlEscapeJsonString(
+      JSON.stringify([INLINE_FLIGHT_PAYLOAD_DATA, chunk])
+    )
+  } else {
+    // The chunk cannot be embedded as a UTF-8 string in the script tag.
+    // Instead let's inline it in base64.
+    // Credits to Devon Govett (devongovett) for the technique.
+    // https://github.com/devongovett/rsc-html-stream
+    const base64 = btoa(String.fromCodePoint(...chunk))
+    htmlInlinedData = htmlEscapeJsonString(
+      JSON.stringify([INLINE_FLIGHT_PAYLOAD_BINARY, base64])
+    )
+  }
+
+  readable.push(
+    encoder.encode(
+      `${scriptStart}self.__next_f.push(${htmlInlinedData})</script>`
+    )
+  )
+}
diff --git a/packages/next/src/server/stream-utils/stream-utils.node.ts b/packages/next/src/server/stream-utils/stream-utils.node.ts
index 7da3e21f8117d..107b3d9456afb 100644
--- a/packages/next/src/server/stream-utils/stream-utils.node.ts
+++ b/packages/next/src/server/stream-utils/stream-utils.node.ts
@@ -10,7 +10,10 @@ import {
   Writable,
   pipeline,
 } from 'node:stream'
-import type { Options as RenderToPipeableStreamOptions } from 'react-dom/server.node'
+import type {
+  PipeableStream,
+  Options as RenderToPipeableStreamOptions,
+} from 'react-dom/server.node'
 import isError from '../../lib/is-error'
 import {
   indexOfUint8Array,
@@ -18,7 +21,6 @@ import {
   removeFromUint8Array,
 } from './uint8array-helpers'
 import { ENCODED_TAGS } from './encodedTags'
-import { createMergedTransformStream } from './stream-utils.edge'
 
 export * from './stream-utils.edge'
 
@@ -414,7 +416,7 @@ export function continueFizzStream(
     serverInsertedHTMLToHead,
     validateRootLayout,
   }: {
-    inlinedDataStream?: ReadableStream<Uint8Array>
+    inlinedDataStream?: Readable
     isStaticGeneration: boolean
     getServerInsertedHTML?: () => Promise<string>
     serverInsertedHTMLToHead: boolean
@@ -443,8 +445,7 @@ export function continueFizzStream(
   }
 
   if (inlinedDataStream) {
-    // @ts-ignore
-    streams.push(createMergedTransformStream(inlinedDataStream))
+    streams.push(inlinedDataStream.pipe(new PassThrough()))
   }
 
   if (validateRootLayout) {
@@ -476,42 +477,11 @@ export function convertReadable(
     : stream
 }
 
-// export function continueDynamicPrerender(prerenderStream: Readable, { getServerInsertedHTML }: { getServerInsertedHTML: ()=> Promise<string>}) {
-//   const pt = new PassThrough();
-
-//   return new Promise((resolve, reject) => {
-//     // @ts-expect-error
-//     pipeline([
-//       prerenderStream,
-//       createBufferedTransformStream(),
-//       createStripDocumentClosingTagsTransform(),
-//       createHeadInsertionTransformStream(getServerInsertedHTML),
-//       pt
-//     ], (error) => {
-//       if (error) return reject(error)
-//         else return resolve(pt)
-//     })
-//   })
-// }
-
-// export function continueStaticPrerender(
-//   prerenderStream: Readable,
-//   {
-//     inlinedDataStream,
-//     getServerInsertedHTML
-//   }: {
-//     inlinedDataStream: Readable,
-//     getServerInsertedHTML: () => Promise<string>
-//   }
-// ) {
-//   const pt = new PassThrough();
-//   return new Promise((resolve, reject) => {
-//     pipeline(
-//       prerenderStream,
-//       createBufferedTransformStream(),
-//       createHeadInsertionTransformStream(getServerInsertedHTML),
-//       inlinedDataStream,
-//       createMoveSuffixStream('</body></html>')
-//     )
-//   })
-// }
+export function teeReadable(
+  stream: PipeableStream | Readable
+): [Readable, Readable] {
+  if (!(stream instanceof Readable)) {
+    stream = stream.pipe(new PassThrough())
+  }
+  return [stream.pipe(new PassThrough()), stream.pipe(new PassThrough())]
+}