implement contents lock to prevent overlapping reads/writes

Author: dlqqq

jupyter_server_documents/rooms/yroom.py

@@ -308,7 +308,7 @@ async def _process_message_queue(self) -> None:
             # queue is empty in `self.stop()`.
             self._message_queue.task_done()
 
-        self.log.info(
+        self.log.debug(
             "Stopped `self._process_message_queue()` background task "
             f"for YRoom '{self.room_id}'."
         )
@@ -709,13 +709,13 @@ def stop(self, close_code: int = 1001, immediately: bool = False):
         if immediately:
             self._clear_ydoc()
         else:
-            # TODO: how to handle restarts after this safely?
             self._loop.create_task(
                 self._save_then_clear_ydoc()
             )
 
         self._stopped = True
 
+
     def _clear_ydoc(self):
         """
         Clears the YDoc, awareness, and JupyterYDoc, freeing their memory to the
@@ -727,18 +727,27 @@ def _clear_ydoc(self):
             ydoc=self._ydoc,
             awareness=self._awareness
         )
+    
 
     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
+        lock to prevent overlapping reads & writes to a single file.
+        """
         await self.file_api.save(self._jupyter_ydoc)
         self._clear_ydoc()
-    
+
+
     @property
     def stopped(self) -> bool:
         """
         Returns whether the room is stopped.
         """
         return self._stopped
 
+
     def restart(self, close_code: int = 1001, immediately: bool = False):
         """
         Restarts the YRoom. This method re-initializes & reloads the YDoc,

jupyter_server_documents/rooms/yroom_file_api.py

@@ -36,7 +36,18 @@ class YRoomFileAPI:
     log: logging.Logger
 
     _fileid_manager: BaseFileIdManager
+    """
+    Stores a reference to the Jupyter Server's File ID Manager.
+    """
+
     _contents_manager: AsyncContentsManager | ContentsManager
+    """
+    Stores a reference to the Jupyter Server's ContentsManager.
+
+    NOTE: any calls made on this attribute should acquire & release the
+    `_content_lock`. See `_content_lock` for more info.
+    """
+
     _loop: asyncio.AbstractEventLoop
     _save_scheduled: bool
     _content_loading: bool
@@ -81,6 +92,13 @@ class YRoomFileAPI:
     not running.
     """
 
+    _content_lock: asyncio.Lock
+    """
+    An `asyncio.Lock` that ensures `ContentsManager` calls reading/writing for a
+    single file do not overlap. This prevents file corruption scenarios like
+    dual-writes or dirty-reads.
+    """
+
     def __init__(
         self,
         *,
@@ -108,9 +126,10 @@ def __init__(
         self._last_modified = None
         self._stopped = False
 
-        # Initialize loading & loaded states
+        # Initialize content-related primitives
         self._content_loading = False
         self._content_load_event = asyncio.Event()
+        self._content_lock = asyncio.Lock()
 
 
     def get_path(self) -> str | None:
@@ -169,11 +188,12 @@ async def _load_content(self, jupyter_ydoc: YBaseDoc) -> None:
 
         # Load the content of the file from the path
         self.log.info(f"Loading content for room ID '{self.room_id}', found at path: '{path}'.")
-        file_data = await ensure_async(self._contents_manager.get(
-            path,
-            type=self.file_type,
-            format=self.file_format
-        ))
+        async with self._content_lock:
+            file_data = await ensure_async(self._contents_manager.get(
+                path,
+                type=self.file_type,
+                format=self.file_format
+            ))
 
         # Set JupyterYDoc content and set `dirty = False` to hide the "unsaved
         # changes" icon in the UI
@@ -289,9 +309,10 @@ async def _check_file(self):
         # If this raises `HTTPError(404)`, that indicates the file was
         # moved/deleted out-of-band.
         try:
-            file_data = await ensure_async(self._contents_manager.get(
-                path=path, format=file_format, type=file_type, content=False
-            ))
+            async with self._content_lock:
+                file_data = await ensure_async(self._contents_manager.get(
+                    path=path, format=file_format, type=file_type, content=False
+                ))
         except HTTPError as e:
             # If not 404, re-raise the exception as it is unknown
             if (e.status_code != 404):
@@ -344,14 +365,15 @@ async def save(self, jupyter_ydoc: YBaseDoc):
             self._save_scheduled = False
 
             # Save the YDoc via the ContentsManager
-            file_data = await ensure_async(self._contents_manager.save(
-                {
-                    "format": file_format,
-                    "type": file_type,
-                    "content": content,
-                },
-                path
-            ))
+            async with self._content_lock:
+                file_data = await ensure_async(self._contents_manager.save(
+                    {
+                        "format": file_format,
+                        "type": file_type,
+                        "content": content,
+                    },
+                    path
+                ))
 
             # Set most recent `last_modified` timestamp
             if file_data['last_modified']: