diff --git a/backend/psn_server.py b/backend/psn_server.py index 93abb54..60847be 100644 --- a/backend/psn_server.py +++ b/backend/psn_server.py @@ -3,10 +3,11 @@ import json import logging import socket import time +import weakref from dataclasses import dataclass import psn -from aiohttp import web +from aiohttp import web, WSCloseCode PSN_DEFAULT_UDP_PORT = 56565 PSN_DEFAULT_UDP_MCAST_ADDRESS = "236.10.10.10" @@ -88,18 +89,20 @@ async def handle_websocket(request): elif msg.type == web.WSMsgType.ERROR: logging.error("ws connection closed with exception %s" % ws.exception()) - print("ws connection closed with exception %s" % ws.exception()) + except Exception as e: logging.error(f"Websocket exception: {e}") finally: logging.debug("Websocket connection closing") - await ws.close() - - request.app["ws_clients"].remove(ws) + request.app["ws_clients"].discard(ws) return ws +async def on_shutdown(app): + for ws in set(app["ws_clients"]): + await ws.close(code=WSCloseCode.GOING_AWAY, message="Server shutdown") + async def handle_root(request): return web.FileResponse("./static/index.html") @@ -133,13 +136,14 @@ def create_app(): app.router.add_static("/", "./static") # Setup app state - app["ws_clients"] = set() + app["ws_clients"] = weakref.WeakSet() app["trackers"] = {} app["sock"] = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) for i in range(NUM_TRACKERS): app["trackers"][i] = TrackerData(i, 0, 0) + app.on_shutdown.append(on_shutdown) app.cleanup_ctx.append(background_tasks) return app