feat: better ws cleanup

backend/controllers/assembly.ts

@@ -4,6 +4,7 @@ import { User } from "../models/user";
 import { Votation, Option } from "../models/vote";
 import { RequestWithNtnuiNo } from "../utils/request";
 import { AssemblyResponseType } from "../types/assembly";
+import { organizerConnections } from "../utils/socketNotifier";
 
 export async function createAssembly(req: RequestWithNtnuiNo, res: Response) {
   if (!req.ntnuiNo) {
@@ -109,7 +110,10 @@ export async function deleteAssembly(req: RequestWithNtnuiNo, res: Response) {
         }
         await Votation.findByIdAndDelete(vote);
       });
+
       await Assembly.deleteOne({ _id: assembly._id });
+      // Clean up websocket connections for this assembly
+      organizerConnections.delete(group);
       return res.status(200).json({ message: "Assembly successfully deleted" });
     }
   }

backend/utils/socketNotifier.ts

@@ -16,6 +16,8 @@ export const organizerConnections = new Map<string, Map<number, WebSocket>>();
 // Store all active participant connections, for access when sending messages about assembly.
 export const lobbyConnections = new Map<number, WebSocket>();
 
+export const waitingForPongLobby = new Map<number, boolean>();
+
 const sendPing = (ws: WebSocket) => {
   if (ws.readyState === WebSocket.OPEN) {
     ws.ping();
@@ -24,7 +26,14 @@ const sendPing = (ws: WebSocket) => {
 
 // Send ping to all participants to check if they are still connected and prevent the connection from closing.
 export const startHeartbeatInterval = setInterval(() => {
-  lobbyConnections.forEach((ws: WebSocket) => {
+  lobbyConnections.forEach((ws: WebSocket, userID: number) => {
+    // If the participant has not responded to the last ping, close and remove the connection.
+    if (waitingForPongLobby.get(userID) === true) {
+      lobbyConnections.delete(userID);
+      ws.close();
+      return;
+    }
+    waitingForPongLobby.set(userID, true);
     sendPing(ws);
   });
 
@@ -48,6 +57,7 @@ export const storeLobbyConnectionByCookie = (
     }
     // Store socket connection on NTNUI ID.
     lobbyConnections.set(ntnuiNo, ws);
+    console.log("User " + ntnuiNo + " connected to lobby");
   }
 };
 

backend/wsServers/lobby.ts

@@ -1,8 +1,9 @@
 import { WebSocketServer } from "ws";
 import {
-  removeLobbyConnectionByCookie,
   storeLobbyConnectionByCookie,
+  waitingForPongLobby,
 } from "../utils/socketNotifier";
+import { NTNUINoFromRequest } from "../utils/wsCookieRetriever";
 
 export const lobbyWss = new WebSocketServer({ noServer: true });
 
@@ -12,9 +13,6 @@ lobbyWss.on("connection", function connection(ws, req) {
 
   ws.on("pong", () => {
     // The client responded to the ping, so the connection is still active.
-  });
-
-  ws.on("close", () => {
-    removeLobbyConnectionByCookie(req);
+    waitingForPongLobby.set(NTNUINoFromRequest(req) || 0, false);
   });
 });

backend/wsServers/organizer.ts

@@ -52,8 +52,4 @@ organizerWss.on("connection", function connection(ws, req) {
   ws.on("pong", () => {
     // The client responded to the ping, so the connection is still active.
   });
-
-  ws.on("close", () => {
-    removeOrganizerConnectionByCookie(req);
-  });
 });