diff --git a/src/index.ts b/src/index.ts index 87a25eb8..c8b7d6ed 100644 --- a/src/index.ts +++ b/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 - >[] = [] + private promisedModules = new PromiseGroup() constructor(config?: ElysiaConfig) { 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) } } diff --git a/src/utils.ts b/src/utils.ts index 818bad79..d56699f4 100644 --- a/src/utils.ts +++ b/src/utils.ts @@ -719,3 +719,53 @@ export const isNumericString = (message: string): boolean => { return false } + +export class PromiseGroup implements PromiseLike { + private root: Promise | null = null + private promises: Promise[] = [] + + 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(promise: Promise) { + 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( + onfulfilled?: + | ((value: void) => TResult1 | PromiseLike) + | undefined + | null, + onrejected?: + | ((reason: any) => TResult2 | PromiseLike) + | undefined + | null + ): PromiseLike { + return (this.root ?? Promise.resolve()).then(onfulfilled, onrejected) + } +}