stop inactive rooms instead of deleting

Author: dlqqq

jupyter_server_documents/rooms/yroom.py

@@ -76,12 +76,6 @@ class YRoom:
     _ydoc_subscription: pycrdt.Subscription
     """Subscription to YDoc changes."""
 
-    _on_stopping: callable[[], Any] | None
-    """
-    Callback to run as soon as `stop()` or `stop_immediately()` are called. Only
-    set in the constructor.
-    """
-
     _stopped: bool = False
     """
     Whether the YRoom is stopped. Set to `True` when `stop()` is called and set
@@ -101,15 +95,13 @@ def __init__(
         fileid_manager: BaseFileIdManager,
         contents_manager: AsyncContentsManager | ContentsManager,
         event_logger: EventLogger,
-        on_stopping: callable[[], Any] | None = None,
     ):
         # Bind instance attributes
         self.room_id = room_id
         self.log = log
         self._loop = loop
         self._fileid_manager = fileid_manager
         self._contents_manager = contents_manager
-        self._on_stopping = on_stopping
         self._stopped = False
 
         # Initialize YjsClientGroup, YDoc, and Awareness
@@ -671,11 +663,6 @@ def stop(self, close_code: int = 1001, immediately: bool = False):
         - Clears the YDoc, Awareness, and JupyterYDoc, freeing their memory to
         the server. This deletes the YDoc history.
         """
-
-        # TODO: delete dis?
-        if self._on_stopping:
-            self._on_stopping()
-        
         # Disconnect all clients with the given close code
         self.clients.stop(close_code=close_code)
 
@@ -733,7 +720,7 @@ async def _save_then_clear_ydoc(self):
         """
         Saves the JupyterYDoc, then calls `self._clear_ydoc()`.
 
-        This can be run safely in the background because the FileAPI uses an
+        This can be run safely in the background because the FileAPI uses a
         lock to prevent overlapping reads & writes to a single file.
         """
         await self.file_api.save(self._jupyter_ydoc)

jupyter_server_documents/rooms/yroom_manager.py

@@ -13,28 +13,25 @@
 class YRoomManager():
     """
     A singleton that manages all `YRoom` instances in the server extension. This
-    automatically deletes `YRoom` instances if they have had no connected
-    clients or active kernel for >10 seconds.
-
-    Because rooms may be deleted due to inactivity, consumers should only store
-    a reference to the room ID and call `get_room(room_id)` each time a
-    reference to the room is needed. See `get_room()` for more details.
+    automatically stops `YRoom` instances if they have had no connected clients
+    or active kernel for >10 seconds.
     """
 
     _rooms_by_id: dict[str, YRoom]
     """
-    Dictionary of active `YRoom` instances, keyed by room ID.
+    Dictionary of active `YRoom` instances, keyed by room ID. Rooms are never
+    deleted from this dictionary, even if stopped due to inactivity.
 
-    It is guaranteed that if a room is present in the dictionary, the room is
-    not currently stopping. This is ensured by `_handle_yroom_stopping()`.
+    TODO: Delete a room if its file was deleted in/out-of-band or moved
+    out-of-band. See #116.
     """
 
     _inactive_rooms: set[str]
     """
     Set of room IDs that were marked inactive on the last iteration of
     `_watch_rooms()`. If a room is inactive and its ID is present in this set,
-    then the room has been inactive for >10 seconds, and the room should be
-    deleted in `_watch_rooms()`.
+    then the room the room should be stopped as it has been inactive for >10
+    seconds.
     """
 
     _get_fileid_manager: callable[[], BaseFileIdManager]
@@ -80,23 +77,16 @@ def get_room(self, room_id: str) -> YRoom | None:
         """
         Returns the `YRoom` instance for a given room ID. If the instance does
         not exist, this method will initialize one and return it. Otherwise,
-        this method returns the instance from its cache, ensuring that this
-        method is fast in almost all cases.
-
-        Consumers should always call this method each time a reference to the
-        `YRoom` is needed, since rooms may be deleted due to inactivity.
-
-        This method also ensures that the returned room will be alive for >10
-        seconds. This prevents the room from being deleted shortly after the
-        consumer receives it via this method, even if it was inactive.
+        this method returns the instance from its cache.
         """
         # First, ensure this room stays open for >10 seconds by removing it from
         # the inactive set of rooms if it is present.
         self._inactive_rooms.discard(room_id)
 
         # If room exists, return the room
-        if room_id in self._rooms_by_id:
-            return self._rooms_by_id[room_id]
+        yroom = self._rooms_by_id.get(room_id, None)
+        if yroom:
+            return yroom
 
         # Otherwise, create a new room
         try:
@@ -107,7 +97,6 @@ def get_room(self, room_id: str) -> YRoom | None:
                 loop=self.loop,
                 fileid_manager=self.fileid_manager,
                 contents_manager=self.contents_manager,
-                on_stopping=lambda: self._handle_yroom_stopping(room_id),
                 event_logger=self.event_logger,
             )
             self._rooms_by_id[room_id] = yroom
@@ -119,31 +108,21 @@ def get_room(self, room_id: str) -> YRoom | None:
             )
             return None
 
+
     def has_room(self, room_id: str) -> bool:
         """
         Returns whether a `YRoom` instance with a matching `room_id` already
         exists.
         """
         return room_id in self._rooms_by_id
 
-    def _handle_yroom_stopping(self, room_id: str) -> None:
-        """
-        Callback that is run when the YRoom starts stopping. This callback:
-        
-        - Ensures the room is removed from the dictionary, even if the room was
-        stopped directly without `YRoomManager.delete_room()`.
 
-        - Prevents future connections to the stopping room and allows its memory
-        to be freed once complete.
-        """
-        self._rooms_by_id.pop(room_id, None)
-
-        
     def delete_room(self, room_id: str) -> None:
         """
         Gracefully deletes a YRoom given a room ID. This stops the YRoom,
-        closing all Websockets, applying remaining updates, and saves the final
-        content of the YDoc in a background task.
+        closing all Websockets with close code 1001 (server shutting down),
+        applying remaining updates, and saving the final content of the YDoc in
+        a background task.
 
         Returns `True` if the room was deleted successfully. Returns `False` if
         an exception was raised.
@@ -165,65 +144,68 @@ def delete_room(self, room_id: str) -> None:
     async def _watch_rooms(self) -> None:
         """
         Background task that checks all `YRoom` instances every 10 seconds,
-        deleting any rooms that have been inactive for >10 seconds.
+        stopping any rooms that have been inactive for >10 seconds.
 
-        - For rooms providing notebooks: This task deletes the room if it has no
+        - For rooms providing notebooks: This task stops the room if it has no
         connected clients and its kernel execution status is either 'idle' or
         'dead'.
 
-        - For all other rooms: This task deletes the room if it has no connected
+        - For all other rooms: This task stops the room if it has no connected
         clients.
         """
         while True:
             # Check every 10 seconds
             await asyncio.sleep(10)
 
-            # Get all room IDs from the room dictionary in advance, as the
-            # dictionary will mutate if rooms are deleted.
+            # Get all room IDs, except for the global awareness room
             room_ids = set(self._rooms_by_id.keys())
-
-            # Remove the global awareness room ID from this set, as that room
-            # should not be stopped until the server extension is stopped.
             room_ids.discard("JupyterLab:globalAwareness")
 
-            # Iterate through all rooms. If any rooms are empty, stop the rooms.
+            # Iterate through all rooms. If any rooms are empty, stop the room.
             for room_id in room_ids:
-                # Continue if `room_id` is not in the rooms dictionary. This
-                # happens if the room was stopped by something else while this
-                # `for` loop is still running, so we must check.
-                if room_id not in self._rooms_by_id:
-                    self._inactive_rooms.discard(room_id)
-                    continue
-
-                # Continue if the room has any connected clients.
-                room = self._rooms_by_id[room_id]
-                if room.clients.count != 0:
-                    self._inactive_rooms.discard(room_id)
-                    continue
-                
-                # Continue if the room contains a notebook with kernel execution
-                # state neither 'idle' nor 'dead'.
-                # In this case, the notebook kernel may still be running code
-                # cells, so the room should not be closed.
-                awareness = room.get_awareness().get_local_state() or {}
-                execution_state = awareness.get("kernel", {}).get("execution_state", None)
-                if execution_state not in { "idle", "dead", None }:
-                    self._inactive_rooms.discard(room_id)
-                    continue
-
-                # The room is inactive if this statement is reached
-                # Delete the room if was marked as inactive in the last
-                # iteration, otherwise mark it as inactive.
-                if room_id in self._inactive_rooms:
-                    self.log.info(
-                        f"YRoom '{room_id}' has been inactive for >10 seconds. "
-                    )
-                    self.loop.create_task(self.delete_room(room_id))
-                    self._inactive_rooms.discard(room_id)
-                else:
-                    self._inactive_rooms.add(room_id)
+                self._check_room(room_id)
 
 
+    def _check_room(self, room_id: str) -> None:
+        """
+        Checks a room for inactivity.
+
+        - If a room is inactive and not in `_inactive_rooms`, this method adds
+        the room to `_inactive_rooms`. 
+
+        - If a room is inactive and is listed in `_inactive_rooms`, this method
+        stops the room, as it has been inactive for 2 consecutive iterations of
+        `_watch_rooms()`.
+        """
+        # Do nothing if the room has any connected clients.
+        room = self._rooms_by_id[room_id]
+        if room.clients.count != 0:
+            self._inactive_rooms.discard(room_id)
+            return
+        
+        # Do nothing if the room contains a notebook with kernel execution state
+        # neither 'idle' nor 'dead'.
+        # In this case, the notebook kernel may still be running code cells, so
+        # the room should not be closed.
+        awareness = room.get_awareness().get_local_state() or {}
+        execution_state = awareness.get("kernel", {}).get("execution_state", None)
+        if execution_state not in { "idle", "dead", None }:
+            self._inactive_rooms.discard(room_id)
+            return
+
+        # The room is inactive if this line is reached.
+        # Stop the room if was marked as inactive in the last iteration,
+        # otherwise mark it as inactive.
+        if room_id in self._inactive_rooms:
+            self.log.info(
+                f"YRoom '{room_id}' has been inactive for >10 seconds. "
+            )
+            room.stop()
+            self._inactive_rooms.discard(room_id)
+        else:
+            self._inactive_rooms.add(room_id)
+
+
     def stop(self) -> None:
         """
         Gracefully deletes each `YRoom`. See `delete_room()` for more info.