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)})