handle exception when websocket server start room failed

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

@@ -11,6 +11,7 @@
 
 from jupyter_server.auth import authorized
 from jupyter_server.base.handlers import APIHandler, JupyterHandler
+from jupyter_server.utils import ensure_async
 from jupyter_ydoc import ydocs as YDOCS
 from pycrdt_websocket.websocket_server import YRoom
 from pycrdt_websocket.ystore import BaseYStore
@@ -70,6 +71,7 @@ def create_task(self, aw):
         task.add_done_callback(self._background_tasks.discard)
 
     async def prepare(self):
+        await ensure_async(super().prepare())
         if not self._websocket_server.started.is_set():
             self.create_task(self._websocket_server.start())
             await self._websocket_server.started.wait()
@@ -110,12 +112,27 @@ async def prepare(self):
                     # it is a transient document (e.g. awareness)
                     self.room = TransientRoom(self._room_id, self.log)
 
-            await self._websocket_server.start_room(self.room)
-            self._websocket_server.add_room(self._room_id, self.room)
-
-        res = super().prepare()
-        if res is not None:
-            return await res
+                try:
+                    await self._websocket_server.start_room(self.room)
+                except Exception as e:
+                    self.log.error("Room %s failed to start on websocket server", self._room_id)
+                    # Clean room
+                    self.room.stop()
+                    self.log.info("Room %s deleted", self._room_id)
+                    self._emit(LogLevel.INFO, "clean", "Room deleted.")
+
+                    # if websocket server failed, then room and file will be garbage collected and file
+                    # Clean the file loader in file loader mapping if there are not rooms using it
+                    _, _, file_id = decode_file_path(self._room_id)
+                    file = self._file_loaders[file_id]
+                    if file.number_of_subscriptions == 0 or (
+                        file.number_of_subscriptions == 1 and self._room_id in file._subscriptions
+                    ):
+                        self.log.info("Deleting file %s", file.path)
+                        await self._file_loaders.remove(file_id)
+                        self._emit(LogLevel.INFO, "clean", "file loader removed.")
+                    raise e
+                self._websocket_server.add_room(self._room_id, self.room)
 
     def initialize(
         self,
@@ -227,7 +244,7 @@ async def send(self, message):
         try:
             self.write_message(message, binary=True)
         except Exception as e:
-            self.log.debug("Failed to write message", exc_info=e)
+            self.log.error("Failed to write message", exc_info=e)
 
     async def recv(self):
         """
@@ -292,13 +309,15 @@ def on_close(self) -> None:
         """
         # stop serving this client
         self._message_queue.put_nowait(b"")
-        if isinstance(self.room, DocumentRoom) and self.room.clients == [self]:
-            # no client in this room after we disconnect
-            # keep the document for a while in case someone reconnects
-            self.log.info("Cleaning room: %s", self._room_id)
-            self.room.cleaner = asyncio.create_task(self._clean_room())
-        if self._room_id != "JupyterLab:globalAwareness":
-            self._emit_awareness_event(self.current_user.username, "leave")
+        if hasattr(self, "room"):
+            if isinstance(self.room, DocumentRoom) and self.room.clients == [self]:
+                # no client in this room after we disconnect
+                # keep the document for a while in case someone reconnects
+                self.log.info("Cleaning room: %s", self._room_id)
+                self.room.cleaner = asyncio.create_task(self._clean_room())
+        if hasattr(self, "_room_id"):
+            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:
         _, _, file_id = decode_file_path(self._room_id)

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

@@ -74,7 +74,10 @@ async def clean(self) -> None:
         if self._watcher is not None:
             if not self._watcher.cancelled():
                 self._watcher.cancel()
-            await self._watcher
+            try:
+                await self._watcher
+            except asyncio.CancelledError:
+                self._log.info(f"file watcher for '{self.file_id}' is cancelled now")
 
     def observe(self, id: str, callback: Callable[[], Coroutine[Any, Any, None]]) -> None:
         """

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

@@ -161,7 +161,10 @@ def stop(self) -> None:
 
         Cancels the save task and unsubscribes from the file.
         """
-        super().stop()
+        try:
+            super().stop()
+        except RuntimeError:
+            pass
         # TODO: Should we cancel or wait ?
         if self._saving_document:
             self._saving_document.cancel()
@@ -299,3 +302,12 @@ async def _broadcast_updates(self):
             await super()._broadcast_updates()
         except asyncio.CancelledError:
             pass
+
+    def stop(self) -> None:
+        """
+        Stop the room.
+        """
+        try:
+            super().stop()
+        except RuntimeError:
+            pass

tests/test_handlers.py

@@ -131,3 +131,87 @@ async def my_listener(logger: EventLogger, schema_id: str, data: dict) -> None:
     assert collected_data[1]["action"] == "leave"
     assert collected_data[1]["roomid"] == "text:file:" + fim.get_id("test.txt")
     assert collected_data[1]["username"] is not None
+
+
+async def test_room_handler_doc_client_should_cleanup_room_file(
+    rtc_create_file, rtc_connect_doc_client, jp_serverapp
+):
+    path, _ = await rtc_create_file("test.txt", "test")
+
+    event = Event()
+
+    def _on_document_change(target: str, e: Any) -> None:
+        if target == "source":
+            event.set()
+
+    doc = YUnicode()
+    doc.observe(_on_document_change)
+
+    async with await rtc_connect_doc_client("text", "file", path) as ws, WebsocketProvider(
+        doc.ydoc, ws
+    ):
+        await event.wait()
+        await sleep(0.1)
+
+    # kill websocketserver to minik task group inactive failure
+    jp_serverapp.web_app.settings["jupyter_server_ydoc"].ywebsocket_server.stop()
+
+    listener_was_called = False
+    collected_data = []
+
+    async def my_listener(logger: EventLogger, schema_id: str, data: dict) -> None:
+        nonlocal listener_was_called
+        collected_data.append(data)
+        listener_was_called = True
+
+    event_logger = jp_serverapp.event_logger
+    event_logger.add_listener(
+        schema_id="https://schema.jupyter.org/jupyter_collaboration/session/v1",
+        listener=my_listener,
+    )
+
+    path2, _ = await rtc_create_file("test2.txt", "test2")
+
+    try:
+        async with await rtc_connect_doc_client("text2", "file2", path2) as ws, WebsocketProvider(
+            doc.ydoc, ws
+        ):
+            await event.wait()
+            await sleep(0.1)
+    except Exception:
+        pass
+
+    try:
+        async with await rtc_connect_doc_client("text2", "file2", path2) as ws, WebsocketProvider(
+            doc.ydoc, ws
+        ):
+            await event.wait()
+            await sleep(0.1)
+    except Exception:
+        pass
+
+    fim = jp_serverapp.web_app.settings["file_id_manager"]
+
+    assert listener_was_called is True
+    assert len(collected_data) == 4
+    # no two collaboration events are emitted.
+    # [{'level': 'WARNING', 'msg': 'There is another collaborative session accessing the same file.\nThe synchronization bet...ou might lose some of your changes.', 'path': 'test2.txt', 'room': 'text2:file2:51b7e24f-f534-47fb-882f-5cc45ba867fe'}]
+    assert collected_data[0]["path"] == "test2.txt"
+    assert collected_data[0]["room"] == "text2:file2:" + fim.get_id("test2.txt")
+    assert collected_data[0]["action"] == "clean"
+    assert collected_data[0]["msg"] == "Room deleted."
+    assert collected_data[1]["path"] == "test2.txt"
+    assert collected_data[1]["room"] == "text2:file2:" + fim.get_id("test2.txt")
+    assert collected_data[1]["action"] == "clean"
+    assert collected_data[1]["msg"] == "file loader removed."
+    assert collected_data[2]["path"] == "test2.txt"
+    assert collected_data[2]["room"] == "text2:file2:" + fim.get_id("test2.txt")
+    assert collected_data[2]["action"] == "clean"
+    assert collected_data[2]["msg"] == "Room deleted."
+    assert collected_data[3]["path"] == "test2.txt"
+    assert collected_data[3]["room"] == "text2:file2:" + fim.get_id("test2.txt")
+    assert collected_data[3]["action"] == "clean"
+    assert collected_data[3]["msg"] == "file loader removed."
+
+    await jp_serverapp.web_app.settings["jupyter_server_ydoc"].stop_extension()
+    del jp_serverapp.web_app.settings["file_id_manager"]