Merge pull request #44 from bholmesdev/feat/stream-parallel

stream: Support parallel Suspense

.changeset/large-cobras-return.md

@@ -0,0 +1,5 @@
+---
+"simple-stack-stream": minor
+---
+
+Update simple:stream internals to support parallel Suspense boundaries and fix nested Suspense edge cases.

packages/stream/components/Suspense.astro

@@ -1,14 +1,8 @@
 ---
-import type { LocalsWithStreamInternals } from "./types";
-
-const slotPromise = Astro.slots.render('default');
-
-const { stream } =  Astro.locals as LocalsWithStreamInternals;
-
-const idx = stream.components.length;
-stream.components.push(slotPromise);
+const idx = Astro.locals.suspend(Astro.slots.render("default"));
 ---
 
-<simple-suspense data-suspense-id={idx}>
+<div style="display: contents" data-suspense-fallback={idx}>
   <slot name="fallback" />
-</simple-suspense>
+</div>
+

packages/stream/components/env.d.ts

@@ -1 +1,7 @@
 /// <reference types="astro/client" />
+
+declare namespace App {
+	interface Locals {
+		suspend(promise: Promise<string>): number;
+	}
+}

packages/stream/src/env.d.ts

@@ -0,0 +1,7 @@
+/// <reference types="astro/client" />
+
+declare namespace App {
+	interface Locals {
+		suspend(promise: Promise<string>): number;
+	}
+}

packages/stream/src/middleware.ts

@@ -1,28 +1,63 @@
-import { defineMiddleware } from "astro/middleware";
+import { defineMiddleware } from "astro:middleware";
 
-export const onRequest = defineMiddleware(async ({ locals }, next) => {
-	locals.stream = {
-		components: [],
-	};
-	let response = await next();
+type SuspendedChunk = {
+	chunk: string;
+	idx: number;
+};
+
+export const onRequest = defineMiddleware(async (ctx, next) => {
+	const response = await next();
+	// ignore non-HTML responses
 	if (!response.headers.get("content-type")?.startsWith("text/html")) {
 		return response;
 	}
 
+	let streamController: ReadableStreamDefaultController<SuspendedChunk>;
+
 	async function* render() {
+		// Thank you owoce!
+		// https://gist.github.com/lubieowoce/05a4cb2e8cd252787b54b7c8a41f09fc
+		const stream = new ReadableStream<SuspendedChunk>({
+			start(controller) {
+				streamController = controller;
+			},
+		});
+
+		let curId = 0;
+		const pending = new Set<Promise<string>>();
+
+		ctx.locals.suspend = (promise) => {
+			const idx = curId++;
+			pending.add(promise);
+			promise
+				.then((chunk) => {
+					streamController.enqueue({ chunk, idx });
+					pending.delete(promise);
+				})
+				.catch((e) => {
+					streamController.error(e);
+				});
+			return idx;
+		};
+
 		// @ts-expect-error ReadableStream does not have asyncIterator
-		for await (let chunk of response.body) {
+		for await (const chunk of response.body) {
 			yield chunk;
 		}
-		for (let [idx, component] of locals.stream.components.entries()) {
-			yield `<template data-suspense-id=${JSON.stringify(
-				idx,
-			)}>${await component}</template>
+
+		if (!pending.size) return streamController.close();
+
+		// @ts-expect-error ReadableStream does not have asyncIterator
+		for await (const { chunk, idx } of stream) {
+			yield `<template data-suspense=${JSON.stringify(idx)}>${chunk}</template>
 <script>
-	const template = document.querySelector(\`template[data-suspense-id="${idx}"]\`).content;
-	const dest = document.querySelector(\`simple-suspense[data-suspense-id="${idx}"]\`);
+(() => {
+	const template = document.querySelector(\`[data-suspense="${idx}"]\`).content;
+	const dest = document.querySelector(\`[data-suspense-fallback="${idx}"]\`);
 	dest.replaceWith(template);
+})();
 </script>`;
+			if (!pending.size) return streamController.close();
 		}
 	}