ws

2023-03-10 11:26:39 +08:00
parent 363130ddd0
commit 0428a1b155
10 changed files with 270 additions and 20 deletions
--- a/main.py
+++ b/main.py
@@ -0,0 +1,88 @@
+import uvicorn
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request, Depends, HTTPException
+from fastapi.responses import HTMLResponse
+import ConnectionManager
+
+
+'''
+1. 创建一个任务, 则分发任务给worker节点 
+2. 客户端建立一个websocket连接, 提供所有的任务的状态更新
+'''
+
+
+html = """
+<!DOCTYPE html>
+<html>
+    <head>
+        <title>Chat</title>
+    </head>
+    <body>
+        <h1>WebSocket Chat</h1>
+        <h2>Your ID: <span id="ws-id"></span></h2>
+        <form action="" onsubmit="sendMessage(event)">
+            <input type="text" id="messageText" autocomplete="off"/>
+            <button>Send</button>
+        </form>
+        <ul id='messages'>
+        </ul>
+        <script>
+            var client_id = Date.now()
+            document.querySelector("#ws-id").textContent = client_id;
+            var ws = new WebSocket(`ws://localhost:8000/ws/${client_id}`);
+            ws.onmessage = function(event) {
+                var messages = document.getElementById('messages')
+                var message = document.createElement('li')
+                var content = document.createTextNode(event.data)
+                message.appendChild(content)
+                messages.appendChild(message)
+            };
+            function sendMessage(event) {
+                var input = document.getElementById("messageText")
+                ws.send(input.value)
+                input.value = ''
+                event.preventDefault()
+            }
+        </script>
+    </body>
+</html>
+"""
+
+
+app = FastAPI()
+
+
+@app.get("/")
+async def get():
+    return HTMLResponse(html)
+
+manager = ConnectionManager.ConnectionManager()
+
+
+# 接收客户端的websocket连接
+@app.websocket("/ws/{client_id}")
+async def websocket_endpoint(websocket: WebSocket, client_id: int):
+    # TODO: 验证客户端的身份(使用TOKEN)
+    # 获取TOKEN (从数据库中获取用户的信息)
+    # token = request.headers.get('Authorization')
+    # if token is None:
+    #     raise HTTPException(status_code=401, detail="Unauthorized")
+
+    await manager.connect(websocket=websocket, client_id=client_id)
+    try:
+        while True:
+            data = await websocket.receive_text()
+            await manager.send_personal_message(f"You wrote: {data}", client_id=client_id)
+            await manager.broadcast(f"Client #{client_id} says: {data}")
+            # TODO: 处理客户端的请求变化(理论上并没有)
+    except WebSocketDisconnect:
+        manager.disconnect(websocket)
+        await manager.broadcast(f"Client #{client_id} left the chat")
+
+# 通知所有的ws客户端
+@app.post("/notify")
+async def notify():
+    pass
+
+
+# 维护一个任务队列, 任务队列中的任务会被分发给worker节点
+# 任务状态变化时通知对应的客户端