refactor: improve readability of connection manager code

backend/server/src/ConnectionManager.ts

@@ -18,31 +18,39 @@ class Counter {
 
 // Owner Connection Management
 export class ConnectionManager {
+    // Counts how many times the owner is connected to a sandbox
     private ownerConnections: Record<string, Counter> = {}
+    // Stores all sockets connected to a given sandbox
     private sockets: Record<string, Set<Socket>> = {}
 
-    ownerConnected(sandboxId: string) {
+    // Checks if the owner of a sandbox is connected
-        this.ownerConnections[sandboxId] ??= new Counter()
-        this.ownerConnections[sandboxId].increment()
-    }
-
-    ownerDisconnected(sandboxId: string) {
-        this.ownerConnections[sandboxId]?.decrement()
-    }
-
     ownerIsConnected(sandboxId: string): boolean {
         return this.ownerConnections[sandboxId]?.getValue() > 0
     }
 
-    addConnectionForSandbox(socket: Socket, sandboxId: string) {
+    // Adds a connection for a sandbox
+    addConnectionForSandbox(socket: Socket, sandboxId: string, isOwner: boolean) {
         this.sockets[sandboxId] ??= new Set()
         this.sockets[sandboxId].add(socket)
+
+        // If the connection is for the owner, increments the owner connection counter
+        if (isOwner) {
+            this.ownerConnections[sandboxId] ??= new Counter()
+            this.ownerConnections[sandboxId].increment()
+        }
     }
 
-    removeConnectionForSandbox(socket: Socket, sandboxId: string) {
+    // Removes a connection for a sandbox
+    removeConnectionForSandbox(socket: Socket, sandboxId: string, isOwner: boolean) {
         this.sockets[sandboxId]?.delete(socket)
+
+        // If the connection being removed is for the owner, decrements the owner connection counter
+        if (isOwner) {
+            this.ownerConnections[sandboxId]?.decrement()
+        }
     }
 
+    // Returns the set of sockets connected to a given sandbox
     connectionsForSandbox(sandboxId: string): Set<Socket> {
         return this.sockets[sandboxId] ?? new Set();
     }

backend/server/src/index.ts

@@ -98,16 +98,14 @@ io.on("connection", async (socket) => {
       isOwner: boolean
     }
 
+    // Register the connection
+    connections.addConnectionForSandbox(socket, data.sandboxId, data.isOwner)
+
     // Disable access unless the sandbox owner is connected
-    if (data.isOwner) {
+    if (!data.isOwner && !connections.ownerIsConnected(data.sandboxId)) {
-      connections.ownerConnected(data.sandboxId)
+      socket.emit("disableAccess", "The sandbox owner is not connected.")
-    } else {
+      return
-      if (!connections.ownerIsConnected(data.sandboxId)) {
-        socket.emit("disableAccess", "The sandbox owner is not connected.")
-        return
-      }
     }
-    connections.addConnectionForSandbox(socket, data.sandboxId)
 
     try {
       // Create or retrieve the sandbox manager for the given sandbox ID
@@ -119,6 +117,7 @@ io.on("connection", async (socket) => {
       )
       sandboxes[data.sandboxId] = sandboxManager
 
+      // This callback recieves an update when the file list changes, and notifies all relevant connections.
       const sendFileNotifications = (files: (TFolder | TFile)[]) => {
         connections.connectionsForSandbox(data.sandboxId).forEach((socket: Socket) => {
           socket.emit("loaded", files);
@@ -131,6 +130,8 @@ io.on("connection", async (socket) => {
       socket.emit("loaded", sandboxManager.fileManager?.files)
 
       // Register event handlers for the sandbox
+      // For each event handler, listen on the socket for that event
+      // Pass connection-specific information to the handlers
       Object.entries(sandboxManager.handlers({
         userId: data.userId,
         isOwner: data.isOwner,
@@ -149,19 +150,17 @@ io.on("connection", async (socket) => {
       // Handle disconnection event
       socket.on("disconnect", async () => {
         try {
-          connections.removeConnectionForSandbox(socket, data.sandboxId)
+          // Deregister the connection
+          connections.removeConnectionForSandbox(socket, data.sandboxId, data.isOwner)
 
-          if (data.isOwner) {
+          // If the owner has disconnected from all sockets, close open terminals and file watchers.o
-            connections.ownerDisconnected(data.sandboxId)
+          // The sandbox itself will timeout after the heartbeat stops.
-            // If the owner has disconnected from all sockets, close open terminals and file watchers.o
+          if (data.isOwner && !connections.ownerIsConnected(data.sandboxId)) {
-            // The sandbox itself will timeout after the heartbeat stops.
+            await sandboxManager.disconnect()
-            if (!connections.ownerIsConnected(data.sandboxId)) {
+            socket.broadcast.emit(
-              await sandboxManager.disconnect()
+              "disableAccess",
-              socket.broadcast.emit(
+              "The sandbox owner has disconnected."
-                "disableAccess",
+            )
-                "The sandbox owner has disconnected."
-              )
-            }
           }
         } catch (e: any) {
           handleErrors("Error disconnecting:", e, socket);