diff --git a/app/signalr/hub/hub.py b/app/signalr/hub/hub.py index 3c8fe98..52d6840 100644 --- a/app/signalr/hub/hub.py +++ b/app/signalr/hub/hub.py @@ -77,9 +77,13 @@ class Client: await asyncio.sleep(settings.signalr_ping_interval) except WebSocketDisconnect: break - except Exception as e: - logger.error(f"Error in ping task for {self.connection_id}: {e}") - break + except RuntimeError as e: + if "disconnect message" in str(e): + break + else: + logger.error(f"Error in ping task for {self.connection_id}: {e}") + except Exception: + logger.exception(f"Error in client {self.connection_id}") class Hub[TState: UserState]: @@ -158,6 +162,8 @@ class Hub[TState: UserState]: return client async def remove_client(self, client: Client) -> None: + if client.connection_token not in self.clients: + return del self.clients[client.connection_token] if client._listen_task: client._listen_task.cancel() @@ -186,7 +192,22 @@ class Hub[TState: UserState]: await method(client) async def send_packet(self, client: Client, packet: Packet) -> None: - await client.send_packet(packet) + try: + await client.send_packet(packet) + except WebSocketDisconnect as e: + logger.info( + f"Client {client.connection_id} disconnected: {e.code}, {e.reason}" + ) + await self.remove_client(client) + except RuntimeError as e: + if "disconnect message" in str(e): + logger.info(f"Client {client.connection_id} closed the connection.") + else: + logger.exception(f"RuntimeError in client {client.connection_id}: {e}") + await self.remove_client(client) + except Exception: + logger.exception(f"Error in client {client.connection_id}") + await self.remove_client(client) async def broadcast_call(self, method: str, *args: Any) -> None: tasks = [] diff --git a/app/signalr/hub/spectator.py b/app/signalr/hub/spectator.py index 0bf3cd4..5485732 100644 --- a/app/signalr/hub/spectator.py +++ b/app/signalr/hub/spectator.py @@ -14,6 +14,7 @@ from app.database.score_token import ScoreToken from app.dependencies.database import engine from app.dependencies.fetcher import get_fetcher from app.dependencies.storage import get_storage_service +from app.exception import InvokeException from app.models.mods import mods_to_int from app.models.score import LegacyReplaySoloScoreInfo, ScoreStatistics from app.models.spectator_hub import ( @@ -30,6 +31,7 @@ from app.utils import unix_timestamp_to_windows from .hub import Client, Hub +from httpx import HTTPError from sqlalchemy.orm import joinedload from sqlmodel import select from sqlmodel.ext.asyncio.session import AsyncSession @@ -190,9 +192,12 @@ class SpectatorHub(Hub[StoreClientState]): fetcher = await get_fetcher() async with AsyncSession(engine) as session: async with session.begin(): - beatmap = await Beatmap.get_or_fetch( - session, fetcher, bid=state.beatmap_id - ) + try: + beatmap = await Beatmap.get_or_fetch( + session, fetcher, bid=state.beatmap_id + ) + except HTTPError: + raise InvokeException(f"Beatmap {state.beatmap_id} not found.") user = ( await session.exec(select(User).where(User.id == user_id)) ).first()