fix: ensure we use res.end() to flush buffers

vercel · Oct 14, 2023 · 2bdfb4a · 2bdfb4a
1 parent e297e07
commit 2bdfb4a
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 12 deletions.
diff --git a/packages/next/src/lib/download-swc.ts b/packages/next/src/lib/download-swc.ts
@@ -5,9 +5,8 @@ import tar from 'next/dist/compiled/tar'
 const { fetch } = require('next/dist/compiled/undici') as {
   fetch: typeof global.fetch
 }
-const { WritableStream } = require('node:stream/web') as {
-  WritableStream: typeof global.WritableStream
-}
+import { WritableStream } from 'next/dist/compiled/@edge-runtime/ponyfill'
+
 import { getRegistry } from './helpers/get-registry'
 import { getCacheDirectory } from './helpers/get-cache-directory'
 

diff --git a/packages/next/src/server/lib/incremental-cache/index.ts b/packages/next/src/server/lib/incremental-cache/index.ts
@@ -9,6 +9,7 @@ import FileSystemCache from './file-system-cache'
 import path from '../../../shared/lib/isomorphic/path'
 import { normalizePagePath } from '../../../shared/lib/page-path/normalize-page-path'
 
+import { WritableStream } from 'next/dist/compiled/@edge-runtime/ponyfill'
 import {
   CACHE_ONE_YEAR,
   NEXT_CACHE_REVALIDATED_TAGS_HEADER,

diff --git a/packages/next/src/server/pipe-readable.ts b/packages/next/src/server/pipe-readable.ts
@@ -1,5 +1,7 @@
 import type { ServerResponse } from 'node:http'
 
+import { WritableStream } from 'next/dist/compiled/@edge-runtime/ponyfill'
+
 export function isAbortError(e: any): e is Error & { name: 'AbortError' } {
   return e?.name === 'AbortError'
 }
@@ -10,12 +12,25 @@ function createWriterFromResponse(
 ): WritableStream<Uint8Array> {
   let started = false
   let finished = false
+  let aborted = false
+
+  controller.signal.addEventListener('abort', () => {
+    // If we've already finished, we don't need to do anything.
+    if (aborted) return
+
+    // We haven't finished writing, so signal to abort.
+    aborted = true
+    finished = true
+    res.end()
+  })
 
   function onClose() {
     // If we've already finished, we don't need to do anything.
     if (finished) return
 
     // We haven't finished writing, so signal to abort.
+    aborted = true
+    finished = true
     controller.abort()
   }
 
@@ -34,22 +49,30 @@ function createWriterFromResponse(
         res.flushHeaders()
       }
 
-      const ok = res.write(chunk)
-      if (!ok) {
-        // If the write returns false, it means there's some backpressure, so
-        // wait until it's streamed before continuing.
-        await new Promise((resolve) => {
-          res.once('drain', resolve)
-        })
+      try {
+        const ok = res.write(chunk)
+        if (!ok) {
+          // If the write returns false, it means there's some backpressure, so
+          // wait until it's streamed before continuing.
+          await new Promise((resolve) => {
+            res.once('drain', resolve)
+          })
+        }
+      } catch (err: any) {
+        res.end()
+
+        throw err
       }
     },
-    abort: (err) => {
+    abort: () => {
       finished = true
-      res.destroy(err)
+      res.end()
     },
     close: () => {
       finished = true
 
+      if (res.writableFinished) return
+
       // Create a promise that will resolve once the response has finished.
       //
       // See: https://nodejs.org/api/http.html#event-finish_1

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
@@ -1,5 +1,7 @@
 import type { FlightRouterState } from '../app-render/types'
 
+import { WritableStream } from 'next/dist/compiled/@edge-runtime/ponyfill'
+
 import { nonNullable } from '../../lib/non-nullable'
 import { getTracer } from '../lib/trace/tracer'
 import { AppRenderSpan } from '../lib/trace/constants'