socket_server.py 1.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041
  1. import json
  2. import asyncio
  3. from models import *
  4. from .connect_manager import ConnectionManager
  5. from .message_handler import *
  6. import time
  7. from .socket_client import socket_manager
  8. manager = ConnectionManager()
  9. active_connections = set()
  10. @app.websocket("/ws")
  11. async def websocket_endpoint(websocket: WebSocket):
  12. await manager.connect(websocket)
  13. active_connections.add(websocket)
  14. try:
  15. await socket_manager.connect()
  16. async def handler_messages():
  17. while True:
  18. byteDats = await websocket.receive()
  19. # print("byteDats", byteDats)
  20. await handlerSend(manager, json.loads(byteDats), websocket)
  21. await asyncio.gather(handler_messages())
  22. except WebSocketDisconnect:
  23. socket_manager.close()
  24. print("Client disconnected")
  25. finally:
  26. active_connections.discard(websocket)
  27. # if websocket:
  28. # await websocket.close()
  29. @app.on_event("shutdown")
  30. async def shutdown_event():
  31. print("Shutting down...")
  32. socket_manager.close()
  33. # 清理操作
  34. for connection in list(active_connections):
  35. try:
  36. await connection.close()
  37. except Exception as e:
  38. print(f"Error closing connection: {e}")