Simplify the html insertion html stream handling for ppr

huozhi · huozhi · commit 621754015968 · 2025-02-03T16:22:24.000+01:00
diff --git a/packages/next/src/server/stream-utils/node-web-streams-helper.ts b/packages/next/src/server/stream-utils/node-web-streams-helper.ts
@@ -196,7 +196,6 @@ function createHeadInsertionTransformStream(
   insert: () => Promise<string>
 ): TransformStream<Uint8Array, Uint8Array> {
   let inserted = false
-  let freezing = false
 
   // We need to track if this transform saw any bytes because if it didn't
   // we won't want to insert any server HTML at all
@@ -205,32 +204,35 @@ function createHeadInsertionTransformStream(
   return new TransformStream({
     async transform(chunk, controller) {
       hasBytes = true
-      // While react is flushing chunks, we don't apply insertions
-      if (freezing) {
-        controller.enqueue(chunk)
-        return
-      }
 
       const insertion = await insert()
-
       if (inserted) {
         if (insertion) {
           const encodedInsertion = encoder.encode(insertion)
           controller.enqueue(encodedInsertion)
         }
         controller.enqueue(chunk)
-        freezing = true
       } else {
         // TODO (@Ethan-Arrowood): Replace the generic `indexOfUint8Array` method with something finely tuned for the subset of things actually being checked for.
         const index = indexOfUint8Array(chunk, ENCODED_TAGS.CLOSED.HEAD)
+        // In fully static rendering or non PPR rendering cases:
+        // `/head>` will always be found in the chunk in first chunk rendering.
         if (index !== -1) {
           if (insertion) {
             const encodedInsertion = encoder.encode(insertion)
+            // Get the total count of the bytes in the chunk and the insertion
+            // e.g.
+            // chunk = <head><meta charset="utf-8"></head>
+            // insertion = <script>...</script>
+            // output = <head><meta charset="utf-8"> [ <script>...</script> ] </head>
             const insertedHeadContent = new Uint8Array(
               chunk.length + encodedInsertion.length
             )
+            // Append the first part of the chunk, before the head tag
             insertedHeadContent.set(chunk.slice(0, index))
+            // Append the server inserted content
             insertedHeadContent.set(encodedInsertion, index)
+            // Append the rest of the chunk
             insertedHeadContent.set(
               chunk.slice(index),
               index + encodedInsertion.length
@@ -239,18 +241,21 @@ function createHeadInsertionTransformStream(
           } else {
             controller.enqueue(chunk)
           }
-          freezing = true
+          inserted = true
+        } else {
+          // This will happens in PPR rendering during next start, when the page is partially rendered.
+          // When the page resumes, the head tag will be found in the middle of the chunk.
+          // Where we just need to append the insertion and chunk to the current stream.
+          // e.g.
+          // PPR-static: <head>...</head><body> [ resume content ] </body>
+          // PPR-resume: [ insertion ] [ rest content ]
+          if (insertion) {
+            controller.enqueue(encoder.encode(insertion))
+          }
+          controller.enqueue(chunk)
           inserted = true
         }
       }
-
-      if (!inserted) {
-        controller.enqueue(chunk)
-      } else {
-        scheduleImmediate(() => {
-          freezing = false
-        })
-      }
     },
     async flush(controller) {
       // Check before closing if there's anything remaining to insert.
diff --git a/test/production/app-dir/ppr-use-server-inserted-html/app/layout.tsx b/test/production/app-dir/ppr-use-server-inserted-html/app/layout.tsx
@@ -0,0 +1,13 @@
+import React from 'react'
+
+export default async function Root({
+  children,
+}: {
+  children: React.ReactNode
+}) {
+  return (
+    <html>
+      <body>{children}</body>
+    </html>
+  )
+}
diff --git a/test/production/app-dir/ppr-use-server-inserted-html/app/partial-resume/client.tsx b/test/production/app-dir/ppr-use-server-inserted-html/app/partial-resume/client.tsx
@@ -0,0 +1,24 @@
+'use client'
+
+import { useRef } from 'react'
+import { useServerInsertedHTML } from 'next/navigation'
+
+export function InsertHtml({ id, data }: { id: string; data: string }) {
+  const insertRef = useRef(false)
+  useServerInsertedHTML(() => {
+    // only insert the style tag once
+    if (insertRef.current) {
+      return
+    }
+    insertRef.current = true
+    const value = (
+      <style
+        data-test-id={id}
+      >{`[data-inserted-${data}] { content: ${data} }`}</style>
+    )
+    console.log(`testing-log-insertion:${data}`)
+    return value
+  })
+
+  return <div>Loaded: {data}</div>
+}
diff --git a/test/production/app-dir/ppr-use-server-inserted-html/app/partial-resume/page.tsx b/test/production/app-dir/ppr-use-server-inserted-html/app/partial-resume/page.tsx
@@ -0,0 +1,22 @@
+import React, { Suspense } from 'react'
+import { headers } from 'next/headers'
+import { InsertHtml } from './client'
+
+async function Dynamic() {
+  await headers()
+
+  return (
+    <div>
+      <h3>dynamic</h3>
+      <InsertHtml id={'inserted-html'} data={'dynamic-data'} />
+    </div>
+  )
+}
+
+export default function Page() {
+  return (
+    <Suspense fallback={<div>Loading...</div>}>
+      <Dynamic />
+    </Suspense>
+  )
+}
diff --git a/test/production/app-dir/ppr-use-server-inserted-html/next.config.js b/test/production/app-dir/ppr-use-server-inserted-html/next.config.js
@@ -0,0 +1,10 @@
+/**
+ * @type {import('next').NextConfig}
+ */
+const nextConfig = {
+  experimental: {
+    ppr: true,
+  },
+}
+
+module.exports = nextConfig
diff --git a/test/production/app-dir/ppr-use-server-inserted-html/ppr-use-server-inserted-html.test.ts b/test/production/app-dir/ppr-use-server-inserted-html/ppr-use-server-inserted-html.test.ts
@@ -0,0 +1,32 @@
+import { nextTestSetup } from 'e2e-utils'
+
+describe('ppr-use-server-inserted-html', () => {
+  const { next, isNextStart } = nextTestSetup({
+    files: __dirname,
+  })
+
+  if (isNextStart) {
+    it('should mark the route as ppr rendered', async () => {
+      const prerenderManifest = JSON.parse(
+        await next.readFile('.next/prerender-manifest.json')
+      )
+      expect(prerenderManifest.routes['/partial-resume'].renderingMode).toBe(
+        'PARTIALLY_STATIC'
+      )
+    })
+  }
+
+  it('should not log insertion in build', async () => {
+    const output = next.cliOutput
+    expect(output).not.toContain('testing-log-insertion:')
+  })
+
+  it('should insert the html insertion into html body', async () => {
+    const $ = await next.render$('/partial-resume')
+    const output = next.cliOutput
+    expect(output).toContain('testing-log-insertion:dynamic-data')
+
+    expect($('head [data-test-id]').length).toBe(0)
+    expect($('body [data-test-id]').length).toBe(1)
+  })
+})
