Revert "Revert "Fix stream cancellation in RenderResult.pipe() and sendResponse()" (#52277)"

This reverts commit 5046ee1.

Author: jridgewell

packages/next-swc/crates/next-core/js/src/entry/app/edge-page-bootstrap.ts

@@ -80,11 +80,20 @@ async function render(request: NextRequest, event: NextFetchEvent) {
   response.headers.append('Vary', RSC_VARY_HEADER)
 
   const writer = tranform.writable.getWriter()
-  result.pipe({
+  const target = {
     write: (chunk: Uint8Array) => writer.write(chunk),
     end: () => writer.close(),
     destroy: (reason?: Error) => writer.abort(reason),
-  })
+    closed: false,
+  }
+  const onClose = () => {
+    target.closed = true
+  }
+  // No, this cannot be replaced with `finally`, because early cancelling
+  // the stream will create a rejected promise, and finally will create an
+  // unhandled rejection.
+  writer.closed.then(onClose, onClose)
+  result.pipe(target)
 
   return response
 }

packages/next/src/server/render-result.ts

@@ -16,6 +16,7 @@ export interface PipeTarget {
   end: () => unknown
   flush?: () => unknown
   destroy: (err?: Error) => unknown
+  get closed(): boolean
 }
 
 export default class RenderResult {
@@ -111,31 +112,25 @@ export default class RenderResult {
         : () => {}
     const reader = this.response.getReader()
 
-    let shouldFatalError = false
     try {
-      let result = await reader.read()
-      if (!result.done) {
-        // As we're going to write to the response, we should destroy the
-        // response if an error occurs.
-        shouldFatalError = true
-      }
+      while (true) {
+        const result = await reader.read()
+
+        if (res.closed || result.done) {
+          break
+        }
 
-      while (!result.done) {
         // Write the data to the response.
         res.write(result.value)
 
         // Flush it to the client (if it supports flushing).
         flush()
-
-        // Read the next chunk.
-        result = await reader.read()
       }
 
       // We're done writing to the response, so we can end it.
       res.end()
     } catch (err) {
-      // If we've written to the response, we should destroy it.
-      if (shouldFatalError) {
+      if (!res.closed) {
         res.destroy(err as any)
       }
 

packages/next/src/server/send-response.ts

@@ -49,6 +49,7 @@ export async function sendResponse(
       const iterator = consumeUint8ArrayReadableStream(response.body)
       try {
         for await (const chunk of iterator) {
+          if (originalResponse.closed) break
           originalResponse.write(chunk)
         }
       } finally {

packages/next/src/server/stream-utils/node-web-streams-helper.ts

@@ -25,8 +25,9 @@ export const streamToBufferedResult = async (
     },
     end() {},
     destroy() {},
+    closed: false,
   }
-  await renderResult.pipe(writable as any)
+  await renderResult.pipe(writable)
   return renderChunks.join('')
 }
 

packages/next/src/server/web-server.ts

@@ -417,14 +417,20 @@ export default class NextWebServer extends BaseServer<WebServerOptions> {
 
     if (options.result.isDynamic) {
       const writer = res.transformStream.writable.getWriter()
-      options.result.pipe({
+      const target = {
         write: (chunk: Uint8Array) => writer.write(chunk),
         end: () => writer.close(),
         destroy: (err: Error) => writer.abort(err),
-        cork: () => {},
-        uncork: () => {},
-        // Not implemented: on/removeListener
-      } as any)
+        closed: false,
+      } as any
+      const onClose = () => {
+        target.closed = true
+      }
+      // No, this cannot be replaced with `finally`, because early cancelling
+      // the stream will create a rejected promise, and finally will create an
+      // unhandled rejection.
+      writer.closed.then(onClose, onClose)
+      options.result.pipe(target)
     } else {
       const payload = await options.result.toUnchunkedString()
       res.setHeader('Content-Length', String(byteLength(payload)))

test/e2e/cancel-request/app/edge-route/route.ts

@@ -0,0 +1,26 @@
+import { Streamable } from '../../streamable'
+
+export const runtime = 'edge'
+
+let streamable
+let requestAborted = false
+
+export function GET(req: Request): Response {
+  // The 2nd request should render the stats. We don't use a query param
+  // because edge rendering will create a different bundle for that.
+  if (streamable) {
+    return new Response(
+      JSON.stringify({
+        requestAborted,
+        i: streamable.i,
+        streamCleanedUp: streamable.streamCleanedUp,
+      })
+    )
+  }
+
+  streamable = Streamable()
+  req.signal.onabort = () => {
+    requestAborted = true
+  }
+  return new Response(streamable.stream)
+}

test/e2e/cancel-request/app/node-route/route.ts

@@ -0,0 +1,28 @@
+import { Streamable } from '../../streamable'
+
+export const runtime = 'nodejs'
+// Next thinks it can statically compile this route, which breaks the test.
+export const dynamic = 'force-dynamic'
+
+let streamable
+let requestAborted = false
+
+export function GET(req: Request): Response {
+  // The 2nd request should render the stats. We don't use a query param
+  // because edge rendering will create a different bundle for that.
+  if (streamable) {
+    return new Response(
+      JSON.stringify({
+        requestAborted,
+        i: streamable.i,
+        streamCleanedUp: streamable.streamCleanedUp,
+      })
+    )
+  }
+
+  streamable = Streamable()
+  req.signal.onabort = () => {
+    requestAborted = true
+  }
+  return new Response(streamable.stream)
+}

test/e2e/cancel-request/middleware.ts

@@ -0,0 +1,28 @@
+import { Streamable } from './streamable'
+
+export const config = {
+  matcher: '/middleware',
+}
+
+let streamable
+let requestAborted = false
+
+export default function handler(req: Request): Response {
+  // The 2nd request should render the stats. We don't use a query param
+  // because edge rendering will create a different bundle for that.
+  if (streamable) {
+    return new Response(
+      JSON.stringify({
+        requestAborted,
+        i: streamable.i,
+        streamCleanedUp: streamable.streamCleanedUp,
+      })
+    )
+  }
+
+  streamable = Streamable()
+  req.signal.onabort = () => {
+    requestAborted = true
+  }
+  return new Response(streamable.stream)
+}

test/e2e/cancel-request/pages/api/edge-api.ts

@@ -0,0 +1,28 @@
+import { Streamable } from '../../streamable'
+
+export const config = {
+  runtime: 'edge',
+}
+
+let streamable
+let requestAborted = false
+
+export default function handler(req: Request): Response {
+  // The 2nd request should render the stats. We don't use a query param
+  // because edge rendering will create a different bundle for that.
+  if (streamable) {
+    return new Response(
+      JSON.stringify({
+        requestAborted,
+        i: streamable.i,
+        streamCleanedUp: streamable.streamCleanedUp,
+      })
+    )
+  }
+
+  streamable = Streamable()
+  req.signal.onabort = () => {
+    requestAborted = true
+  }
+  return new Response(streamable.stream)
+}

test/e2e/cancel-request/pages/api/node-api.ts

@@ -0,0 +1,36 @@
+import { IncomingMessage, ServerResponse } from 'http'
+import { pipeline } from 'stream'
+import { Readable } from '../../readable'
+
+export const config = {
+  runtime: 'nodejs',
+}
+
+let readable
+let requestAborted = false
+
+export default function handler(
+  _req: IncomingMessage,
+  res: ServerResponse
+): void {
+  // The 2nd request should render the stats. We don't use a query param
+  // because edge rendering will create a different bundle for that.
+  if (readable) {
+    res.end(
+      JSON.stringify({
+        requestAborted,
+        i: readable.i,
+        streamCleanedUp: readable.streamCleanedUp,
+      })
+    )
+    return
+  }
+
+  readable = Readable()
+  res.on('close', () => {
+    requestAborted = true
+  })
+  pipeline(readable.stream, res, () => {
+    res.end()
+  })
+}