socket_server.py 1.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152
  1. import json
  2. import asyncio
  3. from models import *
  4. from .connect_manager import ConnectionManager
  5. def jsonMessage(code=0, msg="", data: object = None):
  6. """json字符串数据"""
  7. jsonData = {"code": code, "msg": msg, "data": data}
  8. return json.dumps(jsonData)
  9. manager = ConnectionManager()
  10. active_connections = set()
  11. @app.websocket("/ws")
  12. async def websocket_endpoint(websocket: WebSocket):
  13. await manager.connect(websocket)
  14. active_connections.add(websocket)
  15. try:
  16. async def handler_messages():
  17. while True:
  18. data = await websocket.receive_json()
  19. await handlerSend(data,websocket)
  20. await asyncio.gather( handler_messages(),)
  21. except WebSocketDisconnect:
  22. print("Client disconnected")
  23. finally:
  24. active_connections.discard(websocket)
  25. # if websocket:
  26. # await websocket.close()
  27. @app.on_event("shutdown")
  28. async def shutdown_event():
  29. print("Shutting down...")
  30. # 清理操作
  31. for connection in list(active_connections):
  32. try:
  33. await connection.close()
  34. except Exception as e:
  35. print(f"Error closing connection: {e}")
  36. async def handlerSend(receiveData,websocket):
  37. # 处理消息发送逻辑
  38. jsonType = receiveData.get("type")
  39. match jsonType:
  40. case 'ping':
  41. '''发送心跳'''
  42. data = jsonMessage("pong")
  43. await manager.send_personal_message(data,websocket)
  44. case 'pong':
  45. '''发送心跳'''
  46. pass