From 0b8375fe82a15bfff3f517f98de6454adb2779f1 Mon Sep 17 00:00:00 2001
From: Bjorn Lu <bjornlu.dev@gmail.com>
Date: Tue, 1 Aug 2023 21:30:47 +0800
Subject: [PATCH] Fix streaming Astro components (#7895)

---
 .changeset/new-ants-speak.md                  |  5 ++++
 .../src/runtime/server/render/component.ts    | 29 ++++++++++++-------
 packages/astro/test/streaming.test.js         |  2 +-
 3 files changed, 24 insertions(+), 12 deletions(-)
 create mode 100644 .changeset/new-ants-speak.md

diff --git a/.changeset/new-ants-speak.md b/.changeset/new-ants-speak.md
new file mode 100644
index 000000000..71ba9b807
--- /dev/null
+++ b/.changeset/new-ants-speak.md
@@ -0,0 +1,5 @@
+---
+'astro': patch
+---
+
+Fix streaming Astro components
diff --git a/packages/astro/src/runtime/server/render/component.ts b/packages/astro/src/runtime/server/render/component.ts
index 33de90a26..f78f11d51 100644
--- a/packages/astro/src/runtime/server/render/component.ts
+++ b/packages/astro/src/runtime/server/render/component.ts
@@ -413,28 +413,35 @@ async function renderHTMLComponent(
 	};
 }
 
-async function renderAstroComponent(
+function renderAstroComponent(
 	result: SSRResult,
 	displayName: string,
 	Component: AstroComponentFactory,
 	props: Record<string | number, any>,
 	slots: any = {}
-): Promise<RenderInstance> {
+): RenderInstance {
 	const instance = createAstroComponentInstance(result, displayName, Component, props, slots);
 
-	// Eagerly render the component so they are rendered in parallel
-	const chunks: RenderDestinationChunk[] = [];
-	const temporaryDestination: RenderDestination = {
-		write: (chunk) => chunks.push(chunk),
+	// Eagerly render the component so they are rendered in parallel.
+	// Render to buffer for now until our returned render function is called.
+	const bufferChunks: RenderDestinationChunk[] = [];
+	const bufferDestination: RenderDestination = {
+		write: (chunk) => bufferChunks.push(chunk),
 	};
-	await instance.render(temporaryDestination);
+	// Don't await for the render to finish to not block streaming
+	const renderPromise = instance.render(bufferDestination);
 
 	return {
-		render(destination) {
-			// The real render function will simply pass on the results from the temporary destination
-			for (const chunk of chunks) {
+		async render(destination) {
+			// Write the buffered chunks to the real destination
+			for (const chunk of bufferChunks) {
 				destination.write(chunk);
 			}
+			// Free memory
+			bufferChunks.length = 0;
+			// Re-assign the real destination so `instance.render` will continue and write to the new destination
+			bufferDestination.write = (chunk) => destination.write(chunk);
+			await renderPromise;
 		},
 	};
 }
@@ -460,7 +467,7 @@ export async function renderComponent(
 	}
 
 	if (isAstroComponentFactory(Component)) {
-		return await renderAstroComponent(result, displayName, Component, props, slots);
+		return renderAstroComponent(result, displayName, Component, props, slots);
 	}
 
 	return await renderFrameworkComponent(result, displayName, Component, props, slots);
diff --git a/packages/astro/test/streaming.test.js b/packages/astro/test/streaming.test.js
index e3627d7ba..c7b835de1 100644
--- a/packages/astro/test/streaming.test.js
+++ b/packages/astro/test/streaming.test.js
@@ -48,7 +48,7 @@ describe('Streaming', () => {
 				let chunk = decoder.decode(bytes);
 				chunks.push(chunk);
 			}
-			expect(chunks.length).to.equal(2);
+			expect(chunks.length).to.equal(3);
 		});
 	});