Start y Websocket Server earlier and make it resilient to crashes

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

@@ -10,7 +10,7 @@
 from jupyter_ydoc.ybasedoc import YBaseDoc
 from pycrdt import Doc
 from pycrdt_websocket.ystore import BaseYStore
-from traitlets import Bool, Float, Type
+from traitlets import Bool, Float, Type, Set, Instance
 
 from .handlers import DocSessionHandler, YDocWebSocketHandler
 from .loaders import FileLoaderMapping
@@ -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,13 @@ def initialize_settings(self):
             }
         )
 
+        self.ywebsocket_server = JupyterWebsocketServer(
+            rooms_ready=False,
+            auto_clean_rooms=False,
+            ystore_class=self.ystore_class,
+            log=self.log,
+        )
+
     def initialize_handlers(self):
         self.serverapp.web_app.settings.setdefault(
             "page_config_data",
@@ -103,13 +114,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 +138,33 @@ def initialize_handlers(self):
             ]
         )
 
+    async def start_extension(self):
+        """Start the y-websocket server.
+        """
+        self.log.info("Starting the Collaborative Document Server.")
+
+        def _restart_or_teardown_yserver(_):
+            """If the y websocket server task stopped due to an exception, restart it.
+            
+            If the y-websocket server was cancelled on purpose, tear it down.
+            """
+            self.ywebsocket_server._started = None
+            self.ywebsocket_server._starting = False
+            self.ywebsocket_server._task_group = None
+            # If an exception was raised, restart the websocket server.
+            if self._running_ywebsocket_server.exception():
+                self.log.error(self._running_ywebsocket_server.exception())
+                self.log.warning("Restarting the Collaborative Document Server.")
+                self._running_ywebsocket_server = asyncio.create_task(self.ywebsocket_server.start())
+                self._running_ywebsocket_server.add_done_callback(_restart_or_teardown_yserver)
+                return
+            self.log.info("Stopping the Collaborative Document Server.")
+
+        # Start the websocket server
+        self._running_ywebsocket_server = asyncio.create_task(self.ywebsocket_server.start())
+        # If the websocket crashes for any reason, let's restart it automatically and log errors.
+        self._running_ywebsocket_server.add_done_callback(_restart_or_teardown_yserver)
+
     async def get_document(
         self: YDocExtension,
         *,

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)