From 6b2b8700200d99b2ed257e43042cc6b08e7e12cc Mon Sep 17 00:00:00 2001
From: James Murdza <james@jamesmurdza.com>
Date: Fri, 25 Oct 2024 19:23:11 -0600
Subject: [PATCH] refactor: improve readability of the connection manager code

---
 backend/server/src/ConnectionManager.ts | 30 ++++++++++++-------
 backend/server/src/index.ts             | 39 ++++++++++++-------------
 2 files changed, 38 insertions(+), 31 deletions(-)

diff --git a/backend/server/src/ConnectionManager.ts b/backend/server/src/ConnectionManager.ts
index 51683d4..45b5432 100644
--- a/backend/server/src/ConnectionManager.ts
+++ b/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) {
-        this.ownerConnections[sandboxId] ??= new Counter()
-        this.ownerConnections[sandboxId].increment()
-    }
-
-    ownerDisconnected(sandboxId: string) {
-        this.ownerConnections[sandboxId]?.decrement()
-    }
-
+    // Checks if the owner of a sandbox is connected
     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();
     }
diff --git a/backend/server/src/index.ts b/backend/server/src/index.ts
index 08ede15..ac84fc5 100644
--- a/backend/server/src/index.ts
+++ b/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) {
-      connections.ownerConnected(data.sandboxId)
-    } else {
-      if (!connections.ownerIsConnected(data.sandboxId)) {
-        socket.emit("disableAccess", "The sandbox owner is not connected.")
-        return
-      }
+    if (!data.isOwner && !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) {
-            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)) {
-              await sandboxManager.disconnect()
-              socket.broadcast.emit(
-                "disableAccess",
-                "The sandbox owner has disconnected."
-              )
-            }
+          // 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 (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);