Log (instead of raise) exceptions when running as a server extension

Author: Zsailer

projects/jupyter-server-ydoc/jupyter_server_ydoc/app.py

@@ -22,7 +22,7 @@
     encode_file_path,
     room_id_from_encoded_path,
 )
-from .websocketserver import JupyterWebsocketServer, RoomNotFound
+from .websocketserver import JupyterWebsocketServer, RoomNotFound, exception_logger
 
 
 class YDocExtension(ExtensionApp):
@@ -107,6 +107,9 @@ def initialize_handlers(self):
             rooms_ready=False,
             auto_clean_rooms=False,
             ystore_class=self.ystore_class,
+            # Log exceptions, because we don't want the websocket server
+            # to _ever_ crash permanently in a live jupyter_server.
+            exception_handler=exception_logger,
             log=self.log,
         )
 

projects/jupyter-server-ydoc/jupyter_server_ydoc/handlers.py

@@ -7,6 +7,7 @@
 import json
 import time
 import uuid
+from logging import Logger
 from typing import Any
 
 from jupyter_server.auth import authorized
@@ -105,6 +106,26 @@ async def prepare(self):
                         self._document_save_delay,
                     )
 
+                    def exception_logger(exception: Exception, log: Logger) -> bool:
+                        """A function that catches any exceptions raised in the websocket
+                        server and logs them.
+
+                        The protects the y-room's task group from cancelling
+                        anytime an exception is raised.
+                        """
+                        room_id = "unknown"
+                        if self.room.room_id:
+                            room_id = self.room.room_id
+                        log.error(
+                            f"Document Room Exception, (room_id={room_id}: ",
+                            exc_info=exception,
+                        )
+                        return True
+
+                    # Logging exceptions, instead of raising them here to ensure
+                    # that the y-rooms stay alive even after an exception is seen.
+                    self.room.exception_handler = exception_logger
+
                 else:
                     # TransientRoom
                     # it is a transient document (e.g. awareness)
@@ -203,9 +224,12 @@ async def open(self, room_id):
                     self.log.error(f"File {file.path} not found.\n{e!r}", exc_info=e)
                     self.close(1004, f"File {file.path} not found.")
                 else:
-                    self.log.error(f"Error initializing: {file.path}\n{e!r}", exc_info=e)
+                    self.log.error(
+                        f"Error initializing: {file.path}\n{e!r}", exc_info=e
+                    )
                     self.close(
-                        1003, f"Error initializing: {file.path}. You need to close the document."
+                        1003,
+                        f"Error initializing: {file.path}. You need to close the document.",
                     )
 
                 # Clean up the room and delete the file loader
@@ -272,16 +296,24 @@ async def on_message(self, message):
 
             user = self.current_user
             data = json.dumps(
-                {"sender": user.username, "timestamp": time.time(), "content": json.loads(msg)}
+                {
+                    "sender": user.username,
+                    "timestamp": time.time(),
+                    "content": json.loads(msg),
+                }
             ).encode("utf8")
 
             for client in self.room.clients:
                 if client != self:
                     task = asyncio.create_task(
-                        client.send(bytes([MessageType.CHAT]) + write_var_uint(len(data)) + data)
+                        client.send(
+                            bytes([MessageType.CHAT]) + write_var_uint(len(data)) + data
+                        )
                     )
                     self._websocket_server.background_tasks.add(task)
-                    task.add_done_callback(self._websocket_server.background_tasks.discard)
+                    task.add_done_callback(
+                        self._websocket_server.background_tasks.discard
+                    )
 
         self._message_queue.put_nowait(message)
         self._websocket_server.ypatch_nb += 1
@@ -300,7 +332,9 @@ def on_close(self) -> None:
         if self._room_id != "JupyterLab:globalAwareness":
             self._emit_awareness_event(self.current_user.username, "leave")
 
-    def _emit(self, level: LogLevel, action: str | None = None, msg: str | None = None) -> None:
+    def _emit(
+        self, level: LogLevel, action: str | None = None, msg: str | None = None
+    ) -> None:
         _, _, file_id = decode_file_path(self._room_id)
         path = self._file_id_manager.get_path(file_id)
 
@@ -312,12 +346,16 @@ def _emit(self, level: LogLevel, action: str | None = None, msg: str | None = No
 
         self.event_logger.emit(schema_id=JUPYTER_COLLABORATION_EVENTS_URI, data=data)
 
-    def _emit_awareness_event(self, username: str, action: str, msg: str | None = None) -> None:
+    def _emit_awareness_event(
+        self, username: str, action: str, msg: str | None = None
+    ) -> None:
         data = {"roomid": self._room_id, "username": username, "action": action}
         if msg:
             data["msg"] = msg
 
-        self.event_logger.emit(schema_id=JUPYTER_COLLABORATION_AWARENESS_EVENTS_URI, data=data)
+        self.event_logger.emit(
+            schema_id=JUPYTER_COLLABORATION_AWARENESS_EVENTS_URI, data=data
+        )
 
     async def _clean_room(self) -> None:
         """
@@ -387,7 +425,12 @@ async def put(self, path):
             # index already exists
             self.log.info("Request for Y document '%s' with room ID: %s", path, idx)
             data = json.dumps(
-                {"format": format, "type": content_type, "fileId": idx, "sessionId": SERVER_SESSION}
+                {
+                    "format": format,
+                    "type": content_type,
+                    "fileId": idx,
+                    "sessionId": SERVER_SESSION,
+                }
             )
             self.set_status(200)
             return self.finish(data)
@@ -401,7 +444,12 @@ async def put(self, path):
         # index successfully created
         self.log.info("Request for Y document '%s' with room ID: %s", path, idx)
         data = json.dumps(
-            {"format": format, "type": content_type, "fileId": idx, "sessionId": SERVER_SESSION}
+            {
+                "format": format,
+                "type": content_type,
+                "fileId": idx,
+                "sessionId": SERVER_SESSION,
+            }
         )
         self.set_status(201)
         return self.finish(data)

projects/jupyter-server-ydoc/jupyter_server_ydoc/rooms.py

@@ -123,7 +123,9 @@ async def initialize(self) -> None:
                 if self._document.source != model["content"]:
                     # TODO: Delete document from the store.
                     self._emit(
-                        LogLevel.INFO, "initialize", "The file is out-of-sync with the ystore."
+                        LogLevel.INFO,
+                        "initialize",
+                        "The file is out-of-sync with the ystore.",
                     )
                     self.log.info(
                         "Content in file %s is out-of-sync with the ystore %s",
@@ -135,7 +137,9 @@ async def initialize(self) -> None:
             if read_from_source:
                 self._emit(LogLevel.INFO, "load", "Content loaded from disk.")
                 self.log.info(
-                    "Content in room %s loaded from file %s", self._room_id, self._file.path
+                    "Content in room %s loaded from file %s",
+                    self._room_id,
+                    self._file.path,
                 )
                 self._document.source = model["content"]
 
@@ -146,7 +150,9 @@ async def initialize(self) -> None:
             self.ready = True
             self._emit(LogLevel.INFO, "initialize", "Room initialized")
 
-    def _emit(self, level: LogLevel, action: str | None = None, msg: str | None = None) -> None:
+    def _emit(
+        self, level: LogLevel, action: str | None = None, msg: str | None = None
+    ) -> None:
         data = {"level": level.value, "room": self._room_id, "path": self._file.path}
         if action:
             data["action"] = action
@@ -180,8 +186,12 @@ async def _on_outofband_change(self) -> None:
         """
         Called when the file got out-of-band changes.
         """
-        self.log.info("Out-of-band changes. Overwriting the content in room %s", self._room_id)
-        self._emit(LogLevel.INFO, "overwrite", "Out-of-band changes. Overwriting the room.")
+        self.log.info(
+            "Out-of-band changes. Overwriting the content in room %s", self._room_id
+        )
+        self._emit(
+            LogLevel.INFO, "overwrite", "Out-of-band changes. Overwriting the room."
+        )
 
         try:
             model = await self._file.load_content(self._file_format, self._file_type)
@@ -257,9 +267,13 @@ async def _maybe_save_document(self, saving_document: asyncio.Task | None) -> No
             return
 
         except OutOfBandChanges:
-            self.log.info("Out-of-band changes. Overwriting the content in room %s", self._room_id)
+            self.log.info(
+                "Out-of-band changes. Overwriting the content in room %s", self._room_id
+            )
             try:
-                model = await self._file.load_content(self._file_format, self._file_type)
+                model = await self._file.load_content(
+                    self._file_format, self._file_type
+                )
             except Exception as e:
                 msg = f"Error loading content from file: {self._file.path}\n{e!r}"
                 self.log.error(msg, exc_info=e)

projects/jupyter-server-ydoc/jupyter_server_ydoc/websocketserver.py

@@ -5,7 +5,7 @@
 
 import asyncio
 from logging import Logger
-from typing import Any
+from typing import Any, Callable
 
 from pycrdt_websocket.websocket_server import WebsocketServer, YRoom
 from pycrdt_websocket.ystore import BaseYStore
@@ -16,6 +16,17 @@ class RoomNotFound(LookupError):
     pass
 
 
+def exception_logger(exception: Exception, log: Logger) -> bool:
+    """A function that catches any exceptions raised in the websocket
+    server and logs them.
+
+    The protects the websocket server's task group from cancelling
+    anytime an exception is raised.
+    """
+    log.error("Jupyter Websocket Server: ", exc_info=exception)
+    return True
+
+
 class JupyterWebsocketServer(WebsocketServer):
     """Ypy websocket server.
 
@@ -30,9 +41,15 @@ def __init__(
         ystore_class: BaseYStore,
         rooms_ready: bool = True,
         auto_clean_rooms: bool = True,
+        exception_handler: Callable[[Exception, Logger], bool] | None = None,
         log: Logger | None = None,
     ):
-        super().__init__(rooms_ready, auto_clean_rooms, log)
+        super().__init__(
+            rooms_ready=rooms_ready,
+            auto_clean_rooms=auto_clean_rooms,
+            exception_handler=exception_handler,
+            log=log,
+        )
         self.ystore_class = ystore_class
         self.ypatch_nb = 0
         self.connected_users: dict[Any, Any] = {}

projects/jupyter-server-ydoc/pyproject.toml

@@ -30,7 +30,7 @@ authors = [
 dependencies = [
     "jupyter_server>=2.4.0,<3.0.0",
     "jupyter_ydoc>=2.0.0,<3.0.0",
-    "pycrdt-websocket>=0.13.0,<0.14.0",
+    "pycrdt-websocket>=0.13.1,<0.14.0",
     "jupyter_events>=0.10.0",
     "jupyter_server_fileid>=0.7.0,<1",
     "jsonschema>=4.18.0"