fix(edge): run after() if request is cancelled mid-streaming (#76013)

- Fixes an issue where `after()` in an edge page would not run if the
request was cancelled/aborted
- Unskips the `runs callbacks if redirect() was called` test which was
disabled to not block other things

This was initially hit in #75882, during which the `runs callbacks if
redirect() was called` test started failing when using experimental
react + turbo in dev mode. Turns out that this happenws because
something got slower and we weren't finishing the redirect response in
time, i.e. before the browser disconnected and started loading the page
it got redirected to.

It's relevant that the response didn't finish streaming, because in
`edge`, we use that as the trigger to start running `after()` callbacks.
In particular, we instrument the response stream using
`trackStreamConsumed`. The problem was that this function didn't handle
the stream being cancelled, which is what happens when a request is
aborted mid-streaming, so `after()` never ended up executing.

This PR fixes that and adds some tests for cancellation and interrupted
streaming.

Author: lubieowoce

packages/next/src/server/web/web-on-close.test.ts

@@ -0,0 +1,106 @@
+import { DetachedPromise } from '../../lib/detached-promise'
+import { trackStreamConsumed } from './web-on-close'
+
+describe('trackStreamConsumed', () => {
+  it('calls onEnd when the stream finishes', async () => {
+    const endPromise = new DetachedPromise<void>()
+    const onEnd = jest.fn(endPromise.resolve)
+
+    const { stream: inputStream, controller } =
+      readableStreamWithController<string>()
+    const trackedStream = trackStreamConsumed(inputStream, onEnd)
+
+    const reader = trackedStream.getReader()
+    controller.enqueue('one')
+    controller.enqueue('two')
+    await reader.read()
+    await reader.read()
+    expect(onEnd).not.toHaveBeenCalled()
+
+    controller.close()
+
+    await expect(reader.read()).resolves.toEqual({
+      done: true,
+      value: undefined,
+    })
+
+    await endPromise.promise
+    expect(onEnd).toHaveBeenCalledTimes(1)
+  })
+
+  it('calls onEnd when the stream errors', async () => {
+    const endPromise = new DetachedPromise<void>()
+    const onEnd = jest.fn(endPromise.resolve)
+
+    const { stream: inputStream, controller } =
+      readableStreamWithController<string>()
+    const trackedStream = trackStreamConsumed(inputStream, onEnd)
+
+    const reader = trackedStream.getReader()
+    controller.enqueue('one')
+    controller.enqueue('two')
+    await reader.read()
+    await reader.read()
+    expect(onEnd).not.toHaveBeenCalled()
+
+    const error = new Error('kaboom')
+    controller.error(error)
+
+    // if the underlying stream errors, we should error as well
+    await expect(reader.read()).rejects.toThrow(error)
+
+    await endPromise.promise
+    expect(onEnd).toHaveBeenCalledTimes(1)
+  })
+
+  it('calls onEnd when the stream is cancelled', async () => {
+    const endPromise = new DetachedPromise<void>()
+    const onEnd = jest.fn(endPromise.resolve)
+
+    const cancelledPromise = new DetachedPromise<unknown>()
+    const onCancel = jest.fn(cancelledPromise.resolve)
+
+    const { stream: inputStream, controller } =
+      readableStreamWithController<string>(onCancel)
+    const trackedStream = trackStreamConsumed(inputStream, onEnd)
+
+    const reader = trackedStream.getReader()
+    controller.enqueue('one')
+    controller.enqueue('two')
+    await reader.read()
+    await reader.read()
+    expect(onEnd).not.toHaveBeenCalled()
+
+    const cancellationReason = new Error('cancelled')
+    await reader.cancel(cancellationReason)
+
+    // from a reader's perspective, a cancelled stream behaves like it's done
+    // (which is a bit weird honestly?)
+    await expect(reader.read()).resolves.toEqual({
+      done: true,
+      value: undefined,
+    })
+
+    await endPromise.promise
+    expect(onEnd).toHaveBeenCalledTimes(1)
+
+    //  the cancellation should propagate to back to the underlying stream
+    await cancelledPromise.promise
+    expect(onCancel).toHaveBeenCalledWith(cancellationReason)
+  })
+})
+
+function readableStreamWithController<TChunk>(
+  onCancel?: (reason: unknown) => void
+) {
+  let controller: ReadableStreamDefaultController<TChunk> = undefined!
+  const stream = new ReadableStream<TChunk>({
+    start(_controller) {
+      controller = _controller
+    },
+    cancel(reason) {
+      onCancel?.(reason)
+    },
+  })
+  return { controller, stream }
+}

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

@@ -22,12 +22,17 @@ export function trackStreamConsumed<TChunk>(
   stream: ReadableStream<TChunk>,
   onEnd: () => void
 ): ReadableStream<TChunk> {
-  const closePassThrough = new TransformStream<TChunk, TChunk>({
-    flush: () => {
-      return onEnd()
-    },
-  })
-  return stream.pipeThrough(closePassThrough)
+  // NOTE: This function must handle `stream` being aborted or cancelled,
+  // so it can't just be this:
+  //
+  //   return stream.pipeThrough(new TransformStream({ flush() { onEnd() } }))
+  //
+  // because that doesn't handle cancellations.
+  // (and cancellation handling via `Transformer.cancel` is only available in node >20)
+  const dest = new TransformStream()
+  const runOnEnd = () => onEnd()
+  stream.pipeTo(dest.writable).then(runOnEnd, runOnEnd)
+  return dest.readable
 }
 
 export class CloseController {

test/e2e/app-dir/next-after-app/app/edge/interrupted/incomplete-stream/end/page.js

@@ -0,0 +1 @@
+export { default } from '../../../../nodejs/interrupted/incomplete-stream/end/page'

test/e2e/app-dir/next-after-app/app/edge/interrupted/incomplete-stream/hang/page.js

@@ -0,0 +1 @@
+export { default } from '../../../../nodejs/interrupted/incomplete-stream/hang/page'

test/e2e/app-dir/next-after-app/app/edge/interrupted/incomplete-stream/start/page.js

@@ -0,0 +1 @@
+export { default } from '../../../../nodejs/interrupted/incomplete-stream/start/page'

test/e2e/app-dir/next-after-app/app/nodejs/interrupted/incomplete-stream/end/page.js

@@ -0,0 +1,7 @@
+export default function Page() {
+  return (
+    <main>
+      <h1>End</h1>
+    </main>
+  )
+}

test/e2e/app-dir/next-after-app/app/nodejs/interrupted/incomplete-stream/hang/page.js

@@ -0,0 +1,34 @@
+import { connection, after } from 'next/server'
+import { Suspense } from 'react'
+import { cliLog } from '../../../../../utils/log'
+
+export default async function Page() {
+  await connection()
+  after(() => {
+    cliLog({
+      source: '[page] /interrupted/incomplete-stream/hang',
+    })
+  })
+  return (
+    <main>
+      <h1>Hanging forever</h1>
+      <Suspense
+        fallback={
+          <div id="loading-fallback">
+            {
+              // we're going to look for this string in the streamed response,
+              // make sure it doesn't show up in the literal form someplace else
+              'Loading' + (Math.random() > 1 ? 'impossible' : '') + '...'
+            }
+          </div>
+        }
+      >
+        <HangForever />
+      </Suspense>
+    </main>
+  )
+}
+
+async function HangForever() {
+  await new Promise((_resolve) => {})
+}

test/e2e/app-dir/next-after-app/app/nodejs/interrupted/incomplete-stream/start/page.js

@@ -0,0 +1,14 @@
+'use client'
+
+import Link from 'next/link'
+import { usePathname } from 'next/navigation'
+
+export default function Page() {
+  const pathname = usePathname()
+  return (
+    <main>
+      <h1>Start</h1>
+      <Link href={pathname.replace('/start', '/hang')}>hang</Link>
+    </main>
+  )
+}

test/e2e/app-dir/next-after-app/index.test.ts

@@ -124,6 +124,64 @@ describe.each(runtimes)('after() in %s runtime', (runtimeValue) => {
         source: '[page] /interrupted/throws-error',
       })
     })
+
+    it('runs callbacks if a request is aborted before the page finishes streaming', async () => {
+      const abortController = new AbortController()
+      const res = await next.fetch(
+        pathPrefix + '/interrupted/incomplete-stream/hang',
+        { signal: abortController.signal }
+      )
+      expect(res.status).toBe(200)
+
+      const textDecoder = new TextDecoder()
+      for await (const rawChunk of res.body) {
+        const chunk =
+          typeof rawChunk === 'string' ? rawChunk : textDecoder.decode(rawChunk)
+        // we found the loading fallback for the part that hangs forever, so we know we won't progress any further
+        if (chunk.includes('Loading...')) {
+          break
+        }
+      }
+      abortController.abort()
+
+      await retry(() => {
+        expect(getLogs()).toContainEqual({
+          source: '[page] /interrupted/incomplete-stream/hang',
+        })
+      })
+    })
+
+    it('runs callbacks if the browser disconnects before the page finishes streaming', async () => {
+      // `next.browser()` always waits for the `load` event, which we don't want here.
+      // (because the page hangs forever while streaming and will thus never fire `load`)
+      // but we can't easily bypass that, so go to a dummy page first
+      const browser = await next.browser(
+        pathPrefix + '/interrupted/incomplete-stream/start'
+      )
+      expect(await browser.elementByCss('h1').text()).toEqual('Start')
+
+      // navigate to a page that hangs forever while streaming...
+      // NOTE: this needs to be a soft navigation (using Link), playwright seems to hang otherwise
+      await browser.elementByCss('a').click()
+      await retry(async () => {
+        expect(await browser.hasElementByCssSelector('#loading-fallback')).toBe(
+          true
+        )
+      })
+
+      // ...but navigate away before streaming is finished (it hangs forever, so it will never finish)
+      await browser.get(
+        new URL(pathPrefix + '/interrupted/incomplete-stream/end', next.url)
+          .href
+      )
+      expect(await browser.elementByCss('h1').text()).toEqual('End')
+
+      await retry(async () => {
+        expect(getLogs()).toContainEqual({
+          source: '[page] /interrupted/incomplete-stream/hang',
+        })
+      })
+    })
   })
 
   it('runs in middleware', async () => {