import uvicorn
import sys
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
from ObjectManager import TaskManager, 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)
task_manager = TaskManager()
# 通知所有的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(task)
'''
监听任务进度
可能有多个客户端监听同一个任务(向任务的观察者列表中添加websocket连接)
可能有多个任务被同一客户端监听(向客户端的观察目标列表中添加任务)
应检查目标任务是否存在, 连接断开时应从观察者列表中移除
'''
@app.websocket("/tasks/{task_id}")
async def task_endpoint(websocket: WebSocket, task_id: str):
await websocket.accept()
if not task_manager.has_task(task_id):
await websocket.close()
return
task_manager.add_observer(task_id, websocket)
try:
while True:
data = await websocket.receive_text()
print(f"收到客户端消息: {data}")
except WebSocketDisconnect:
task_manager.remove_observer(task_id, websocket)
print(f"close websocket: {task_id}")
# 启动服务
if __name__ == '__main__':
port = 8000 if len(sys.argv) < 2 else int(sys.argv[1])
uvicorn.run(app='main:app', host='0.0.0.0',
port=port, reload=True, workers=1)