From da47c2dfa38f7c742ef30c3758392795967edfde Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Fri, 7 Jun 2024 17:35:01 -0700
Subject: [PATCH] refac: active users

---
 backend/apps/socket/main.py | 23 +++++++++++++++++++----
 1 file changed, 19 insertions(+), 4 deletions(-)

diff --git a/backend/apps/socket/main.py b/backend/apps/socket/main.py
index 0bc45287a..debbe22ea 100644
--- a/backend/apps/socket/main.py
+++ b/backend/apps/socket/main.py
@@ -10,7 +10,7 @@ app = socketio.ASGIApp(sio, socketio_path="/ws/socket.io")
 
 # Dictionary to maintain the user pool
 
-
+SESSION_POOL = {}
 USER_POOL = {}
 USAGE_POOL = {}
 # Timeout duration in seconds
@@ -29,7 +29,12 @@ async def connect(sid, environ, auth):
             user = Users.get_user_by_id(data["id"])
 
         if user:
-            USER_POOL[sid] = user.id
+            SESSION_POOL[sid] = user.id
+            if user.id in USER_POOL:
+                USER_POOL[user.id].append(sid)
+            else:
+                USER_POOL[user.id] = [sid]
+
             print(f"user {user.name}({user.id}) connected with session ID {sid}")
 
             print(len(set(USER_POOL)))
@@ -50,7 +55,13 @@ async def user_join(sid, data):
             user = Users.get_user_by_id(data["id"])
 
         if user:
-            USER_POOL[sid] = user.id
+
+            SESSION_POOL[sid] = user.id
+            if user.id in USER_POOL:
+                USER_POOL[user.id].append(sid)
+            else:
+                USER_POOL[user.id] = [sid]
+
             print(f"user {user.name}({user.id}) connected with session ID {sid}")
 
             print(len(set(USER_POOL)))
@@ -124,7 +135,11 @@ async def remove_after_timeout(sid, model_id):
 @sio.event
 async def disconnect(sid):
     if sid in USER_POOL:
-        disconnected_user = USER_POOL.pop(sid)
+        disconnected_user = SESSION_POOL.pop(sid)
+        USER_POOL[disconnected_user].remove(sid)
+        if len(USER_POOL[disconnected_user]) == 0:
+            del USER_POOL[disconnected_user]
+
         print(f"user {disconnected_user} disconnected with session ID {sid}")
 
         await sio.emit("user-count", {"count": len(USER_POOL)})