WS change

loon/web/users/router.py

@@ -1,12 +1,9 @@
-import asyncio
-import threading
-
 from fastapi import APIRouter, Request
 from starlette.websockets import WebSocket, WebSocketDisconnect
 
 from loon.web import get_mqtt_manager
 from loon.web.auth.middleware import authenticated
-from loon.web.users.state import user_threads, user_world_requests
+from loon.web.users.state import subscribe_player, get_user_messages, unsubscribe_player
 
 router = APIRouter(prefix="/users", tags=["users"])
 
@@ -47,16 +44,12 @@ async def websocket_endpoint(websocket: WebSocket):
     await websocket.accept()
     user = websocket.user.user
 
-    queue = asyncio.Queue()
-    user_threads[user.uuid] = queue
-    user_world_requests[user.uuid] = set()
+    subscribe_player(user)
 
     try:
-        while True:
-            message = await queue.get()
+        async for message in get_user_messages(user):
             await websocket.send_text(message)
     except WebSocketDisconnect:
         pass
     finally:
-        user_threads.pop(user.uuid, None)
-        user_world_requests.pop(user.uuid, None)
+        unsubscribe_player(user)

loon/web/users/state.py

@@ -1,2 +1,38 @@
+import asyncio
+
+from loon.web.users.models import User
+
 user_threads = {}
 user_world_requests = {}
+
+
+async def send_to_player(user: User, msg):
+    queue = user_threads.get(user.uuid)
+    if queue is not None:
+        await queue.put(msg)
+
+
+def get_player_thread(user: User):
+    return user_threads.get(user.uuid)
+
+
+async def get_user_messages(user: User):
+    queue = user_threads.get(user.uuid)
+    if queue is None:
+        return
+
+    while True:
+        yield await queue.get()
+
+
+def subscribe_player(user: User):
+    queue = asyncio.Queue()
+    user_threads[user.uuid] = queue
+    user_world_requests[user.uuid] = set()
+
+    return user_threads[user.uuid]
+
+
+def unsubscribe_player(user: User):
+    user_threads.pop(user.uuid, None)
+    user_world_requests.pop(user.uuid, None)