diff --git a/nicegui/client.py b/nicegui/client.py index 9b0135cc4f..578deb3653 100644 --- a/nicegui/client.py +++ b/nicegui/client.py @@ -56,6 +56,9 @@ class Client: instances: ClassVar[dict[str, Client]] = {} '''Maps client IDs to clients.''' + sid_to_client: ClassVar[dict[str, Client]] = {} + '''Maps socket IDs to clients.''' + shared_head_html = '' '''HTML to be inserted in the of every page template.''' @@ -85,6 +88,9 @@ def __init__(self, page: page, *, request: Request | None = None) -> None: self.page = page self.outbox = Outbox(self) + self._last_ping = 0.0 + self._latency = 0.0 + if self._request is not None: self._request.scope['nicegui_page_path'] = self.page.path @@ -290,6 +296,7 @@ def handle_handshake(self, socket_id: str, document_id: str, next_message_id: in self._waiting_for_connection.clear() self._connected.set() self._socket_to_document_id[socket_id] = document_id + self.sid_to_client[socket_id] = self self._cancel_delete_task(document_id) self._num_connections[document_id] += 1 if next_message_id is not None: @@ -314,6 +321,7 @@ def handle_disconnect(self, socket_id: str) -> None: self.safe_invoke(t) for t in core.app._disconnect_handlers: # pylint: disable=protected-access self.safe_invoke(t) + self.sid_to_client.pop(socket_id, None) async def delete_content() -> None: await asyncio.sleep(self.page.resolve_reconnect_timeout()) diff --git a/nicegui/nicegui.py b/nicegui/nicegui.py index 0ea74ebb2c..a6be8ce44e 100644 --- a/nicegui/nicegui.py +++ b/nicegui/nicegui.py @@ -8,6 +8,10 @@ from typing import Any import socketio +from engineio.async_server import AsyncServer # type: ignore[import-untyped] +from engineio.async_socket import AsyncSocket # type: ignore[import-untyped] +from engineio.base_server import BaseServer # type: ignore[import-untyped] +from engineio.packet import PING, PONG, Packet # type: ignore[import-untyped] from fastapi import HTTPException, Request from fastapi.responses import FileResponse, Response @@ -25,6 +29,43 @@ from .staticfiles import CacheControlledStaticFiles from .version import __version__ +BaseServer.event_names.append('ping') +BaseServer.event_names.append('pong') + + +_original_asyncsocket_init = AsyncSocket.__init__ + + +def _patched_asyncsocket_init(self: AsyncSocket, *args: Any, **kwargs: Any) -> None: + _original_asyncsocket_init(self, *args, **kwargs) + + original_send = self.send + + async def patched_send(pkt: Packet) -> None: + try: + if pkt.packet_type == PING: # PING + assert isinstance(self.server, AsyncServer) + await self.server._trigger_event('ping', self.sid, pkt.packet_type, run_async=self.server.async_handlers) # pylint: disable=protected-access + except Exception: + helpers.warn_once('Unable to patch ping handling for AsyncSocket.') + return await original_send(pkt) + self.send = patched_send + + original_receive = self.receive + + async def patched_receive(pkt: Packet) -> None: + try: + if pkt.packet_type == PONG: # PONG + assert isinstance(self.server, AsyncServer) + await self.server._trigger_event('pong', self.sid, pkt.packet_type, run_async=self.server.async_handlers) # pylint: disable=protected-access + except Exception: + helpers.warn_once('Unable to patch pong handling for AsyncSocket.') + return await original_receive(pkt) + self.receive = patched_receive + + +AsyncSocket.__init__ = _patched_asyncsocket_init + @asynccontextmanager async def _lifespan(_: App): @@ -250,6 +291,20 @@ def _on_log(_: str, msg: dict) -> None: }[msg['level']](msg['message']) +@sio.eio.on('ping') +async def _on_ping(sid: str, _: Any) -> None: + client = Client.sid_to_client.get(sio.manager.sid_from_eio_sid(sid, '/')) + if client: + client._last_ping = time.time() # pylint: disable=protected-access + + +@sio.eio.on('pong') +async def _on_pong(sid: str, _: Any) -> None: + client = Client.sid_to_client.get(sio.manager.sid_from_eio_sid(sid, '/')) + if client: + client._latency = time.time() - client._last_ping # pylint: disable=protected-access + + async def prune_tab_storage(*, force: bool = False) -> None: """Prune tab storage that is older than the configured ``max_tab_storage_age``.""" tab_storages = core.app.storage._tabs # pylint: disable=protected-access