refactor: improve readability of the 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)
-    } else {
-      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,20 +150,18 @@ 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) {
-            connections.ownerDisconnected(data.sandboxId)
           // If the owner has disconnected from all sockets, close open terminals and file watchers.o
           // The sandbox itself will timeout after the heartbeat stops.
-            if (!connections.ownerIsConnected(data.sandboxId)) {
+          if (data.isOwner && !connections.ownerIsConnected(data.sandboxId)) {
             await sandboxManager.disconnect()
             socket.broadcast.emit(
               "disableAccess",
               "The sandbox owner has disconnected."
             )
           }
-          }
         } catch (e: any) {
           handleErrors("Error disconnecting:", e, socket);
         }