Support ypy-websocket v0.12 (#168)

* Support ypy-websocket v0.12

* Add type

* ypy-websocket >=0.12.1

Author: davidbrochart

binder/environment.yml

@@ -17,4 +17,4 @@ dependencies:
   - pip:
     - jupyter_server_fileid >=0.7.0
     - jupyter_ydoc >=1.0.0
-    - ypy-websocket >=0.8.3
+    - ypy-websocket >=0.12.0

jupyter_collaboration/handlers.py

@@ -54,28 +54,23 @@ class YDocWebSocketHandler(WebSocketHandler, JupyterHandler):
     """
 
     _message_queue: asyncio.Queue[Any]
+    _background_tasks: set[asyncio.Task]
 
-    def initialize(
-        self,
-        ywebsocket_server: JupyterWebsocketServer,
-        file_loaders: FileLoaderMapping,
-        ystore_class: type[BaseYStore],
-        document_cleanup_delay: float | None = 60.0,
-        document_save_delay: float | None = 1.0,
-    ) -> None:
-        # File ID manager cannot be passed as argument as the extension may load after this one
-        self._file_id_manager = self.settings["file_id_manager"]
-        self._file_loaders = file_loaders
-        self._cleanup_delay = document_cleanup_delay
-        self._websocket_server = ywebsocket_server
+    def create_task(self, aw):
+        task = asyncio.create_task(aw)
+        self._background_tasks.add(task)
+        task.add_done_callback(self._background_tasks.discard)
 
-        self._message_queue = asyncio.Queue()
+    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()
 
         # Get room
         self._room_id: str = self.request.path.split("/")[-1]
 
         if self._websocket_server.room_exists(self._room_id):
-            self.room: YRoom = self._websocket_server.get_room(self._room_id)
+            self.room: YRoom = await self._websocket_server.get_room(self._room_id)
 
         else:
             if self._room_id.count(":") >= 2:
@@ -92,7 +87,7 @@ def initialize(
                 path = self._file_id_manager.get_path(file_id)
                 path = Path(path)
                 updates_file_path = str(path.parent / f".{file_type}:{path.name}.y")
-                ystore = ystore_class(path=updates_file_path, log=self.log)
+                ystore = self._ystore_class(path=updates_file_path, log=self.log)
                 self.room = DocumentRoom(
                     self._room_id,
                     file_format,
@@ -101,16 +96,37 @@ def initialize(
                     self.event_logger,
                     ystore,
                     self.log,
-                    document_save_delay,
+                    self._document_save_delay,
                 )
 
             else:
                 # TransientRoom
                 # 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)
 
+        return await super().prepare()
+
+    def initialize(
+        self,
+        ywebsocket_server: JupyterWebsocketServer,
+        file_loaders: FileLoaderMapping,
+        ystore_class: type[BaseYStore],
+        document_cleanup_delay: float | None = 60.0,
+        document_save_delay: float | None = 1.0,
+    ) -> None:
+        self._background_tasks = set()
+        # File ID manager cannot be passed as argument as the extension may load after this one
+        self._file_id_manager = self.settings["file_id_manager"]
+        self._file_loaders = file_loaders
+        self._ystore_class = ystore_class
+        self._cleanup_delay = document_cleanup_delay
+        self._document_save_delay = document_save_delay
+        self._websocket_server = ywebsocket_server
+        self._message_queue = asyncio.Queue()
+
     @property
     def path(self):
         """
@@ -150,9 +166,7 @@ async def open(self, room_id):
         """
         On connection open.
         """
-        task = asyncio.create_task(self._websocket_server.serve(self))
-        self._websocket_server.background_tasks.add(task)
-        task.add_done_callback(self._websocket_server.background_tasks.discard)
+        self.create_task(self._websocket_server.serve(self))
 
         if isinstance(self.room, DocumentRoom):
             # Close the connection if the document session expired

jupyter_collaboration/websocketserver.py

@@ -58,21 +58,11 @@ async def clean(self):
         #         self.log.warning(msg)
         #         self.log.debug("Pending tasks: %r", pending)
 
+        self.stop()
         tasks = []
-        for name, room in list(self.rooms.items()):
-            try:
-                self.delete_room(name=name)
-            except Exception as e:  # Capture exception as room may be auto clean
-                msg = f"Failed to delete room {name}"
-                self.log.debug(msg, exc_info=e)
-            else:
-                tasks.append(room._broadcast_task)  # FIXME should be upstreamed
         if self.monitor_task is not None:
             self.monitor_task.cancel()
             tasks.append(self.monitor_task)
-        for task in self.background_tasks:
-            task.cancel()  # FIXME should be upstreamed
-            tasks.append(task)
 
         if tasks:
             _, pending = await asyncio.wait(tasks, timeout=3)
@@ -103,7 +93,7 @@ def add_room(self, path: str, room: YRoom) -> None:
         """
         self.rooms[path] = room
 
-    def get_room(self, path: str) -> YRoom:
+    async def get_room(self, path: str) -> YRoom:
         """
         Returns the room for the specified room ID or raises a RoomNotFound
         error if the room doesn't exist.
@@ -121,7 +111,9 @@ def get_room(self, path: str) -> YRoom:
             # Document rooms need a file
             raise RoomNotFound
 
-        return self.rooms[path]
+        room = self.rooms[path]
+        await self.start_room(room)
+        return room
 
     async def serve(self, websocket: WebSocketHandler) -> None:
         # start monitoring here as the event loop is not yet available when initializing the object

pyproject.toml

@@ -29,7 +29,7 @@ classifiers = [
 dependencies = [
     "jupyter_server>=2.0.0,<3.0.0",
     "jupyter_ydoc>=1.0.1,<2.0.0",
-    "ypy-websocket>=0.8.3,<0.9.0",
+    "ypy-websocket>=0.12.1,<0.13.0",
     "jupyter_events",
     "jupyter_server_fileid>=0.6.0,<1"
 ]