diff --git a/projects/jupyter-server-ydoc/jupyter_server_ydoc/app.py b/projects/jupyter-server-ydoc/jupyter_server_ydoc/app.py index 8ee23ee3..2b8b0301 100644 --- a/projects/jupyter-server-ydoc/jupyter_server_ydoc/app.py +++ b/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, ) diff --git a/projects/jupyter-server-ydoc/jupyter_server_ydoc/handlers.py b/projects/jupyter-server-ydoc/jupyter_server_ydoc/handlers.py index d74bd6df..1f976eca 100644 --- a/projects/jupyter-server-ydoc/jupyter_server_ydoc/handlers.py +++ b/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 @@ -81,6 +82,21 @@ async def prepare(self): if self._websocket_server.room_exists(self._room_id): self.room: YRoom = await self._websocket_server.get_room(self._room_id) else: + # Logging exceptions, instead of raising them here to ensure + # that the y-rooms stay alive even after an exception is seen. + 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. + """ + log.error( + f"Document Room Exception, (room_id={self._room_id or 'unknown'}): ", + exc_info=exception, + ) + return True + if self._room_id.count(":") >= 2: # DocumentRoom file_format, file_type, file_id = decode_file_path(self._room_id) @@ -102,13 +118,18 @@ async def prepare(self): self.event_logger, ystore, self.log, - self._document_save_delay, + exception_handler=exception_logger, + save_delay=self._document_save_delay, ) else: # TransientRoom # it is a transient document (e.g. awareness) - self.room = TransientRoom(self._room_id, self.log) + self.room = TransientRoom( + self._room_id, + log=self.log, + exception_handler=exception_logger, + ) await self._websocket_server.start_room(self.room) self._websocket_server.add_room(self._room_id, self.room) @@ -205,7 +226,8 @@ async def open(self, room_id): else: 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,7 +294,11 @@ 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: @@ -387,7 +413,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 +432,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) diff --git a/projects/jupyter-server-ydoc/jupyter_server_ydoc/rooms.py b/projects/jupyter-server-ydoc/jupyter_server_ydoc/rooms.py index 691943c5..00787dbe 100644 --- a/projects/jupyter-server-ydoc/jupyter_server_ydoc/rooms.py +++ b/projects/jupyter-server-ydoc/jupyter_server_ydoc/rooms.py @@ -5,7 +5,7 @@ import asyncio from logging import Logger -from typing import Any +from typing import Any, Callable from jupyter_events import EventLogger from jupyter_ydoc import ydocs as YDOCS @@ -31,8 +31,9 @@ def __init__( ystore: BaseYStore | None, log: Logger | None, save_delay: float | None = None, + exception_handler: Callable[[Exception, Logger], bool] | None = None, ): - super().__init__(ready=False, ystore=ystore, log=log) + super().__init__(ready=False, ystore=ystore, exception_handler=exception_handler, log=log) self._room_id: str = room_id self._file_format: str = file_format @@ -123,7 +124,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 +138,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"] @@ -281,8 +286,13 @@ async def _maybe_save_document(self, saving_document: asyncio.Task | None) -> No class TransientRoom(YRoom): """A Y room for sharing state (e.g. awareness).""" - def __init__(self, room_id: str, log: Logger | None): - super().__init__(log=log) + def __init__( + self, + room_id: str, + log: Logger | None = None, + exception_handler: Callable[[Exception, Logger], bool] | None = None, + ): + super().__init__(log=log, exception_handler=exception_handler) self._room_id = room_id diff --git a/projects/jupyter-server-ydoc/jupyter_server_ydoc/websocketserver.py b/projects/jupyter-server-ydoc/jupyter_server_ydoc/websocketserver.py index 365da078..3137e1f8 100644 --- a/projects/jupyter-server-ydoc/jupyter_server_ydoc/websocketserver.py +++ b/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. + + This 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] = {} diff --git a/projects/jupyter-server-ydoc/pyproject.toml b/projects/jupyter-server-ydoc/pyproject.toml index 9ab08633..d96af6cf 100644 --- a/projects/jupyter-server-ydoc/pyproject.toml +++ b/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"