simplify create_async_iterator

Author: Rich-Harris

packages/kit/src/runtime/server/page/data_serializer.js

@@ -17,9 +17,8 @@ import {
  */
 export function server_data_serializer(event, event_state, options) {
 	let promise_id = 1;
-	let count = 0;
 
-	const { iterator, push, done } = create_async_iterator();
+	const { iterator, add } = create_async_iterator();
 	const global = get_global_name(options);
 
 	/** @type {(nonce: string) => void} */
@@ -30,9 +29,8 @@ export function server_data_serializer(event, event_state, options) {
 	function replacer(thing) {
 		if (typeof thing?.then === 'function') {
 			const id = promise_id++;
-			count += 1;
 
-			thing
+			const promise = thing
 				.then(/** @param {any} data */ (data) => ({ data }))
 				.catch(
 					/** @param {any} error */ async (error) => ({
@@ -44,8 +42,6 @@ export function server_data_serializer(event, event_state, options) {
 					 * @param {{data: any; error: any}} result
 					 */
 					async ({ data, error }) => {
-						count -= 1;
-
 						let str;
 						try {
 							str = devalue.uneval(error ? [, error] : [data], replacer);
@@ -60,13 +56,12 @@ export function server_data_serializer(event, event_state, options) {
 							str = devalue.uneval([, error], replacer);
 						}
 
-						push(
-							`<script${await nonce}>${global}.resolve(${id}, ${str.includes('app.decode') ? `(app) => ${str}` : `() => ${str}`})</script>\n`
-						);
-						if (count === 0) done();
+						return `<script${await nonce}>${global}.resolve(${id}, ${str.includes('app.decode') ? `(app) => ${str}` : `() => ${str}`})</script>\n`;
 					}
 				);
 
+			add(promise);
+
 			return `${global}.defer(${id})`;
 		} else {
 			for (const key in options.hooks.transport) {
@@ -124,57 +119,56 @@ export function server_data_serializer(event, event_state, options) {
  */
 export function server_data_serializer_json(event, event_state, options) {
 	let promise_id = 1;
-	let count = 0;
 
-	const { iterator, push, done } = create_async_iterator();
+	const { iterator, add } = create_async_iterator();
 
 	const reducers = {
 		...Object.fromEntries(
 			Object.entries(options.hooks.transport).map(([key, value]) => [key, value.encode])
 		),
 		/** @param {any} thing */
 		Promise: (thing) => {
-			if (typeof thing?.then === 'function') {
-				const id = promise_id++;
-				count += 1;
+			if (typeof thing?.then !== 'function') {
+				return;
+			}
+
+			const id = promise_id++;
+
+			/** @type {'data' | 'error'} */
+			let key = 'data';
 
-				/** @type {'data' | 'error'} */
-				let key = 'data';
+			const promise = thing
+				.catch(
+					/** @param {any} e */ async (e) => {
+						key = 'error';
+						return handle_error_and_jsonify(event, event_state, options, /** @type {any} */ (e));
+					}
+				)
+				.then(
+					/** @param {any} value */
+					async (value) => {
+						let str;
+						try {
+							str = devalue.stringify(value, reducers);
+						} catch {
+							const error = await handle_error_and_jsonify(
+								event,
+								event_state,
+								options,
+								new Error(`Failed to serialize promise while rendering ${event.route.id}`)
+							);
 
-				thing
-					.catch(
-						/** @param {any} e */ async (e) => {
 							key = 'error';
-							return handle_error_and_jsonify(event, event_state, options, /** @type {any} */ (e));
+							str = devalue.stringify(error, reducers);
 						}
-					)
-					.then(
-						/** @param {any} value */
-						async (value) => {
-							let str;
-							try {
-								str = devalue.stringify(value, reducers);
-							} catch {
-								const error = await handle_error_and_jsonify(
-									event,
-									event_state,
-									options,
-									new Error(`Failed to serialize promise while rendering ${event.route.id}`)
-								);
-
-								key = 'error';
-								str = devalue.stringify(error, reducers);
-							}
-
-							count -= 1;
-
-							push(`{"type":"chunk","id":${id},"${key}":${str}}\n`);
-							if (count === 0) done();
-						}
-					);
 
-				return id;
-			}
+						return `{"type":"chunk","id":${id},"${key}":${str}}\n`;
+					}
+				);
+
+			add(promise);
+
+			return id;
 		}
 	};
 

packages/kit/src/utils/streaming.js

@@ -18,11 +18,12 @@ function defer() {
  * Create an async iterator and a function to push values into it
  * @returns {{
  *   iterator: AsyncIterable<any>;
- *   push: (value: any) => void;
- *   done: () => void;
+ *   add: (promise: Promise<any>) => void;
  * }}
  */
 export function create_async_iterator() {
+	let count = 0;
+
 	const deferred = [defer()];
 
 	return {
@@ -37,15 +38,20 @@ export function create_async_iterator() {
 				};
 			}
 		},
-		push: (value) => {
-			deferred[deferred.length - 1].fulfil({
-				value,
-				done: false
+		add: (promise) => {
+			count += 1;
+
+			void promise.then((value) => {
+				deferred[deferred.length - 1].fulfil({
+					value,
+					done: false
+				});
+				deferred.push(defer());
+
+				if (--count === 0) {
+					deferred[deferred.length - 1].fulfil({ done: true });
+				}
 			});
-			deferred.push(defer());
-		},
-		done: () => {
-			deferred[deferred.length - 1].fulfil({ done: true });
 		}
 	};
 }

packages/kit/src/utils/streaming.spec.js

@@ -2,14 +2,13 @@ import { expect, test } from 'vitest';
 import { create_async_iterator } from './streaming.js';
 
 test('works with fast consecutive promise resolutions', async () => {
-	const iterator = create_async_iterator();
+	const { iterator, add } = create_async_iterator();
 
-	void Promise.resolve(1).then((n) => iterator.push(n));
-	void Promise.resolve(2).then((n) => iterator.push(n));
-	void Promise.resolve().then(() => iterator.done());
+	add(Promise.resolve(1));
+	add(Promise.resolve(2));
 
 	const actual = [];
-	for await (const value of iterator.iterator) {
+	for await (const value of iterator) {
 		actual.push(value);
 	}