fix: Cloudflare streaming response backpressure

.changeset/cloudflare-stream-backpressure.md

@@ -0,0 +1,5 @@
+---
+"@opennextjs/aws": patch
+---
+
+Bound Cloudflare Node response streaming with backpressure and stop retaining an unused duplicate response body.

packages/open-next/src/core/requestHandler.ts

@@ -1,4 +1,6 @@
 import { AsyncLocalStorage } from "node:async_hooks";
+import { Readable } from "node:stream";
+import { pipeline } from "node:stream/promises";
 
 import type { OpenNextNodeResponse } from "http/index.js";
 import { IncomingMessage } from "http/index.js";
@@ -159,12 +161,13 @@ export async function openNextHandler(
           );
           response.statusCode = routingResult.statusCode;
           response.flushHeaders();
-          const [bodyToConsume, bodyToReturn] = routingResult.body.tee();
-          for await (const chunk of bodyToConsume) {
-            response.write(chunk);
+          let bodyToConsume = routingResult.body;
+          if (options.streamCreator.retainChunks !== false) {
+            const [bodyToStream, bodyToReturn] = routingResult.body.tee();
+            bodyToConsume = bodyToStream;
+            routingResult.body = bodyToReturn;
           }
-          response.end();
-          routingResult.body = bodyToReturn;
+          await pipeline(Readable.fromWeb(bodyToConsume), response);
         }
         return routingResult;
       }

packages/open-next/src/overrides/wrappers/cloudflare-node.ts

@@ -5,7 +5,7 @@
 } from "types/open-next";
 import type { Wrapper, WrapperHandler } from "types/overrides";
 
-import { Writable } from "node:stream";
+import { type Readable, Writable } from "node:stream";
 
 // Response with null body status (101, 204, 205, or 304) cannot have a body.
 const NULL_BODY_STATUSES = new Set([101, 204, 205, 304]);
@@ -31,7 +31,7 @@
     const url = new URL(request.url);
 
     const { promise: promiseResponse, resolve: resolveResponse } =
       Promise.withResolvers<Response>();
 
     const streamCreator: StreamCreator = {
       writeHeaders(prelude: {
@@ -69,10 +69,70 @@
         }
 
         let controller: ReadableStreamDefaultController<Uint8Array>;
+        let controllerClosed = false;
+        let producer: Readable | undefined;
+        let completed = false;
+        let destructionError: Error | undefined;
+        let parkedWriteCallback:
+          | ((error?: Error | null | undefined) => void)
+          | undefined;
+
+        const settleParkedWrite = (error?: Error | null) => {
+          const callback = parkedWriteCallback;
+          parkedWriteCallback = undefined;
+          callback?.(error);
+        };
+
+        const destroyProducer = () => {
+          if (!completed && producer && !producer.destroyed) {
+            producer.destroy();
+          }
+        };
+
+        const closeController = () => {
+          if (controllerClosed) {
+            return;
+          }
+          controllerClosed = true;
+          try {
+            controller.close();
+          } catch {
+            // The Web stream may already have been closed by its consumer.
+          }
+        };
+
+        const errorController = (error: Error) => {
+          if (controllerClosed) {
+            return;
+          }
+          controllerClosed = true;
+          try {
+            controller.error(error);
+          } catch {
+            // The Web stream may already have been closed by its consumer.
+          }
+        };
+
+        const normalizeError = (reason: unknown) =>
+          reason == null
+            ? undefined
+            : reason instanceof Error
+              ? reason
+              : new Error(String(reason));
+
         const readable = new ReadableStream({
           start(c) {
             controller = c;
           },
+          pull() {
+            settleParkedWrite();
+          },
+          cancel(reason) {
+            controllerClosed = true;
+            destructionError = normalizeError(reason);
+            destroyProducer();
+            bridge.destroy();
+          },
         });
 
         const response = new Response(readable, {
@@ -81,32 +141,50 @@
         });
         resolveResponse(response);
 
-        return new Writable({
+        const bridge = new Writable({
+          // Keep Node-side buffering minimal; write callbacks are released from pull().
+          highWaterMark: 1,
           write(chunk, encoding, callback) {
+            if (controllerClosed) {
+              callback(
+                destructionError ?? new Error("Response stream is closed"),
+              );
+              return;
+            }
             try {
               controller.enqueue(chunk);
             } catch (e: any) {
-              return callback(e);
+              callback(e);
+              return;
             }
-            callback();
+            parkedWriteCallback = callback;
           },
           final(callback) {
-            controller.close();
+            completed = true;
+            closeController();
             callback();
           },
           destroy(error, callback) {
+            destroyProducer();
+            settleParkedWrite(error);
             if (error) {
-              controller.error(error);
+              destructionError = error;
+              errorController(error);
             } else {
-              try {
-                controller.close();
-              } catch {
-                // Ignore "This ReadableStream is closed" error
-              }
+              closeController();
             }
             callback(error);
           },
         });
+
+        bridge.on("pipe", (source: Readable) => {
+          producer = source;
+          if (bridge.destroyed) {
+            destroyProducer();
+          }
+        });
+
+        return bridge;
       },
       // This is for passing along the original abort signal from the initial Request you retrieve in your worker
       // Ensures that the response we pass to NextServer is aborted if the request is aborted

packages/open-next/src/types/open-next.ts

@@ -61,7 +61,8 @@ export interface StreamCreator {
    *
    * However some implementations use a fake `StreamCreator` and expect the chunks to be retained.
    * When your stream controller implementation doesn't need to retain the chunk, you can set this
-   * to `false` to reduce memory usage.
+   * to `false` to reduce memory usage. This also permits the handler to consume the returned response
+   * body directly instead of retaining an unused duplicate with `ReadableStream.tee()`.
    *
    * @see https://github.com/opennextjs/opennextjs-aws/blob/main/packages/open-next/src/overrides/wrappers/aws-lambda.ts
    *

packages/tests-unit/tests/core/requestHandler.test.ts

@@ -0,0 +1,166 @@
+import { Writable } from "node:stream";
+import { ReadableStream } from "node:stream/web";
+
+import { openNextHandler } from "@opennextjs/aws/core/requestHandler.js";
+import routingHandler from "@opennextjs/aws/core/routingHandler.js";
+import type {
+  InternalEvent,
+  InternalResult,
+  StreamCreator,
+} from "@opennextjs/aws/types/open-next.js";
+import { beforeEach, vi } from "vitest";
+
+vi.mock("@opennextjs/aws/adapters/config/index.js", () => ({
+  BuildId: "build-id",
+  HtmlPages: [],
+  NextConfig: {},
+}));
+
+vi.mock("@opennextjs/aws/core/routingHandler.js", () => ({
+  default: vi.fn(),
+  INTERNAL_EVENT_REQUEST_ID: "x-opennext-request-id",
+  INTERNAL_HEADER_INITIAL_URL: "x-opennext-initial-url",
+  INTERNAL_HEADER_RESOLVED_ROUTES: "x-opennext-resolved-routes",
+  INTERNAL_HEADER_REWRITE_STATUS_CODE: "x-opennext-rewrite-status-code",
+  MIDDLEWARE_HEADER_PREFIX: "x-middleware-response-",
+  MIDDLEWARE_HEADER_PREFIX_LEN: "x-middleware-response-".length,
+}));
+
+vi.mock("@opennextjs/aws/core/util.js", () => ({
+  requestHandler: vi.fn(),
+  setNextjsPrebundledReact: vi.fn(),
+}));
+
+const mockedRoutingHandler = vi.mocked(routingHandler);
+const event: InternalEvent = {
+  type: "core",
+  method: "GET",
+  rawPath: "/",
+  url: "https://example.com/",
+  body: Buffer.alloc(0),
+  headers: {},
+  query: {},
+  cookies: {},
+  remoteAddress: "::1",
+};
+
+function nextTurn() {
+  return new Promise<void>((resolve) => setImmediate(resolve));
+}
+
+async function waitFor(predicate: () => boolean) {
+  for (let attempt = 0; attempt < 100; attempt++) {
+    if (predicate()) {
+      return;
+    }
+    await nextTurn();
+  }
+  throw new Error("Timed out waiting for condition");
+}
+
+function createBody(total: number) {
+  let produced = 0;
+  const body = new ReadableStream<Uint8Array>({
+    pull(controller) {
+      if (produced >= total) {
+        controller.close();
+        return;
+      }
+      const chunk = new Uint8Array(64 * 1024);
+      chunk[0] = produced++;
+      controller.enqueue(chunk);
+    },
+  });
+
+  return {
+    body,
+    get produced() {
+      return produced;
+    },
+  };
+}
+
+function createResult(body: ReadableStream<Uint8Array>): InternalResult {
+  return {
+    type: "core",
+    statusCode: 200,
+    headers: {},
+    body,
+    isBase64Encoded: false,
+  };
+}
+
+beforeEach(() => {
+  vi.clearAllMocks();
+  globalThis.openNextConfig = {};
+  globalThis.__next_route_preloader = vi.fn(async () => {});
+});
+
+describe("requestHandler streaming", () => {
+  it("skips tee and remains bounded when chunks do not need retaining", async () => {
+    const source = createBody(20);
+    const tee = vi.spyOn(source.body, "tee");
+    const result = createResult(source.body);
+    mockedRoutingHandler.mockResolvedValue(result);
+
+    const writeCallbacks: Array<() => void> = [];
+    const received: number[] = [];
+    const streamCreator: StreamCreator = {
+      retainChunks: false,
+      writeHeaders: () =>
+        new Writable({
+          highWaterMark: 1,
+          write(chunk: Buffer, _encoding, callback) {
+            received.push(chunk[0]);
+            writeCallbacks.push(callback);
+          },
+        }),
+    };
+
+    const handlerPromise = openNextHandler(event, { streamCreator });
+    await waitFor(() => writeCallbacks.length > 0);
+    await nextTurn();
+    const producedWhileStalled = source.produced;
+    await new Promise((resolve) => setTimeout(resolve, 20));
+
+    expect(source.produced).toBe(producedWhileStalled);
+    expect(source.produced).toBeLessThan(20);
+    expect(tee).not.toHaveBeenCalled();
+
+    while (received.length < 20) {
+      await waitFor(() => writeCallbacks.length > 0);
+      writeCallbacks.shift()!();
+    }
+
+    await expect(handlerPromise).resolves.toBe(result);
+    expect(received).toEqual(Array.from({ length: 20 }, (_, index) => index));
+    expect(source.body.locked).toBe(true);
+  });
+
+  it("retains a tee branch by default", async () => {
+    const source = createBody(3);
+    const tee = vi.spyOn(source.body, "tee");
+    const result = createResult(source.body);
+    mockedRoutingHandler.mockResolvedValue(result);
+
+    const streamed: Buffer[] = [];
+    const streamCreator: StreamCreator = {
+      writeHeaders: () =>
+        new Writable({
+          write(chunk, _encoding, callback) {
+            streamed.push(chunk);
+            callback();
+          },
+        }),
+    };
+
+    const returned = await openNextHandler(event, { streamCreator });
+
+    expect(tee).toHaveBeenCalledOnce();
+    expect(returned.body).not.toBe(source.body);
+    expect(Buffer.concat(streamed).length).toBe(3 * 64 * 1024);
+    expect((await new Response(returned.body).arrayBuffer()).byteLength).toBe(
+      3 * 64 * 1024,
+    );
+  });
+});