fix(bookmark): Run async on_restore callbacks before the session's effects

shiny/bookmark/_bookmark.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import warnings
 from abc import ABC, abstractmethod
 from pathlib import Path
@@ -29,6 +30,13 @@
     ExpressStubSession = Any
 
 
+ON_RESTORE_PRIORITY = 1000000
+"""
+Priority of the effect that invokes the `on_restore` callbacks. The session's effects
+with a lower priority run after the callbacks finish, including their async parts.
+"""
+
+
 class Bookmark(ABC):
     """
     A session's bookmarking interface.
@@ -164,6 +172,12 @@ def on_restore(
         Registers a function that will be called just before restoring state.
 
         This callback will be executed **before** the bookmark state is restored.
+        The session's effects and outputs with a priority below 1000000 (the default
+        is 0) start to run only after all `on_restore` callbacks have finished,
+        including the awaited parts of async callbacks. Because those effects wait,
+        an async callback must not await anything that one of them produces (such
+        as an extended task's result or a value set by an effect); the callback
+        would wait forever.
         """
         return self._on_restore_callbacks.register(wrap_async(callback))
 
@@ -341,6 +355,9 @@ def __init__(self, root_session: AppSession):
         super().__init__()
 
         self._root_session = root_session
+        # Resolved when the `on_restore` callbacks have finished (see
+        # `_create_effects()`).
+        self._restore_done: Optional[asyncio.Future[None]] = None
 
         # # Do not set it to avoid supporting a `None` type.
         # # Instead, only use it after it's been set.
@@ -366,6 +383,15 @@ def _set_restore_context(self, restore_context: RestoreContext):
         """
         self._restore_context_value = restore_context
 
+    def _restore_gate(self, priority: int) -> Optional[asyncio.Future[None]]:
+        """
+        Return the future that an effect with `priority` must wait on before it runs,
+        or `None` if it can run now.
+        """
+        if priority >= ON_RESTORE_PRIORITY:
+            return None
+        return self._restore_done
+
     def _create_effects(self) -> None:
         """
         Create the bookmarking `@reactive.effect`s for the session.
@@ -412,10 +438,16 @@ def init_error_message():
                     )
 
             # Run the on_restore function at the beginning of the flush cycle, but after
-            # the server function has been executed.
-            @reactive.effect(priority=1000000)
-            @otel_suppress
-            async def invoke_on_restore_callbacks():
+            # the server function has been executed. A flush doesn't wait for an
+            # effect's async part, so the session's other effects wait on
+            # `_restore_done` instead (see `_restore_gate()`). Unlike R, whose
+            # `onRestore()` callbacks are synchronous, callbacks can be async.
+            restore_done: asyncio.Future[None] = (
+                asyncio.get_running_loop().create_future()
+            )
+            self._restore_done = restore_done
+
+            async def invoke_callbacks() -> None:
                 if self._on_restore_callbacks.count() == 0:
                     return
 
@@ -436,6 +468,17 @@ async def invoke_on_restore_callbacks():
                             type="error",
                         )
 
+            @reactive.effect(priority=ON_RESTORE_PRIORITY)
+            @otel_suppress
+            async def invoke_on_restore_callbacks():
+                try:
+                    await invoke_callbacks()
+                finally:
+                    # Also when the run is cancelled; `invoke_callbacks()` handles
+                    # the callbacks' errors.
+                    if not restore_done.done():
+                        restore_done.set_result(None)
+
             # Run the on_restored function after the flush cycle completes and
             # information is sent to the client.
             @root_session.on_flushed

shiny/reactive/_reactives.py

@@ -1138,6 +1138,9 @@ def _create_context(self) -> Context:
         # Store the context explicitly in Effect object
         # TODO: More explanation here
         self._ctx = ctx
+        # The priority this context's run was queued at; `set_priority()` doesn't
+        # change it.
+        queued_priority = self._priority
 
         def on_invalidate_cb() -> None:
             # Context is invalidated, so we don't need to store a reference to it
@@ -1151,7 +1154,9 @@ def on_invalidate_cb() -> None:
                 return
 
             def _continue() -> None:
-                ctx.add_pending_flush(self._priority)
+                nonlocal queued_priority
+                queued_priority = self._priority
+                ctx.add_pending_flush(queued_priority)
                 if self._session:
                     self._session._increment_busy_count()
 
@@ -1172,6 +1177,15 @@ async def on_flush_cb() -> None:
             try:
                 if previous is not None and not previous.done():
                     await asyncio.shield(previous)
+                # While the session restores its bookmarked state, wait for the
+                # `on_restore` callbacks to finish.
+                restore = (
+                    self._session._restore_gate(queued_priority)
+                    if self._session
+                    else None
+                )
+                if restore is not None and not restore.done():
+                    await asyncio.shield(restore)
                 if not self._destroyed:
                     await self._run()
             except Exception as e:

shiny/session/_session.py

@@ -842,6 +842,13 @@ def _otel_reactive_update_span(self) -> Optional[Span]:
         """The `reactive_update` span of this session's current cycle, if any."""
         return None
 
+    def _restore_gate(self, priority: int) -> Optional[asyncio.Future[None]]:
+        """
+        The future that this session's effect with `priority` waits on before it runs,
+        while the session restores its bookmarked state; `None` if it can run now.
+        """
+        return None
+
     def run_once_when_idle(self, fn: Callable[[], object]) -> None:
         """
         Run a function once, at the start of this session's next cycle.
@@ -1932,6 +1939,11 @@ def _decrement_busy_count(self) -> None:
     def _otel_reactive_update_span(self) -> Optional[Span]:
         return self._otel_cycle_span
 
+    def _restore_gate(self, priority: int) -> Optional[asyncio.Future[None]]:
+        if isinstance(self.bookmark, BookmarkApp):
+            return self.bookmark._restore_gate(priority)
+        return None
+
     def _end_otel_cycle_span(self) -> None:
         span, self._otel_cycle_span = self._otel_cycle_span, None
         if span is not None:
@@ -2266,6 +2278,9 @@ def _decrement_busy_count(self) -> None:
     def _otel_reactive_update_span(self) -> Optional[Span]:
         return self._root_session._otel_reactive_update_span()
 
+    def _restore_gate(self, priority: int) -> Optional[asyncio.Future[None]]:
+        return self._root_session._restore_gate(priority)
+
     def run_once_when_idle(self, fn: Callable[[], object]) -> None:
         _validate_idle_action(fn)