save references to scheduled tasks

Author: dlqqq

jupyter_collaboration/rooms.py

@@ -45,17 +45,12 @@ def __init__(
 
         self._cleaner: asyncio.Task | None = None
 
+        # the current `self._maybe_save_document()` task.
+        self._maybe_save_task: asyncio.Task | None = None
+
         # the task currently saving to disk via FileLoader, if any.
         self._save_task: asyncio.Task | None = None
 
-        # flag that indicates whether a previous call to
-        # `self._maybe_save_document()` is waiting to save.
-        #
-        # if `self._maybe_save_document()` is called while this flag is `True`,
-        # then the call does nothing, as a previous task will save the Ydoc
-        # within `self._save_delay` seconds.
-        self._waiting_to_save = False
-
         # flag that indicates whether a previous call to
         # `self._maybe_save_document()` should call itself again after
         # `self._save_task` finishes. this is set to `True` when a document
@@ -219,7 +214,18 @@ def _on_document_change(self, target: str, event: Any) -> None:
             document. This tasks are debounced (60 seconds by default) so we
             need to cancel previous tasks before creating a new one.
         """
-        asyncio.create_task(self._maybe_save_document())
+        if self._maybe_save_task and not self._maybe_save_task.done():
+            # only one `self._maybe_save_task` needs to be running.
+            #
+            # if `self._save_task` is already running, then we need to set the
+            # `self._should_resave` flag to `True` to indicate that the
+            # `_maybe_save_document()` coroutine needs to be re-scheduled after
+            # the current `self._save_task` completes.
+            if self._save_task and not self._save_task.done():
+                self._should_resave = True
+            return
+
+        self._maybe_save_task = asyncio.create_task(self._maybe_save_document())
 
     async def _maybe_save_document(self) -> None:
         """
@@ -234,23 +240,8 @@ async def _maybe_save_document(self) -> None:
             # TODO: fix this; if _save_delay is unset, then this never writes to disk
             return
 
-        if self._waiting_to_save:
-            # if a previously spawned `self._maybe_save_document()` task is
-            # waiting to save, then that task will save the Ydoc within
-            # `self._save_delay` seconds. therefore we can return early.
-            return
-        
-        if self._save_task and not self._save_task.done():
-            # if we are currently saving, then set the `_should_resave`
-            # flag. this indicates to the currently running `self._save_task`
-            # that it should re-call this method after it completes.
-            self._should_resave = True
-            return
-        
         # save after `self._save_delay` seconds of inactivity
-        self._waiting_to_save = True
         await asyncio.sleep(self._save_delay)
-        self._waiting_to_save = False
 
         # all async code (i.e. await statements) must be part of this try/except block
         # because this coroutine is run in a cancellable task and cancellation is handled here
@@ -271,7 +262,7 @@ async def _maybe_save_document(self) -> None:
 
             if self._should_resave:
                 self._should_resave = False
-                asyncio.create_task(self._maybe_save_document())
+                self._maybe_save_task = asyncio.create_task(self._maybe_save_document())
 
         except asyncio.CancelledError:
             return