diff --git a/.changeset/cloudflare-stream-backpressure.md b/.changeset/cloudflare-stream-backpressure.md new file mode 100644 index 000000000..f737c5138 --- /dev/null +++ b/.changeset/cloudflare-stream-backpressure.md @@ -0,0 +1,5 @@ +--- +"@opennextjs/aws": patch +--- + +Bound Cloudflare Node response streaming with backpressure. diff --git a/packages/open-next/src/core/requestHandler.ts b/packages/open-next/src/core/requestHandler.ts index b66eda577..6c3a68251 100644 --- a/packages/open-next/src/core/requestHandler.ts +++ b/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"; @@ -160,10 +162,9 @@ 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); - } - response.end(); + // Use pipeline so the streamCreator Writable can apply backpressure + // instead of eagerly consuming the whole Web stream. + await pipeline(Readable.fromWeb(bodyToConsume), response); routingResult.body = bodyToReturn; } return routingResult; diff --git a/packages/open-next/src/overrides/wrappers/cloudflare-node.ts b/packages/open-next/src/overrides/wrappers/cloudflare-node.ts index d706e91f1..df0878754 100644 --- a/packages/open-next/src/overrides/wrappers/cloudflare-node.ts +++ b/packages/open-next/src/overrides/wrappers/cloudflare-node.ts @@ -5,7 +5,7 @@ import type { } 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]); @@ -30,8 +30,12 @@ const handler: WrapperHandler = const internalEvent = await converter.convertFrom(request); const url = new URL(request.url); - const { promise: promiseResponse, resolve: resolveResponse } = - Promise.withResolvers(); + // writeHeaders resolves this once Next has chosen status/headers; the body is + // streamed later through the Writable returned from streamCreator. + let resolveResponse!: (response: Response) => void; + const promiseResponse = new Promise((resolve) => { + resolveResponse = resolve; + }); const streamCreator: StreamCreator = { writeHeaders(prelude: { @@ -68,11 +72,82 @@ const handler: WrapperHandler = }); } + // Bridge Next's Node stream into Cloudflare's Web stream: + // - Next writes chunks to the Writable below. + // - Cloudflare reads chunks from the Response's ReadableStream. + // - Each Node write callback is held until the Web stream pulls again, + // which propagates reader backpressure to the Node producer. let controller: ReadableStreamDefaultController; + let controllerClosed = false; + let producer: Readable | undefined; + let completed = false; + let destructionError: Error | undefined; + let parkedWriteCallback: + | ((error?: Error | null | undefined) => void) + | undefined; + + // Completing the parked callback tells Node the previous write finished, + // allowing a piped source to produce the next chunk. + 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)); + + // pull() means the Response consumer has capacity for more data; use that + // signal to release exactly one pending Node write. const readable = new ReadableStream({ start(c) { controller = c; }, + pull() { + settleParkedWrite(); + }, + cancel(reason) { + // The response reader went away. Destroy the original producer too, + // otherwise it may keep generating chunks into a closed bridge. + controllerClosed = true; + destructionError = normalizeError(reason); + destroyProducer(); + bridge.destroy(); + }, }); const response = new Response(readable, { @@ -81,32 +156,55 @@ const handler: WrapperHandler = }); resolveResponse(response); - return new Writable({ + const bridge = new Writable({ + // Keep the Writable queue small; backpressure is enforced by delaying + // each write callback until the Web stream pulls again. + 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(); + // Do not call the callback yet: parking it pauses the Node producer + // until pull() shows the Web stream wants another chunk. + 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); }, }); + + // The source Readable is only exposed through the pipe event; keep it so + // Web stream cancellation can tear down upstream work. + 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 diff --git a/packages/tests-unit/tests/core/requestHandler.test.ts b/packages/tests-unit/tests/core/requestHandler.test.ts new file mode 100644 index 000000000..735f67383 --- /dev/null +++ b/packages/tests-unit/tests/core/requestHandler.test.ts @@ -0,0 +1,173 @@ +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((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({ + 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): 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("retains a tee branch 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; + for (let i = 0; i < 25; i++) { + await nextTurn(); + } + + expect(source.produced).toBe(producedWhileStalled); + expect(source.produced).toBeLessThan(20); + expect(tee).toHaveBeenCalledOnce(); + + while (received.length < 20) { + await waitFor(() => writeCallbacks.length > 0); + writeCallbacks.shift()!(); + } + + const returned = await handlerPromise; + expect(returned).toBe(result); + expect(received).toEqual(Array.from({ length: 20 }, (_, index) => index)); + expect(returned.body).not.toBe(source.body); + expect((await new Response(returned.body).arrayBuffer()).byteLength).toBe( + 20 * 64 * 1024, + ); + 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, + ); + }); +}); diff --git a/packages/tests-unit/tests/overrides/wrappers/cloudflare-node.test.ts b/packages/tests-unit/tests/overrides/wrappers/cloudflare-node.test.ts new file mode 100644 index 000000000..d9d9bafec --- /dev/null +++ b/packages/tests-unit/tests/overrides/wrappers/cloudflare-node.test.ts @@ -0,0 +1,184 @@ +import { Readable, type Writable } from "node:stream"; +import { finished } from "node:stream/promises"; + +import cloudflareNode from "@opennextjs/aws/overrides/wrappers/cloudflare-node.js"; +import type { + InternalEvent, + InternalResult, +} from "@opennextjs/aws/types/open-next.js"; +import type { + Converter, + OpenNextHandler, +} from "@opennextjs/aws/types/overrides.js"; + +const event: InternalEvent = { + type: "core", + method: "GET", + rawPath: "/", + url: "https://example.com/", + body: Buffer.alloc(0), + headers: {}, + query: {}, + cookies: {}, + remoteAddress: "::1", +}; + +const emptyResult: InternalResult = { + type: "core", + statusCode: 200, + headers: {}, + body: new ReadableStream(), + isBase64Encoded: false, +}; + +const converter = { + name: "test", + convertFrom: async () => event, + convertTo: async () => undefined, +} satisfies Converter; + +function nextTurn() { + return new Promise((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 createProducer(total = Number.POSITIVE_INFINITY) { + let produced = 0; + let destroyedWith: Error | null | undefined; + const producer = new Readable({ + highWaterMark: 1, + read() { + if (produced >= total) { + this.push(null); + return; + } + this.push(Buffer.from([produced++])); + }, + destroy(error, callback) { + destroyedWith = error; + callback(error); + }, + }); + + return { + producer, + get produced() { + return produced; + }, + get destroyedWith() { + return destroyedWith; + }, + }; +} + +async function startResponse(producer: Readable) { + let destination: Writable | undefined; + let handlerPromise: Promise | undefined; + const handler: OpenNextHandler = async (_event, options) => { + destination = options?.streamCreator?.writeHeaders({ + statusCode: 200, + cookies: [], + headers: {}, + }); + if (!destination) { + throw new Error("Missing stream destination"); + } + producer.pipe(destination); + await finished(destination); + return emptyResult; + }; + + const wrapped = await cloudflareNode.wrapper(handler, converter); + const response = await wrapped( + new Request("https://example.com/"), + {}, + { + waitUntil(promise: Promise) { + handlerPromise = promise; + promise.catch(() => {}); + }, + }, + new AbortController().signal, + ); + + if (!destination || !handlerPromise) { + throw new Error("Wrapper did not start streaming"); + } + + return { destination, handlerPromise, response }; +} + +describe("cloudflare-node wrapper streaming", () => { + it("bounds production and delivers chunks in order to a slow consumer", async () => { + const source = createProducer(32); + const { handlerPromise, response } = await startResponse(source.producer); + + await waitFor(() => source.produced > 0); + await nextTurn(); + const producedWhileIdle = source.produced; + for (let i = 0; i < 25; i++) { + await nextTurn(); + } + + expect(source.produced).toBe(producedWhileIdle); + expect(source.produced).toBeLessThan(32); + + const reader = response.body!.getReader(); + const received: number[] = []; + while (true) { + const { done, value } = await reader.read(); + if (done) { + break; + } + received.push(value[0]); + expect(source.produced).toBeLessThanOrEqual(received.length + 3); + await nextTurn(); + } + + await expect(handlerPromise).resolves.toBe(emptyResult); + expect(received).toEqual(Array.from({ length: 32 }, (_, index) => index)); + }); + + it("cancels the bridge and its producer while settling the handler", async () => { + const source = createProducer(); + const { destination, handlerPromise, response } = await startResponse( + source.producer, + ); + const reader = response.body!.getReader(); + + await reader.read(); + const cancellation = new Error("consumer cancelled"); + await reader.cancel(cancellation); + + await expect(handlerPromise).rejects.toMatchObject({ + code: "ERR_STREAM_PREMATURE_CLOSE", + }); + expect(destination.destroyed).toBe(true); + expect(source.producer.destroyed).toBe(true); + expect(source.destroyedWith).toBeNull(); + }); + + it("destroys the producer and errors the Web response when destroyed", async () => { + const source = createProducer(); + const { destination, handlerPromise, response } = await startResponse( + source.producer, + ); + const destruction = new Error("bridge destroyed"); + + destination.destroy(destruction); + + await expect(handlerPromise).rejects.toThrow("bridge destroyed"); + await expect(response.text()).rejects.toThrow("bridge destroyed"); + expect(source.producer.destroyed).toBe(true); + expect(source.destroyedWith).toBeNull(); + }); +});