This commit is contained in:
2023-03-10 11:26:39 +08:00
parent 363130ddd0
commit 0428a1b155
10 changed files with 270 additions and 20 deletions

26
ConnectionManager.py Normal file
View File

@@ -0,0 +1,26 @@
from fastapi import WebSocket
# 使用字典来存储websocket连接
class ConnectionManager:
def __init__(self):
self.active_connections: map = {}
# 加入一个新的websocket连接
async def connect(self, websocket: WebSocket, client_id: int):
await websocket.accept()
self.active_connections[client_id] = websocket
# 删除一个websocket连接(要先检查client_id是否存在)
def disconnect(self, client_id: int):
if client_id in self.active_connections:
del self.active_connections[client_id]
# 向指定的websocket连接发送消息(要先检查client_id是否存在)
async def send_personal_message(self, message: str, client_id: int):
if client_id in self.active_connections:
await self.active_connections[client_id].send_text(message)
# 向所有的websocket连接发送消息
async def broadcast(self, message: str):
for client_id, ws in self.active_connections.items():
await ws.send_text(message)

25
TaskManager.py Normal file
View File

@@ -0,0 +1,25 @@
# 使用字典来存储Task任务
# 任务的状态有:未开始、进行中、已完成
# 任务的状态可以通过TaskManager来进行修改
# 任务的状态可以通过TaskManager来进行查询
# 任务的状态可以通过TaskManager来进行删除
# 任务的状态可以通过TaskManager来进行添加
class TaskManager(object):
def __init__(self):
self.tasks = {}
def add(self, task_name):
self.tasks[task_name] = 'not started'
def delete(self, task_name):
del self.tasks[task_name]
def update(self, task_name, status):
self.tasks[task_name] = status
def query(self, task_name):
return self.tasks[task_name]
def query_all(self):
return self.tasks

88
main.py Normal file
View File

@@ -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节点
# 任务状态变化时通知对应的客户端

View File

@@ -13,17 +13,17 @@ export default defineNuxtConfig({
} }
} }
}, },
nitro: { //nitro: {
devProxy: { // devProxy: {
'/api/img': { // '/api/img': {
target: 'http://106.15.192.42:3000/api/img' // target: 'http://106.15.192.42:3000/api/img'
}, // },
'/api/drawing': { // '/api/drawing': {
target: 'http://106.15.192.42:3000/api/drawing' // target: 'http://106.15.192.42:3000/api/drawing'
}, // },
'/api/text': { // '/api/text': {
target: 'http://139.224.128.24:7861/api/text_to_text' // target: 'http://139.224.128.24:7861/api/text_to_text'
}, // },
}, // },
} //}
}) })

18
package-lock.json generated
View File

@@ -814,8 +814,7 @@
"@types/node": { "@types/node": {
"version": "18.11.17", "version": "18.11.17",
"resolved": "https://registry.npmjs.org/@types/node/-/node-18.11.17.tgz", "resolved": "https://registry.npmjs.org/@types/node/-/node-18.11.17.tgz",
"integrity": "sha512-HJSUJmni4BeDHhfzn6nF0sVmd1SMezP7/4F0Lq+aXzmp2xm9O7WXrUtHW/CHlYVtZUbByEvWidHqRtcJXGF2Ng==", "integrity": "sha512-HJSUJmni4BeDHhfzn6nF0sVmd1SMezP7/4F0Lq+aXzmp2xm9O7WXrUtHW/CHlYVtZUbByEvWidHqRtcJXGF2Ng=="
"dev": true
}, },
"@types/resolve": { "@types/resolve": {
"version": "1.20.2", "version": "1.20.2",
@@ -823,6 +822,14 @@
"integrity": "sha512-60BCwRFOZCQhDncwQdxxeOEEkbc5dIMccYLwbxsS4TUNeVECQ/pBJ0j09mrHOl/JJvpRPGwO9SvE4nR2Nb/a4Q==", "integrity": "sha512-60BCwRFOZCQhDncwQdxxeOEEkbc5dIMccYLwbxsS4TUNeVECQ/pBJ0j09mrHOl/JJvpRPGwO9SvE4nR2Nb/a4Q==",
"dev": true "dev": true
}, },
"@types/ws": {
"version": "8.5.4",
"resolved": "https://registry.npmmirror.com/@types/ws/-/ws-8.5.4.tgz",
"integrity": "sha512-zdQDHKUgcX/zBc4GrwsE/7dVdAD8JR4EuiAXiiUhhfyIJXXb2+PrGshFyeXWQPMmmZ2XxgaqclgpIC7eTXc1mg==",
"requires": {
"@types/node": "*"
}
},
"@unhead/dom": { "@unhead/dom": {
"version": "1.0.13", "version": "1.0.13",
"resolved": "https://registry.npmjs.org/@unhead/dom/-/dom-1.0.13.tgz", "resolved": "https://registry.npmjs.org/@unhead/dom/-/dom-1.0.13.tgz",
@@ -6374,10 +6381,9 @@
"dev": true "dev": true
}, },
"ws": { "ws": {
"version": "8.11.0", "version": "8.12.1",
"resolved": "https://registry.npmjs.org/ws/-/ws-8.11.0.tgz", "resolved": "https://registry.npmmirror.com/ws/-/ws-8.12.1.tgz",
"integrity": "sha512-HPG3wQd9sNQoT9xHyNCXoDUa+Xw/VevmY9FoHyQ+g+rrMn4j6FB4np7Z0OhdTgjx6MgQLK7jwSy1YecU1+4Asg==", "integrity": "sha512-1qo+M9Ba+xNhPB+YTWUlK6M17brTut5EXbcBaMRN5pH5dFrXz7lzz1ChFSUq3bOUl8yEvSenhHmYUNJxFzdJew=="
"dev": true
}, },
"xxhashjs": { "xxhashjs": {
"version": "0.2.2", "version": "0.2.2",

View File

@@ -24,6 +24,8 @@
"author": "satori.love", "author": "satori.love",
"license": "MIT", "license": "MIT",
"dependencies": { "dependencies": {
"axios": "^1.3.4" "@types/ws": "^8.5.4",
"axios": "^1.3.4",
"ws": "^8.12.1"
} }
} }

View File

@@ -328,6 +328,31 @@ const TaskSubmit = async () => {
}) })
} }
// 如果列表里有待执行的任务, 使用WS连接检查任务状态
const TaskManagement = () => {
console.log('构建WS连接')
let ws = new WebSocket('ws://'+window.location.host+'/api/sendmessage')
ws.onopen = () => {
console.log('ws open')
ws.send('hello')
}
ws.onmessage = (e) => {
console.log('ws message')
let data = JSON.parse(e.data)
console.log(data.id, data.status, data.progress)
}
ws.onerror = () => {
console.log('ws error')
}
ws.onclose = () => {
console.log('ws close')
}
}
onMounted(() => {
TaskManagement()
})
</script> </script>
<style> <style>

View File

@@ -1,4 +1,5 @@
import axios from 'axios' import axios from 'axios'
export default defineEventHandler(async event => { export default defineEventHandler(async event => {
// 获取任务列表的状态进度(普通用户只能看到自己的) // 获取任务列表的状态进度(普通用户只能看到自己的)

25
server/api/sendmessage.ts Normal file
View File

@@ -0,0 +1,25 @@
export default defineEventHandler(async (event) => {
console.log("Sendingmessage~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~")
//if (event.node.req.socket) {
// console.log("Creating websocket server")
// event.node.req.socket.on("upgrade", function (request, socket, head) {
// console.log("Upgrading")
// })
// event.node.req.socket.on("connection", function (socket) {
// console.log("Connection")
// socket.on("message", function (message: any) {
// console.log("MessageA", message)
// })
// })
// event.node.req.socket.on("message", function (message: any) {
// console.log("MessageB", message)
// })
// event.node.req.socket.on("close", function (message: any) {
// console.log("close", message)
// })
// event.node.req.socket.on("error", function (message: any) {
// console.log("error", message)
// })
//}
return { message:'2333' }
})

View File

@@ -0,0 +1,52 @@
import WebSocket, { WebSocketServer } from 'ws'
type Client = {
id: string
send: (message: string) => void
readyState: number
}
declare global {
var wss: WebSocketServer
var clients: Client[]
}
let wss: WebSocketServer
let clients: Client[] = []
export default defineEventHandler((event) => {
if (!global.wss) {
console.log("Creating websocket server")
wss = new WebSocketServer({ server: event.node.res.socket?.server })
wss.on("connection", function (socket) {
console.log("Client connected")
socket.send("connecte")
socket.on("message", function (message) {
wss.clients.forEach(function (client) {
console.log("Client", client)
client.send('message')
if (client == socket && client.readyState === WebSocket.OPEN) {
clients.push({
id: message.toString(),
send: (data: string) => client.send(data),
readyState: client.readyState,
})
global.clients = clients
}
})
})
socket.on("close", function () {
clients = clients.filter((client) => client.readyState !== WebSocket.CLOSED)
global.clients = clients
console.log("Client disconnected")
})
socket.on("error", function (error) {
console.log("Error", error)
})
socket.on("pong", function (data) {
console.log("Pong", data)
})
global.wss = wss
})
}
})