聊天系统的信息发送
最近使用python写了一个简易聊天室，对于信息的传输写的磕磕绊绊的，因此想整理一下这个过程。
后端
定义WebSocket入口
其实就是使用FastAPI的include_router将WebSocket的相关路由整合起来。
from fastapi import FastAPI
import uvicorn
from backend.api.message import router as message_router
......
from backend.websocket.handlers import router as websocket_router
from fastapi.staticfiles import StaticFiles
app = FastAPI()
app.include_router(message_router, prefix="/api" ,tags=["信息接口"])
......
app.include_router(websocket_router, tags=["WebSocket接口"])
app.mount("/frontend", StaticFiles(directory="frontend"), name="frontend")
if __name__ == '__main__':
uvicorn.run("backend.main:app", port=8000, reload=True)
管理所有连接
ConnectionManager 类是整个后端 WebSocket 机制的核心。它像一个登记处，实时追踪哪个房间有哪些用户，以及他们各自的连接。
数据结构：_RoomMap = Dict[int, Dict[int, Set[WebSocket]]]这是一个嵌套字典，一个房间可以有多个用户，一个用户可以在一个房间内有多个连接（多开）
主要方法：
connect：当一个新用户连接成功时，将其 WebSocket 对象添加到上述数据结构中。disconnect: 当用户断开连接时，将其从数据结构中移除。这里也要考虑用户的多连接，删除用户和删除用户连接是两个东西。broadcast: 这是最重要的广播方法。它会遍历指定房间内的所有连接，并将消息发送给它们。
_RoomMap = Dict[int, Dict[int, Set[WebSocket]]]
class ConnectionManager:
def __init__(self) -> None:
self.rooms: _RoomMap = defaultdict(lambda: defaultdict(set))
self._lock = asyncio.Lock()# 保护 rooms 结构的并发访问
async def connect(self, room_id: int, user_id: int, ws: WebSocket) -> None:
async with self._lock:# 加锁，防止并发修改 rooms 结构
self.rooms[room_id][user_id].add(ws)
async def disconnect(self, room_id: int, user_id: int, ws: WebSocket) -> None:
async with self._lock:
users = self.rooms.get(room_id)
if not users:
return
sockets = users.get(user_id)# 获取该用户的连接集合
if sockets and ws in sockets:
sockets.remove(ws)
if not sockets:# 如果该用户没有其他连接，移除该用户
users.pop(user_id, None)
if not users:
self.rooms.pop(room_id, None)
async def broadcast(self, room_id: int, payload: dict, exclude_user_id: Optional[int] = None) -> None:
# 不加锁读取，发送时失败的 socket 交由 disconnect 清理, 因为发送时可能会很慢
users = self.rooms.get(room_id, {})
for uid, sockets in users.items():
if exclude_user_id is not None and uid == exclude_user_id:
continue
for ws in list(sockets):
try:
await ws.send_json(payload)
except Exception:
# 发送失败时忽略；由 handlers 在断连时清理
pass
处理具体逻辑
当一个客户端连接服务器时，服务器具体该做什么？
-
建立连接与验证: 首先会进行严格的权限检查。
- 从 URL 查询参数中获取
token。 - 验证
token的有效性，并解码出user_id。 - 检查房间是否存在。
- 检查该用户是否是房间成员。
- 验证通过后，调用
manager.connect和websocket.accept注册连接，并广播一条用户上线的系统消息。
- 从 URL 查询参数中获取
-
接收、处理、广播消息: 在
try...while True循环中，服务器等待并处理客户端发来的消息。- 收消息：
await websocket.receive_text(): 等待客户端发送消息。 - 校验：服务器期望收到一个包含
action: 'send'，type: 'text'的 JSON 字符串。 - 入库：
save_message: 消息内容被保存到数据库，获得一个唯一的 ID 和时间戳。 - 广播：
manager.broadcast: 将入库后的完整消息打包成payload，调用manager进行全房间广播。
@router.websocket("/ws/{room_id}") async def ws_room(websocket: WebSocket, room_id: int): # 1) 取 token（使用 query，因为浏览器原生 WS 不能带 Authorization 头） token: Optional[str] = websocket.query_params.get("token") if not token: await _reject(websocket, 4401, "未授权：缺少 token") return # 2) 验证 token -> user_id user_id = decode_token_user_id(token) if user_id is None: await _reject(websocket, 4401, "未授权：无效 token") return # 3) 房间存在性 room = await Room.get_or_none(id=room_id) if not room: await _reject(websocket, 4404, "房间不存在") return # 4) 成员资格（public/private 都检查，保持一致策略） is_member = await RoomMember.filter(room_id=room_id, user_id=user_id).exists() if not is_member: await _reject(websocket, 4403, "你不是该房间成员") return # 5) 接受连接 & 注册 await websocket.accept() await manager.connect(room_id, user_id, websocket) # 可选：上线系统提示（不入库，仅广播） try: user = await User.get(id=user_id) await manager.broadcast(room_id, { "event": "presence", "data": {"user_id": user_id, "username": user.username, "state": "join"} }, exclude_user_id=None) except Exception: pass # 6) 主循环：收消息 -> 校验 -> 入库 -> 广播 try: while True: raw = await websocket.receive_text()# 接收文本消息 try: msg = json.loads(raw) except Exception: await websocket.send_json({"event": "error", "data": {"code": 4410, "message": "非法 JSON"}}) continue action = msg.get("action") if action == "ping": await websocket.send_json({"event": "ack", "data": {"type": "ping"}}) continue if action != "send": await websocket.send_json({"event": "error", "data": {"code": 4410, "message": "不支持的 action"}}) continue msg_type = (msg.get("type") or "text").strip() content = (msg.get("content") or "").strip() client_msg_id = msg.get("client_msg_id") # 类型与内容校验 if msg_type not in {"text"}: await websocket.send_json({"event": "error", "data": {"code": 4410, "message": "仅支持 text 消息"}}) continue if not content or len(content) > 1000: await websocket.send_json({"event": "error", "data": {"code": 4410, "message": "消息内容不能为空且 ≤ 1000 字"}}) continue # 限速 now = time.time() if not _rate_allowed(room_id, user_id, now): await websocket.send_json({"event": "error", "data": {"code": 4429, "message": "发送过于频繁"}}) continue # 再次确认仍是成员（防止中途被踢） still_member = await RoomMember.filter(room_id=room_id, user_id=user_id).exists() if not still_member: await websocket.send_json({"event": "error", "data": {"code": 4403, "message": "你已不在该房间"}}) await websocket.close(code=4403) break # 入库（严格使用 token 中的 user_id，禁用游客/自填名字） saved = await save_message( room_id=room_id, sender_id=user_id, sender_name=user.username, content=content, msg_type=msg_type, ) # 广播（用入库后的 id/时间） payload = { "event": "message", "data": { "id": saved.id, "room_id": room_id, "sender": {"id": user_id, "username": user.username}, "type": saved.msg_type, "content": saved.content, "created_at": saved.created_at.isoformat(), }, } await manager.broadcast(room_id, payload, exclude_user_id=None) # 可选：给发送者回执 if client_msg_id: try: await websocket.send_json({"event": "ack", "data": {"client_msg_id": client_msg_id, "message_id": saved.id}})#意思是给发送者回执 except Exception: pass except WebSocketDisconnect: # 客户端关闭 pass except Exception: # 兜底错误 try: await websocket.send_json({"event": "error", "data": {"code": 1011, "message": "服务器内部错误"}}) except Exception: pass finally: # 7) 断开与下线广播 await manager.disconnect(room_id, user_id, websocket) try: await manager.broadcast(room_id, { "event": "presence", "data": {"user_id": user_id, "username": user.username, "state": "leave"} }, exclude_user_id=None) except Exception: pass - 收消息：
前端
websocket.js: 底层封装，负责连接和原始数据收发。chat.js: 业务逻辑，负责调用websocket.js并处理其事件。ui.js: 视图层，负责将数据显示在页面上。
底层封装
websocketManager 对象将原生的 WebSocket API 封装成更易于使用的接口。
connect(roomId, token): 创建 WebSocket 连接，并设置好onmessage等事件处理器。sendMessage(payload): 将业务逻辑层传来的 JS 对象转换成 JSON 字符串后发送。onMessage(callback): 允许其他模块（如chat.js）注册一个回调函数，当收到消息时该回调会被执行。- 当然，还有
onPresence``onError``onClose等回调函数。
// 处理收到的消息
function handleWebSocketMessage(event) {
try {
const messageData = JSON.parse(event.data);
switch (messageData.event) {
case 'message':
onMessageCallback && onMessageCallback(messageData.data);
break;
case 'presence':
onPresenceCallback && onPresenceCallback(messageData.data);
break;
case 'error':
onErrorCallback && onErrorCallback(messageData.data);
break;
case 'ack':
console.log('Message acknowledgement:', messageData.data);
break;
}
} catch (error) {
console.error("解析WebSocket消息失败", error);
}
}
// WebSocket 管理器
export const websocketManager = {
connect(roomId, token) {
if (ws) {
ws.close();
}
ws = new WebSocket(`ws://${window.location.host}/ws/${roomId}?token=${token}`);
ws.onmessage = handleWebSocketMessage;
ws.onopen = () => console.log('WebSocket打开连接');
ws.onerror = (error) => console.error('WebSocket错误:', error);
ws.onclose = (event) => {
console.log('WebSocket连接关闭', event);
onCloseCallback && onCloseCallback(event);
};
},
disconnect() {
if (ws) {
ws.close();
ws = null;
}
},
sendMessage(payload) {
if (!ws || ws.readyState !== WebSocket.OPEN) {
throw new Error('WebSocket is not connected.');
}
ws.send(JSON.stringify(payload));
},
onMessage(callback) {
onMessageCallback = callback;
},
onPresence(callback) {
onPresenceCallback = callback;
},
onError(callback) {
onErrorCallback = callback;
},
onClose(callback) {
onCloseCallback = callback;
}
};
实现业务逻辑
-
发送消息: 监听消息输入框的
submit事件。当用户发送消息时，调用websocketManager的发送接口。messageForm.addEventListener('submit', (e) => { e.preventDefault(); const content = ui.DOMElements.messageInput.value.trim();//trim用于去除多余空格 if (!content) return ui.showAlert('消息内容不能为空', 'warning'); if (!state.currentRoomId) return ui.showAlert('请先选择一个房间', 'warning'); try { websocketManager.sendMessage({ action: 'send', type: 'text', content }); ui.DOMElements.messageInput.value = ''; } catch (error) { ui.showAlert(error.message, 'danger'); } }); -
接收消息: 在
setupWebSocketHandlers函数中，为websocketManager的onMessage事件注册了一个回调。当websocket.js收到消息并分发时，这个回调被执行，它会立即调用ui.js的方法来更新界面。function setupWebSocketHandlers() { websocketManager.onMessage(message => { ui.displayMessage(message, state.currentUser); }); websocketManager.onPresence(async () => { if (state.currentRoomId) { const onlineUsers = await api.getOnlineUsers(state.currentRoomId); ui.displayUsers(onlineUsers, ui.DOMElements.onlineList, state.currentUser); } }); websocketManager.onClose(event => { if (event.code === 4401) ui.showAlert('认证失败，请重新登录', 'danger'); else if (event.code === 4403) ui.showAlert('你不是该房间成员', 'warning'); else if (event.code === 4404) ui.showAlert('房间不存在', 'warning'); }); websocketManager.onError(error => { ui.showAlert(error.message || 'WebSocket 连接错误', 'danger'); }); }
渲染 UI
负责所有与 DOM 相关的操作
总结
- 连接: 用户点击一个房间，
chat.js调用websocketManager.connect。 - 认证:
websocket.js向后端发起连接，后端handlers.py验证token等信息，成功后将连接注册到manager。 - 发送: 用户在
message-input中输入 “你好” 并回车。chat.js的submit监听器触发，调用websocketManager.sendMessage({ action: 'send', content: '你好' })。websocket.js将其转换为 JSON 字符串"{"action":"send", ...}"并通过ws.send()发出。
- 处理: 后端
handlers.py的receive_text收到字符串。- 解析 JSON，调用
save_message将 “你好” 存入数据库。 - 调用
manager.broadcast，将包含数据库 ID、时间戳等信息的完整消息payload发送给房间内的所有连接。
- 解析 JSON，调用
- 接收与显示:
- 房间内所有客户端的
websocket.js的onmessage事件触发。 handleWebSocketMessage解析payload，发现是message事件，执行onMessageCallback。chat.js中注册的回调函数被执行，它调用ui.displayMessage。ui.js创建新的 HTML 元素，将消息 “你好” 渲染到每个用户的聊天窗口中。
- 房间内所有客户端的
扫一扫
1985
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
