Fix broken html on streaming render for error page

packages/next/server/render.tsx

@@ -1551,6 +1551,7 @@ function renderToNodeStream(
 
     // Based on the suggestion here:
     // https://github.com/reactwg/react-18/discussions/110
+    let suffixFlushed = false
     class NextWritable extends Writable {
       _write(
         chunk: any,
@@ -1572,7 +1573,15 @@ function renderToNodeStream(
 
         if (!shellFlushed) {
           shellFlushed = true
-          underlyingStream.write(suffixUnclosed)
+          // In the first round of streaming, all chunks will be finished in the micro task.
+          // We use setTimeout to guarantee the suffix is flushed after the micro task.
+          setTimeout(() => {
+            // Flush the suffix if stream is not closed.
+            if (underlyingStream) {
+              suffixFlushed = true
+              underlyingStream.write(suffixUnclosed)
+            }
+          })
         }
       }
 
@@ -1603,6 +1612,11 @@ function renderToNodeStream(
         resolved = true
         resolve((res, next) => {
           const doNext = (err?: Error) => {
+            // Some cases when the stream is closed too fast before setTimeout,
+            // have to ensure suffix is flushed anyway.
+            if (!suffixFlushed) {
+              res.write(suffixUnclosed)
+            }
             if (!err) {
               res.write(closeTag)
             }
@@ -1631,6 +1645,7 @@ function renderToNodeStream(
           abort()
         },
         onCompleteShell() {
+          shellFlushed = true
           if (!generateStaticHTML) {
             doResolve(() => pipe(stream))
           }

test/integration/react-streaming-and-server-components/test/index.test.js

@@ -27,6 +27,7 @@ const documentPage = new File(join(appDir, 'pages/_document.jsx'))
 const appPage = new File(join(appDir, 'pages/_app.js'))
 const appServerPage = new File(join(appDir, 'pages/_app.server.js'))
 const error500Page = new File(join(appDir, 'pages/500.js'))
+const error404Page = new File(join(appDir, 'pages/404.js'))
 
 const documentWithGip = `
 import { Html, Head, Main, NextScript } from 'next/document'
@@ -71,6 +72,33 @@ export default function Page500() {
 }
 `
 
+const suspense404 = `
+import { Suspense } from 'react'
+
+let result
+let promise
+function Data() {
+  if (result) return result
+  if (!promise)
+    promise = new Promise((res) => {
+      setTimeout(() => {
+        result = 'next_streaming_data'
+        res()
+      }, 500)
+    })
+  throw promise
+}
+
+export default function Page404() {
+  return (
+    <Suspense fallback={null}>
+      custom-404-page
+      <Data />
+    </Suspense>
+  )
+}
+`
+
 async function nextBuild(dir) {
   return await _nextBuild(dir, [], {
     stdout: true,
@@ -118,12 +146,27 @@ describe('concurrentFeatures - basic', () => {
     expect(stderr).toContain(edgeRuntimeWarning)
     expect(stderr).toContain(rscWarning)
   })
+
   it('should warn user that native node APIs are not supported', async () => {
     const fsImportedErrorMessage =
       'Native Node.js APIs are not supported in the Edge Runtime with `concurrentFeatures` enabled. Found `dns` imported.'
     const { stderr } = await nextBuild(nativeModuleTestAppDir)
     expect(stderr).toContain(fsImportedErrorMessage)
   })
+
+  it('should handle suspense error page correctly (node stream)', async () => {
+    error404Page.write(suspense404)
+    const appPort = await findPort()
+    await nextBuild(appDir)
+    await nextStart(appDir, appPort)
+    const browser = await webdriver(appPort, '/404')
+    const hydrationContent = await browser.eval(
+      `document.querySelector('#__next').textContent`
+    )
+    expect(hydrationContent).toBe('custom-404-pagenext_streaming_data')
+
+    error404Page.restore()
+  })
 })
 
 describe('concurrentFeatures - prod', () => {
@@ -266,7 +309,7 @@ runSuite('document', 'prod', documentSuite)
 
 async function runBasicTests(context, env) {
   const isDev = env === 'dev'
-  it('should render the correct html', async () => {
+  it('should render html correctly', async () => {
     const homeHTML = await renderViaHTTP(context.appPort, '/', null, {
       headers: {
         'x-next-test-client': 'test-util',
@@ -293,6 +336,8 @@ async function runBasicTests(context, env) {
       '/this-is-not-found'
     )
 
+    const page404Content = 'custom-404-page'
+
     expect(homeHTML).toContain('component:index.server')
     expect(homeHTML).toContain('env:env_var_test')
     expect(homeHTML).toContain('header:test-util')
@@ -302,12 +347,14 @@ async function runBasicTests(context, env) {
     expect(dynamicRouteHTML1).toContain('[pid]')
     expect(dynamicRouteHTML2).toContain('[pid]')
 
-    expect(path404HTML).toContain('custom-404-page')
+    const $404 = cheerio.load(path404HTML)
+    expect($404('#__next').text()).toBe(page404Content)
+
     // in dev mode: custom error page is still using default _error
     expect(path500HTML).toContain(
       isDev ? 'Internal Server Error' : 'custom-500-page'
     )
-    expect(pathNotFoundHTML).toContain('custom-404-page')
+    expect(pathNotFoundHTML).toContain(page404Content)
   })
 
   it('should support next/link', async () => {