refactor: Remove retainChunks option and update response handling for improved memory management

Author: conico974

packages/aws/src/overrides/converters/aws-streaming.ts

@@ -31,7 +31,6 @@ const converter: Converter = {
 				responseStream.write(new Uint8Array(8));
 				return writable;
 			},
-			retainChunks: false,
 		};
 
 		return { type: "stream" as const, streamCreator };

packages/aws/src/overrides/converters/response-stream.ts

@@ -39,7 +39,6 @@ export function createBufferedStreamCreator<T>(
 				},
 			});
 		},
-		retainChunks: false,
 	};
 
 	return { streamCreator, output };

packages/aws/src/overrides/wrappers/aws-lambda-streaming.ts

@@ -80,7 +80,7 @@ const handler: WrapperHandler = async (handler, converter) =>
 			}
 
 			const response = await handler(internalEvent, { streamCreator: output.streamCreator });
-			if (globalThis.isEdgeRuntime ?? false) {
+			if ((globalThis.isEdgeRuntime ?? false) && response.body) {
 				const stream = output.streamCreator.writeHeaders({
 					statusCode: response.statusCode,
 					headers: response.headers as Record<string, string>,

packages/core/src/core/requestHandler.ts

@@ -135,12 +135,14 @@ 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);
+					if (routingResult.body) {
+						const [bodyToConsume, bodyToReturn] = routingResult.body.tee();
+						for await (const chunk of bodyToConsume) {
+							response.write(chunk);
+						}
+						routingResult.body = bodyToReturn;
 					}
 					response.end();
-					routingResult.body = bodyToReturn;
 				}
 				return routingResult;
 			}
@@ -181,8 +183,10 @@ export async function openNextHandler(
 			if (routingResult.initialResponse) {
 				res.statusCode = routingResult.initialResponse.statusCode;
 				res.flushHeaders();
-				for await (const chunk of routingResult.initialResponse.body) {
-					res.write(chunk);
+				if (routingResult.initialResponse.body) {
+					for await (const chunk of routingResult.initialResponse.body) {
+						res.write(chunk);
+					}
 				}
 
 				//We create a special response for the PPR resume request
@@ -218,13 +222,12 @@ export async function openNextHandler(
 			});
 			//#endOverride
 
-			const { statusCode, headers: responseHeaders, isBase64Encoded, body } = convertRes(res);
+			const { statusCode, headers: responseHeaders, isBase64Encoded } = convertRes(res);
 
 			const internalResult = {
 				type: internalEvent.type,
 				statusCode,
 				headers: responseHeaders,
-				body,
 				isBase64Encoded,
 			};
 

packages/core/src/core/routing/util.ts

@@ -104,21 +104,10 @@ export function convertRes(res: OpenNextNodeResponse): InternalResult {
 	// Probably some kind of race condition
 	const headers = parseHeaders(res.getFixedHeaders());
 	const isBase64Encoded = isBinaryContentType(headers["content-type"]) || !!headers["content-encoding"];
-	const body = new ReadableStream({
-		pull(controller) {
-			if (!res._chunks || res._chunks.length === 0) {
-				controller.close();
-				return;
-			}
-
-			controller.enqueue(res._chunks.shift());
-		},
-	});
 	return {
 		type: "core",
 		statusCode,
 		headers,
-		body,
 		isBase64Encoded,
 	};
 }

packages/core/src/http/openNextResponse.ts

@@ -18,7 +18,6 @@ export class OpenNextNodeResponse extends Transform implements ServerResponse {
 	statusMessage = "";
 	headers: OutgoingHttpHeaders = {};
 	headersSent = false;
-	_chunks: Buffer[] = [];
 	headersAlreadyFixed = false;
 
 	private _cookies: string[] = [];
@@ -255,21 +254,13 @@ export class OpenNextNodeResponse extends Transform implements ServerResponse {
 		return this.headers;
 	}
 
-	getBody() {
-		return Buffer.concat(this._chunks);
-	}
-
 	private _internalWrite(chunk: Buffer | string, encoding: BufferEncoding) {
 		// When encoding === 'buffer', chunk is already a Buffer
 		// and does not need to be converted again.
 		// @ts-expect-error TS2367 'encoding' can be 'buffer', but it's not in the
 		// official type definition
 		const buffer = encoding === "buffer" ? (chunk as Buffer) : Buffer.from(chunk, encoding);
 		this.bodyLength += buffer.length;
-		if (this.streamCreator?.retainChunks !== false) {
-			// Avoid keeping chunks around when the `StreamCreator` supports it to save memory
-			this._chunks.push(buffer);
-		}
 		// No need to pass the encoding for buffers
 		this.push(buffer);
 		this.streamCreator?.onWrite?.();
@@ -342,9 +333,6 @@ export class OpenNextNodeResponse extends Transform implements ServerResponse {
 	}
 
 	send() {
-		for (const chunk of this._chunks) {
-			this.write(chunk);
-		}
 		this.end();
 	}
 

packages/core/src/overrides/converters/edge.ts

@@ -112,7 +112,6 @@ const converter: Converter<InternalEvent, InternalResult | MiddlewareResult> = {
 				});
 			},
 			abortSignal,
-			retainChunks: false,
 		};
 
 		return {

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

@@ -43,7 +43,7 @@ export type MiddlewareEvent = InternalEvent & {
 export type InternalResult = {
 	statusCode: number;
 	headers: Record<string, string | string[]>;
-	body: ReadableStream;
+	body?: ReadableStream;
 	isBase64Encoded: boolean;
 	rewriteStatusCode?: number;
 } & BaseEventOrResult<"core">;
@@ -72,18 +72,6 @@ export interface StreamCreator {
 	onWrite?: () => void;
 	onFinish?: (length: number) => void;
 	abortSignal?: AbortSignal;
-	/**
-	 * Normally there is no need to retain the chunks that have been pushed to the response stream.
-	 *
-	 * 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.
-	 *
-	 * @see https://github.com/opennextjs/opennextjs-aws/blob/main/packages/open-next/src/overrides/wrappers/aws-lambda.ts
-	 *
-	 * @default true for backward compatibility.
-	 */
-	retainChunks?: boolean;
 }
 
 export type WaitUntil = (promise: Promise<void>) => void;

packages/tests-unit/tests/core/routing/util.test.ts

@@ -49,9 +49,6 @@ function createResponse(res: Partial<Res>) {
 	return {
 		statusCode: res.statusCode,
 		getFixedHeaders: () => res.headers ?? {},
-		body: res.body ?? "",
-		getBody: () => Buffer.from(res.body ?? ""),
-		_chunks: res.body ? [Buffer.from(res.body)] : [],
 	};
 }
 
@@ -237,7 +234,7 @@ describe("convertRes", () => {
 				isBase64Encoded: false,
 			})
 		);
-		expect(await fromReadableStream(result.body)).toEqual(res.getBody().toString());
+		expect(result.body).toBeUndefined();
 	});
 
 	it("convert a response with default status code", async () => {
@@ -260,7 +257,7 @@ describe("convertRes", () => {
 				isBase64Encoded: false,
 			})
 		);
-		expect(await fromReadableStream(result.body)).toEqual(res.getBody().toString());
+		expect(result.body).toBeUndefined();
 	});
 
 	it("convert a response with base64 encoding", async () => {
@@ -283,7 +280,7 @@ describe("convertRes", () => {
 				isBase64Encoded: true,
 			})
 		);
-		expect(await fromReadableStream(result.body)).toEqual(res.getBody().toString());
+		expect(result.body).toBeUndefined();
 	});
 });