elysia 2.0.0-beta.3 → 2.0.0-beta.4

package/dist/ws/route.mjs

@@ -1,5 +1,5 @@
 import { isBun } from "../universal/constants.mjs";
-import { assignOwn, deriveEntryFn, isNotEmpty, nullObject } from "../utils.mjs";
+import { assignOwn, deriveEntryFn, isNotEmpty, isSocketQuiet, nullObject } from "../utils.mjs";
 import { isBridgeLive } from "../type/bridge.mjs";
 import { ElysiaError, ElysiaStatus, ValidationError, internalServerErrorBodyString, problemBody, problemResponse } from "../error.mjs";
 import { mapResponse } from "../adapter/web-standard/handler.mjs";
@@ -13,7 +13,7 @@ import { StandardValidator } from "../validator/index.mjs";
 import { RouteValidator } from "../validator/route.mjs";
 import { buildFrozenRouteValidator, isBridgeNotInitialized } from "../compile/handler/frozen-validator.mjs";
 import { composeRouteHook, localMacroRoot, resolveWSLocalHook } from "../compile/handler/index.mjs";
-import { ElysiaWS, isGeneratorObject } from "./context.mjs";
+import { ElysiaWS, isGeneratorObject, trackWSSettling } from "./context.mjs";
 import { createMessageParser } from "./parser.mjs";
 
 //#region src/ws/route.ts
@@ -569,6 +569,8 @@ function accumulateWSOptions(target, routeOptions, _path) {
 }
 const MAX_INFLIGHT_MESSAGES = 256;
 function buildGlobalWSHandler() {
+	let lifecycle;
+	const runLifecycle = (callback, ...args) => lifecycle ? lifecycle.run(callback, ...args) : callback(...args);
 	function getElysia(ws) {
 		let elysia = ws.data.elysia;
 		if (!elysia) {
@@ -583,9 +585,19 @@ function buildGlobalWSHandler() {
 			ws.send(typeof error?.message === "string" ? error.message : "Internal Server Error");
 		} catch {}
 	};
+	const releaseLifecycle = (ws) => {
+		if (!lifecycle?.closing && isSocketQuiet(ws)) {
+			delete ws.data["~lifecycleRun"];
+			lifecycle?.sockets.delete(ws);
+		}
+	};
 	const release = (ws) => {
-		const inflight = ws.data.inflight;
-		if (inflight) ws.data.inflight = inflight - 1;
+		const data = ws.data;
+		const inflight = data.inflight;
+		if (inflight) {
+			data.inflight = inflight - 1;
+			if (inflight === 1 && data.closeHandlerInvoked) releaseLifecycle(ws);
+		}
 	};
 	function dispatch(ws, message) {
 		let result;
@@ -602,8 +614,9 @@ function buildGlobalWSHandler() {
 		});
 		else release(ws);
 	}
-	return {
+	const handler = {
 		message(ws, message) {
+			if (lifecycle?.closing) return;
 			const data = ws.data;
 			if (!data.message) return;
 			const inflight = (data.inflight ?? 0) + 1;
@@ -616,41 +629,73 @@ function buildGlobalWSHandler() {
 			data.inflight = inflight;
 			const opening = data.opening;
 			if (opening) return void opening.then(() => {
-				if (ws.readyState > 1) return release(ws);
+				if (lifecycle?.closing || ws.readyState > 1) return release(ws);
 				dispatch(ws, message);
 			});
 			dispatch(ws, message);
 		},
 		open(ws) {
-			const result = ws.data.open?.(getElysia(ws));
+			if (lifecycle) Object.defineProperty(ws.data, "~lifecycleRun", {
+				configurable: true,
+				value: lifecycle.run
+			});
+			lifecycle?.sockets.add(ws);
+			if (lifecycle?.closing) {
+				ws.data.closeHandlerInvoked = true;
+				try {
+					ws.terminate();
+				} catch {}
+				return;
+			}
+			const open = ws.data.open;
+			const result = open && runLifecycle(open, getElysia(ws));
 			if (result instanceof Promise) {
 				const clear = () => {
 					ws.data.opening = void 0;
+					releaseLifecycle(ws);
 				};
 				ws.data.opening = result.then(clear, clear);
 			}
 		},
 		drain(ws) {
+			if (lifecycle?.closing) return;
 			const elyWs = getElysia(ws);
 			drainWaiters(elyWs);
-			ws.data.drain?.(elyWs);
+			const drain = ws.data.drain;
+			trackWSSettling(ws.data, drain && runLifecycle(drain, elyWs), () => releaseLifecycle(ws));
 		},
 		close(ws, code, reason) {
 			const elyWs = getElysia(ws);
 			drainWaiters(elyWs);
-			ws.data.opening = void 0;
-			if (ws.data.closeHandlerInvoked) return;
+			if (ws.data.closeHandlerInvoked) return releaseLifecycle(ws);
 			ws.data.closeHandlerInvoked = true;
-			const result = ws.data.close?.(elyWs, code, reason);
-			if (result instanceof Promise) result.catch(() => {});
+			let result;
+			try {
+				const close = ws.data.close;
+				result = close && runLifecycle(close, elyWs, code, reason);
+			} finally {
+				if (!(result instanceof Promise)) releaseLifecycle(ws);
+			}
+			if (result instanceof Promise) {
+				const closing = result.catch(() => {});
+				trackWSSettling(ws.data, closing, () => releaseLifecycle(ws));
+			}
 		},
 		ping(ws, data) {
-			ws.data.ping?.(getElysia(ws), data);
+			if (lifecycle?.closing) return;
+			const ping = ws.data.ping;
+			trackWSSettling(ws.data, ping && runLifecycle(ping, getElysia(ws), data), () => releaseLifecycle(ws));
 		},
 		pong(ws, data) {
-			ws.data.pong?.(getElysia(ws), data);
+			if (lifecycle?.closing) return;
+			const pong = ws.data.pong;
+			trackWSSettling(ws.data, pong && runLifecycle(pong, getElysia(ws), data), () => releaseLifecycle(ws));
 		}
 	};
+	Object.defineProperty(handler, "~lifecycle", { set(value) {
+		lifecycle = value;
+	} });
+	return handler;
 }
 
 //#endregion

package/package.json

@@ -1,7 +1,7 @@
 {
 	"name": "elysia",
 	"description": "Ergonomic Framework for Human",
-	"version": "2.0.0-beta.3",
+	"version": "2.0.0-beta.4",
 	"author": {
 		"name": "saltyAom",
 		"url": "https://github.com/SaltyAom",
@@ -244,7 +244,7 @@
 	],
 	"dependencies": {
 		"deuri": "^3.0.0",
-		"memoirist": "1.2.1"
+		"memoirist": "1.2.2"
 	},
 	"devDependencies": {
 		"@eslint/eslintrc": "^3.3.5",