Merge pull request #549 from aleclarson/fix/plugin-lazyload

fix: await the `.modules` of nested Elysia instances

src/index.ts

@@ -13,7 +13,7 @@ import { sucrose, sucroseTrace, type Sucrose } from './sucrose'
 import { ElysiaWS, websocket } from './ws'
 import type { WS } from './ws/types'
 
-import { mergeDeep } from './utils'
+import { mergeDeep, PromiseGroup } from './utils'
 import {
 	composeHandler,
 	composeGeneralHandler,
@@ -261,9 +261,7 @@ export default class Elysia<
 		trace: Sucrose.TraceInference
 	}
 
-	private lazyLoadModules: Promise<
-		Elysia<any, any, any, any, any, any, any, any>
-	>[] = []
+	private promisedModules = new PromiseGroup()
 
 	constructor(config?: ElysiaConfig<BasePath, Scoped>) {
 		this.config = {
@@ -2540,7 +2538,7 @@ export default class Elysia<
 		}
 
 		if (plugin instanceof Promise) {
-			this.lazyLoadModules.push(
+			this.promisedModules.add(
 				plugin
 					.then((plugin) => {
 						if (typeof plugin === 'function') {
@@ -2575,7 +2573,7 @@ export default class Elysia<
 		if (typeof plugin === 'function') {
 			const instance = plugin(this as unknown as any) as unknown as any
 			if (instance instanceof Promise) {
-				this.lazyLoadModules.push(
+				this.promisedModules.add(
 					instance
 						.then((plugin) => {
 							if (plugin instanceof Elysia) {
@@ -2632,6 +2630,15 @@ export default class Elysia<
 			return instance
 		}
 
+		if (plugin.promisedModules.size) {
+			this.promisedModules.add(
+				plugin.modules
+					.then(() => this._use(plugin))
+					.then((x) => x.compile())
+			)
+			return this
+		}
+
 		const { name, seed } = plugin.config
 
 		plugin.getServer = () => this.getServer()
@@ -4908,7 +4915,7 @@ export default class Elysia<
 			}
 		})
 
-		Promise.all(this.lazyLoadModules).then(() => {
+		this.promisedModules.then(() => {
 			Bun?.gc(false)
 		})
 
@@ -4950,7 +4957,7 @@ export default class Elysia<
 	 * Wait until all lazy loaded modules all load is fully
 	 */
 	get modules() {
-		return Promise.all(this.lazyLoadModules)
+		return Promise.resolve(this.promisedModules)
 	}
 }
 

src/utils.ts

@@ -719,3 +719,53 @@ export const isNumericString = (message: string): boolean => {
 
 	return false
 }
+
+export class PromiseGroup implements PromiseLike<void> {
+	private root: Promise<any> | null = null
+	private promises: Promise<any>[] = []
+
+	constructor(public onError: (error: any) => void = console.error) {}
+
+	/**
+	 * The number of promises still being awaited.
+	 */
+	get size() {
+		return this.promises.length
+	}
+
+	/**
+	 * Add a promise to the group.
+	 * @returns The promise that was added.
+	 */
+	add<T>(promise: Promise<T>) {
+		this.promises.push(promise)
+		this.root ||= this.drain()
+		return promise
+	}
+
+	private async drain() {
+		while (this.promises.length > 0) {
+			try {
+				await this.promises[0]
+				this.promises.shift()
+			} catch (error) {
+				this.onError(error)
+			}
+		}
+		this.root = null
+	}
+
+	// Allow the group to be awaited.
+	then<TResult1 = void, TResult2 = never>(
+		onfulfilled?:
+			| ((value: void) => TResult1 | PromiseLike<TResult1>)
+			| undefined
+			| null,
+		onrejected?:
+			| ((reason: any) => TResult2 | PromiseLike<TResult2>)
+			| undefined
+			| null
+	): PromiseLike<TResult1 | TResult2> {
+		return (this.root ?? Promise.resolve()).then(onfulfilled, onrejected)
+	}
+}