Merge pull request #966 from spacebarchat/feat/refactorIdentify

Rewrite identify handler

3 files changed, 53 insertions, 82 deletions

diff --git a/scripts/benchmark/connections.js b/scripts/benchmark/connections.js
deleted file mode 100644
index 4246c646..00000000
--- a/scripts/benchmark/connections.js
+++ /dev/null
@@ -1,82 +0,0 @@
-/*
-	Spacebar: A FOSS re-implementation and extension of the Discord.com backend.
-	Copyright (C) 2023 Spacebar and Spacebar Contributors
-	
-	This program is free software: you can redistribute it and/or modify
-	it under the terms of the GNU Affero General Public License as published
-	by the Free Software Foundation, either version 3 of the License, or
-	(at your option) any later version.
-	
-	This program is distributed in the hope that it will be useful,
-	but WITHOUT ANY WARRANTY; without even the implied warranty of
-	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-	GNU Affero General Public License for more details.
-	
-	You should have received a copy of the GNU Affero General Public License
-	along with this program.  If not, see <https://www.gnu.org/licenses/>.
-*/
-
-require("dotenv").config();
-const cluster = require("cluster");
-const WebSocket = require("ws");
-const endpoint = process.env.GATEWAY || "ws://localhost:3001";
-const connections = Number(process.env.CONNECTIONS) || 50;
-const token = process.env.TOKEN;
-var cores = 1;
-try {
-	cores = Number(process.env.THREADS) || os.cpus().length;
-} catch {
-	console.log("[Bundle] Failed to get thread count! Using 1...");
-}
-
-if (!token) {
-	console.error("TOKEN env var missing");
-	process.exit();
-}
-
-if (cluster.isMaster) {
-	for (let i = 0; i < cores; i++) {
-		cluster.fork();
-	}
-
-	cluster.on("exit", (worker, code, signal) => {
-		console.log(`worker ${worker.process.pid} died`);
-	});
-} else {
-	for (let i = 0; i < connections; i++) {
-		connect();
-	}
-}
-
-function connect() {
-	const client = new WebSocket(endpoint);
-	client.on("message", (data) => {
-		data = JSON.parse(data);
-
-		switch (data.op) {
-			case 10:
-				client.interval = setInterval(() => {
-					client.send(JSON.stringify({ op: 1 }));
-				}, data.d.heartbeat_interval);
-
-				client.send(
-					JSON.stringify({
-						op: 2,
-						d: {
-							token,
-							properties: {},
-						},
-					}),
-				);
-
-				break;
-		}
-	});
-	client.once("close", (code, reason) => {
-		clearInterval(client.interval);
-		connect();
-	});
-	client.on("error", (err) => {
-		// console.log(err);
-	});
-}
diff --git a/scripts/stress/identify.js b/scripts/stress/identify.js
new file mode 100644
index 00000000..7dea5e4d
--- /dev/null
+++ b/scripts/stress/identify.js
@@ -0,0 +1,53 @@
+/* eslint-env node */
+
+require("dotenv").config();
+const { OPCODES } = require("../../dist/gateway/util/Constants.js");
+const WebSocket = require("ws");
+const ENDPOINT = `ws://localhost:3002?v=9&encoding=json`;
+const TOKEN =
+	"eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpZCI6IjEwOTMxMTgwMjgzNjA1MzYxMDYiLCJpYXQiOjE2ODA2OTE5MDB9.9ByCqDvC4mIutW8nM7WhVCtGuKW08UimPnmBeNw-K0E";
+const TOTAL_ITERATIONS = 500;
+
+const doTimedIdentify = () =>
+	new Promise((resolve) => {
+		let start;
+		const ws = new WebSocket(ENDPOINT);
+		ws.on("message", (data) => {
+			const parsed = JSON.parse(data);
+
+			switch (parsed.op) {
+				case OPCODES.Hello:
+					// send identify
+					start = performance.now();
+					ws.send(
+						JSON.stringify({
+							op: OPCODES.Identify,
+							d: {
+								token: TOKEN,
+								properties: {},
+							},
+						}),
+					);
+					break;
+				case OPCODES.Dispatch:
+					if (parsed.t == "READY") {
+						ws.close();
+						return resolve(performance.now() - start);
+					}
+
+					break;
+			}
+		});
+	});
+
+(async () => {
+	const perfs = [];
+	while (perfs.length < TOTAL_ITERATIONS) {
+		const ret = await doTimedIdentify();
+		perfs.push(ret);
+		// console.log(`${perfs.length}/${TOTAL_ITERATIONS} - this: ${Math.floor(ret)}ms`)
+	}
+
+	const avg = perfs.reduce((prev, curr) => prev + curr) / (perfs.length - 1);
+	console.log(`Average identify time: ${Math.floor(avg * 100) / 100}ms`);
+})();
diff --git a/scripts/benchmark/users.js b/scripts/stress/users.js
index 20f9f7c3..20f9f7c3 100644
--- a/scripts/benchmark/users.js
+++ b/scripts/stress/users.js