Websocket fix (#592)

* Encased 'WebSocket.send()' and 'WebSocket.close()' in try-except blocks
* Added regular pinging when executing 'WebSocketApp.run_forever()'

Author: tieneupin

src/murfey/client/__init__.py

@@ -11,8 +11,9 @@
 import webbrowser
 from datetime import datetime
 from pathlib import Path
+from pprint import pprint
 from queue import Queue
-from typing import List, Literal
+from typing import Literal
 from urllib.parse import ParseResult, urlparse
 
 import requests
@@ -40,7 +41,7 @@ def write_config(config: configparser.ConfigParser):
 
 
 def main_loop(
-    source_watchers: List[murfey.client.watchdir.DirWatcher],
+    source_watchers: list[murfey.client.watchdir.DirWatcher],
     appearance_time: float,
     transfer_all: bool,
 ):
@@ -91,6 +92,7 @@ def _check_for_updates(
 
 
 def run():
+    # Load client config and server information
     config = read_config()
     instrument_name = config["Murfey"]["instrument_name"]
     try:
@@ -109,6 +111,7 @@ def run():
     else:
         known_server = config["Murfey"].get("server")
 
+    # Set up argument parser with dynamic defaults based on client config
     parser = argparse.ArgumentParser(description="Start the Murfey client")
     parser.add_argument(
         "--server",
@@ -194,23 +197,23 @@ def run():
         default=False,
         help="Do not trigger processing for any data directories currently on disk (you may have started processing for them in a previous murfey run)",
     )
-
     args = parser.parse_args()
 
+    # Logic to exit early based on parsed args
     if not args.server:
         exit("Murfey server not set. Please run with --server host:port")
     if not args.server.startswith(("http://", "https://")):
         if "://" in args.server:
             exit("Unknown server protocol. Only http:// and https:// are allowed")
         args.server = f"http://{args.server}"
-
     if args.remove_files:
         remove_prompt = Confirm.ask(
             f"Are you sure you want to remove files from {args.source or Path('.').absolute()}?"
         )
         if not remove_prompt:
             exit("Exiting")
 
+    # If a new server URL is provided, save info to config file
     murfey_url = urlparse(args.server, allow_fragments=False)
     if args.server != known_server:
         # New server specified. Verify that it is real
@@ -232,8 +235,7 @@ def run():
     if args.no_transfer:
         log.info("No files will be transferred as --no-transfer flag was specified")
 
-    from pprint import pprint
-
+    # Check ISPyB (if set up) for ongoing visits
     ongoing_visits = []
     if args.visit:
         ongoing_visits = [args.visit]
@@ -250,35 +252,38 @@ def run():
 
     _enable_webbrowser_in_cygwin()
 
+    # Set up additional log handlers
     log.setLevel(logging.DEBUG)
     log_queue = Queue()
     input_queue = Queue()
 
-    # rich_handler = DirectableRichHandler(log_queue, enable_link_path=False)
+    # Rich-based console handler
     rich_handler = DirectableRichHandler(enable_link_path=False)
     rich_handler.setLevel(logging.DEBUG if args.debug else logging.INFO)
 
+    # Set up websocket app and handler
     client_id = requests.get(f"{murfey_url.geturl()}/new_client_id/").json()
     ws = murfey.client.websocket.WSApp(
         server=args.server,
         id=client_id["new_id"],
     )
+    ws_handler = CustomHandler(ws.send)
 
+    # Add additional handlers and set logging levels
     logging.getLogger().addHandler(rich_handler)
-    handler = CustomHandler(ws.send)
-    logging.getLogger().addHandler(handler)
+    logging.getLogger().addHandler(ws_handler)
     logging.getLogger("murfey").setLevel(logging.INFO)
     logging.getLogger("websocket").setLevel(logging.WARNING)
 
     log.info("Starting Websocket connection")
 
-    status_bar = StatusBar()
-
+    # Load machine data for subsequent sections
     machine_data = requests.get(
         f"{murfey_url.geturl()}/instruments/{instrument_name}/machine"
     ).json()
     gain_ref: Path | None = None
 
+    # Set up Murfey environment instance and map it to websocket app
     instance_environment = MurfeyInstanceEnvironment(
         url=murfey_url,
         client_id=ws.id,
@@ -295,9 +300,10 @@ def run():
             else ""
         ),
     )
-
     ws.environment = instance_environment
 
+    # Set up and run Murfey TUI app
+    status_bar = StatusBar()
     rich_handler.redirect = True
     app = MurfeyTUI(
         environment=instance_environment,

src/murfey/client/websocket.py

@@ -87,7 +87,7 @@ def _run_websocket_event_loop(self):
         backoff = 0
         while True:
             attempt_start = time.perf_counter()
-            connection_failure = self._ws.run_forever()
+            connection_failure = self._ws.run_forever(ping_interval=30, ping_timeout=10)
             if not connection_failure:
                 break
             if (time.perf_counter() - attempt_start) < 5:
@@ -108,7 +108,10 @@ def _send_queue_feeder(self):
                 continue
             while not self._ready:
                 time.sleep(0.3)
-            self._ws.send(element)
+            try:
+                self._ws.send(element)
+            except Exception:
+                log.error("Error sending message through websocket", exc_info=True)
             self._send_queue.task_done()
         log.debug("Websocket send-queue-feeder thread stopped")
 
@@ -135,7 +138,10 @@ def close(self):
             self._send_queue.put(None)
             self._feeder_thread.join()
             self._receiver_thread.join()
-        self._ws.close()
+        try:
+            self._ws.close()
+        except Exception:
+            log.error("Error closing websocket connection", exc_info=True)
 
     def on_message(self, ws: websocket.WebSocketApp, message: str):
         self._receive_queue.put(message)

src/murfey/server/websocket.py

@@ -4,10 +4,10 @@
 import json
 import logging
 from datetime import datetime
-from typing import Any, Dict, TypeVar, Union
+from typing import Any, TypeVar, Union
 
 from fastapi import APIRouter, WebSocket, WebSocketDisconnect
-from sqlmodel import select
+from sqlmodel import Session, select
 
 import murfey.server.prometheus as prom
 from murfey.server.murfey_db import get_murfey_db_session
@@ -22,10 +22,13 @@
 
 class ConnectionManager:
     def __init__(self):
-        self.active_connections: Dict[int | str, WebSocket] = {}
+        self.active_connections: dict[int | str, WebSocket] = {}
 
     async def connect(
-        self, websocket: WebSocket, client_id: int | str, register_client: bool = True
+        self,
+        websocket: WebSocket,
+        client_id: Union[int, str],
+        register_client: bool = True,
     ):
         await websocket.accept()
         self.active_connections[client_id] = websocket
@@ -38,16 +41,17 @@ async def connect(
 
     @staticmethod
     def _register_new_client(client_id: int):
+        log.debug(f"Registering new client with ID {client_id}")
         new_client = ClientEnvironment(client_id=client_id, connected=True)
-        murfey_db = next(get_murfey_db_session())
+        murfey_db: Session = next(get_murfey_db_session())
         murfey_db.add(new_client)
         murfey_db.commit()
         murfey_db.close()
 
-    def disconnect(self, client_id: int | str, unregister_client: bool = True):
+    def disconnect(self, client_id: Union[int, str], unregister_client: bool = True):
         self.active_connections.pop(client_id)
         if unregister_client:
-            murfey_db = next(get_murfey_db_session())
+            murfey_db: Session = next(get_murfey_db_session())
             client_env = murfey_db.exec(
                 select(ClientEnvironment).where(
                     ClientEnvironment.client_id == client_id
@@ -73,7 +77,7 @@ async def websocket_endpoint(websocket: WebSocket, client_id: int):
         while True:
             data = await websocket.receive_text()
             try:
-                json_data = json.loads(data)
+                json_data: dict = json.loads(data)
                 if json_data["type"] == "log":  # and isinstance(json_data, dict)
                     json_data.pop("type")
                     await forward_log(json_data, websocket)
@@ -92,15 +96,16 @@ async def websocket_endpoint(websocket: WebSocket, client_id: int):
 
 @ws.websocket("/connect/{client_id}")
 async def websocket_connection_endpoint(
-    websocket: WebSocket, client_id: Union[int, str]
+    websocket: WebSocket,
+    client_id: Union[int, str],
 ):
     await manager.connect(websocket, client_id, register_client=False)
     await manager.broadcast(f"Client {client_id} joined")
     try:
         while True:
             data = await websocket.receive_text()
             try:
-                json_data = json.loads(data)
+                json_data: dict = json.loads(data)
                 if json_data.get("type") == "log":  # and isinstance(json_data, dict)
                     json_data.pop("type")
                     await forward_log(json_data, websocket)
@@ -115,12 +120,12 @@ async def websocket_connection_endpoint(
         await manager.broadcast(f"Client #{client_id} disconnected")
 
 
-async def check_connections(active_connections):
+async def check_connections(active_connections: list[WebSocket]):
     log.info("Checking connections")
     for connection in active_connections:
         log.info("Checking response")
         try:
-            await asyncio.wait_for(connection.receive(), timeout=6)
+            await asyncio.wait_for(connection.receive(), timeout=10)
         except asyncio.TimeoutError:
             log.info(f"Disconnecting Client {connection[0]}")
             manager.disconnect(connection[0], connection[1])
@@ -139,7 +144,7 @@ async def forward_log(logrecord: dict[str, Any], websocket: WebSocket):
 
 @ws.delete("/test/{client_id}")
 async def close_ws_connection(client_id: int):
-    murfey_db = next(get_murfey_db_session())
+    murfey_db: Session = next(get_murfey_db_session())
     client_env = murfey_db.exec(
         select(ClientEnvironment).where(ClientEnvironment.client_id == client_id)
     ).one()