Start y Websocket Server earlier and make it resilient to crashes

jupyterlab · Apr 26, 2024 · a893b45 · a893b45
1 parent adcde32
commit a893b45
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 12 deletions.
diff --git a/projects/jupyter-server-ydoc/jupyter_server_ydoc/app.py b/projects/jupyter-server-ydoc/jupyter_server_ydoc/app.py
@@ -10,7 +10,8 @@
 from jupyter_ydoc.ybasedoc import YBaseDoc
 from pycrdt import Doc
 from pycrdt_websocket.ystore import BaseYStore
-from traitlets import Bool, Float, Type
+from pycrdt_websocket.websocket_server import exception_logger
+from traitlets import Bool, Float, Type, Set, Instance
 
 from .handlers import DocSessionHandler, YDocWebSocketHandler
 from .loaders import FileLoaderMapping
@@ -24,7 +25,6 @@
 )
 from .websocketserver import JupyterWebsocketServer, RoomNotFound
 
-
 class YDocExtension(ExtensionApp):
     name = "jupyter_server_ydoc"
     app_name = "Collaboration"
@@ -75,6 +75,10 @@ class YDocExtension(ExtensionApp):
         model.""",
     )
 
+    _running_ywebsocket_server = Instance(asyncio.Task, allow_none=True)
+    ywebsocket_server = Instance(JupyterWebsocketServer, allow_none=True)
+
+
     def initialize(self):
         super().initialize()
         self.serverapp.event_logger.register_event_schema(EVENTS_SCHEMA_PATH)
@@ -90,6 +94,14 @@ def initialize_settings(self):
             }
         )
 
+        self.ywebsocket_server = JupyterWebsocketServer(
+            rooms_ready=False,
+            auto_clean_rooms=False,
+            ystore_class=self.ystore_class,
+            log=self.log,
+            exception_logger=exception_logger
+        )
+
     def initialize_handlers(self):
         self.serverapp.web_app.settings.setdefault(
             "page_config_data",
@@ -103,13 +115,6 @@ def initialize_handlers(self):
         for k, v in self.config.get(self.ystore_class.__name__, {}).items():
             setattr(self.ystore_class, k, v)
 
-        self.ywebsocket_server = JupyterWebsocketServer(
-            rooms_ready=False,
-            auto_clean_rooms=False,
-            ystore_class=self.ystore_class,
-            log=self.log,
-        )
-
         # self.settings is local to the ExtensionApp but here we need
         # the global app settings in which the file id manager will register
         # itself maybe at a later time.
@@ -134,6 +139,14 @@ def initialize_handlers(self):
             ]
         )
 
+    async def start_extension(self):
+        """Start the y-websocket server.
+        """
+        self.log.info("Starting the Collaborative Document Server.")
+
+        # Start the websocket server
+        self._running_ywebsocket_server = asyncio.create_task(self.ywebsocket_server.start())
+
     async def get_document(
         self: YDocExtension,
         *,

diff --git a/projects/jupyter-server-ydoc/jupyter_server_ydoc/handlers.py b/projects/jupyter-server-ydoc/jupyter_server_ydoc/handlers.py
@@ -70,9 +70,7 @@ def create_task(self, aw):
         task.add_done_callback(self._background_tasks.discard)
 
     async def prepare(self):
-        if not self._websocket_server.started.is_set():
-            self.create_task(self._websocket_server.start())
-            await self._websocket_server.started.wait()
+        await self._websocket_server.started.wait()
 
         # Get room
         self._room_id: str = room_id_from_encoded_path(self.request.path)