import uvicorn from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request, Depends, HTTPException from fastapi.responses import HTMLResponse from ObjectManager import ConnectionManager, TaskManager, ServerManager, Task from typing import List ''' 1. 创建一个任务, 则分发任务给worker节点 2. 客户端建立一个websocket连接, 提供所有的任务的状态更新 ''' html = """ Chat

WebSocket Chat

Your ID:

""" app = FastAPI() @app.get("/") async def get(): return HTMLResponse(html) connection_manager = ConnectionManager() task_manager = TaskManager() # 接收客户端的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 connection_manager.connect(websocket=websocket, client_id=client_id) try: while True: data = await websocket.receive_text() await connection_manager.send_personal_message(f"You wrote: {data}", client_id=client_id) await connection_manager.broadcast(f"Client #{client_id} says: {data}") # TODO: 处理客户端的请求变化(理论上并没有) except WebSocketDisconnect: connection_manager.disconnect(client_id=client_id) await connection_manager.broadcast(f"Client #{client_id} left the chat") # 通知所有的ws客户端 @app.post("/notify") async def notify(): pass @app.get("/tasks", response_model=List[Task]) async def get_tasks(): return task_manager.query_all() @app.post("/tasks", response_model=Task) async def create_task(task: Task): return task_manager.add(task) # 维护一个任务队列, 任务队列中的任务会被分发给worker节点 # 任务状态变化时通知对应的客户端