diff --git a/shiny/bookmark/_bookmark.py b/shiny/bookmark/_bookmark.py index 1fd5f4dd6..90bc8548a 100644 --- a/shiny/bookmark/_bookmark.py +++ b/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 diff --git a/shiny/reactive/_reactives.py b/shiny/reactive/_reactives.py index 87b7af660..12c152dcf 100644 --- a/shiny/reactive/_reactives.py +++ b/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: diff --git a/shiny/session/_session.py b/shiny/session/_session.py index 092ca1a95..6983240a5 100644 --- a/shiny/session/_session.py +++ b/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) diff --git a/tests/pytest/test_concurrency.py b/tests/pytest/test_concurrency.py index 7751c7d45..5858beef9 100644 --- a/tests/pytest/test_concurrency.py +++ b/tests/pytest/test_concurrency.py @@ -20,6 +20,7 @@ from shiny import App, Inputs, Outputs, Session, module, reactive, render, ui from shiny._connection import MockConnection from shiny._deprecated import ShinyDeprecationWarning +from shiny.bookmark import RestoreState from shiny.bookmark._bookmark import BookmarkApp from shiny.bookmark._restore_state import RestoreContext from shiny.reactive._core import ( @@ -2856,6 +2857,339 @@ async def producer() -> None: await c.close() +# ---------------------------------------------------------------------------- +# Async bookmark `on_restore` callbacks finish before the session's effects run +# ---------------------------------------------------------------------------- + + +class RestoringClient(RecordingClient): + """A RecordingClient for an app that restores its bookmarked state from the URL.""" + + def __init__(self, server: Callable[[Inputs, Outputs, Session], None]) -> None: + self.recording = RecordingConnection() + self.conn = self.recording + app = App(lambda request: ui.TagList(), server, bookmark_store="url") + self.session = app._create_session(self.conn) + self.recording.session = self.session + self.task = asyncio.create_task(self.session._run()) + + def init(self, **data: object) -> None: + """Start the session, restoring `x="restored"`.""" + url_search = "?_inputs_&x=%22restored%22" + self.send( + {"method": "init", "data": {".clientdata_url_search": url_search, **data}} + ) + + def output_values(self, name: str) -> list[object]: + return [m["values"][name] for m in self.sent if name in m.get("values", {})] # type: ignore + + +@pytest.mark.asyncio +async def test_async_on_restore_finishes_before_other_effects_start(): + seen: list[str] = [] + + def server(input: Inputs, output: Outputs, session: Session) -> None: + v = reactive.value("default") + + @session.bookmark.on_restore + async def _(state: RestoreState) -> None: + await asyncio.sleep(0.01) + v.set(state.input["x"]) + + @reactive.effect + def _(): + seen.append(v()) + + @render.text + def out(): + return v() + + c = RestoringClient(server) + try: + c.init(**{".clientdata_output_out_hidden": False}) + assert await wait_until(lambda: c.output_values("out") == ["restored"]) + assert seen == ["restored"] + finally: + await c.close() + + +@pytest.mark.asyncio +async def test_effects_waiting_on_restore_start_in_priority_order(): + log: list[str] = [] + release = asyncio.Event() + + def server(input: Inputs, output: Outputs, session: Session) -> None: + @session.bookmark.on_restore + async def _(state: RestoreState) -> None: + await release.wait() + log.append("restored") + + @reactive.effect(priority=-1) + def _(): + log.append("low") + + @reactive.effect(priority=1) + def _(): + log.append("high") + + c = RestoringClient(server) + try: + c.init() + # Let both effects reach the restore gate before it opens. A short timer + # in the callback isn't enough: before Python 3.13, Windows' clock has a + # 15.6 ms resolution, so a 10 ms sleep expired on the next loop pass, and + # `low` started after the gate had opened, ahead of the `high` still + # being woken. + await asyncio.sleep(0.05) + assert log == [] + release.set() + assert await wait_until(lambda: len(log) == 3) + assert log == ["restored", "high", "low"] + finally: + await c.close() + + +@pytest.mark.asyncio +async def test_effect_above_restore_priority_does_not_wait_for_restore(): + release = asyncio.Event() + log: list[str] = [] + + def server(input: Inputs, output: Outputs, session: Session) -> None: + @session.bookmark.on_restore + async def _(state: RestoreState) -> None: + await release.wait() + log.append("restored") + + @reactive.effect(priority=2_000_000) + def _(): + log.append("above") + + c = RestoringClient(server) + try: + c.init() + assert await wait_until(lambda: log == ["above"]) + release.set() + assert await wait_until(lambda: log == ["above", "restored"]) + finally: + release.set() + await c.close() + + +@pytest.mark.asyncio +async def test_slow_restore_does_not_block_other_sessions(): + release = asyncio.Event() + a_log: list[str] = [] + b_seen: list[object] = [] + + def server_a(input: Inputs, output: Outputs, session: Session) -> None: + @session.bookmark.on_restore + async def _(state: RestoreState) -> None: + await release.wait() + + @reactive.effect + def _(): + a_log.append("ran") + + def server_b(input: Inputs, output: Outputs, session: Session) -> None: + @reactive.effect + def _(): + b_seen.append(input.x()) + + a, b = RestoringClient(server_a), Client(server_b) + try: + a.init() + b.send({"method": "init", "data": {"x": 0}}) + assert await wait_until(lambda: b_seen == [0]) + b.update(x=1) + assert await wait_until(lambda: b_seen == [0, 1]) + assert a_log == [] + release.set() + assert await wait_until(lambda: a_log == ["ran"]) + finally: + release.set() + await a.close() + await b.close() + + +@pytest.mark.asyncio +async def test_effects_run_after_an_on_restore_callback_fails(): + log: list[str] = [] + + def server(input: Inputs, output: Outputs, session: Session) -> None: + @session.bookmark.on_restore + async def _(state: RestoreState) -> None: + await asyncio.sleep(0.01) + log.append("raising") + raise RuntimeError("restore failed") + + @reactive.effect + def _(): + log.append("ran") + + c = RestoringClient(server) + try: + with pytest.warns(UserWarning, match="restore failed"): + c.init() + assert await wait_until(lambda: len(log) == 2) + assert log == ["raising", "ran"] + finally: + await c.close() + + +@pytest.mark.asyncio +async def test_effects_run_after_an_on_restore_callback_is_cancelled(): + # The callback's error handling catches `Exception` only, so a cancellation + # ends the restore effect's run with `CancelledError`. + log: list[str] = [] + awaited: list[asyncio.Task[None]] = [] + + async def wait_forever() -> None: + await asyncio.Event().wait() + + def server(input: Inputs, output: Outputs, session: Session) -> None: + @session.bookmark.on_restore + async def _(state: RestoreState) -> None: + awaited.append(asyncio.create_task(wait_forever())) + log.append("restore start") + await awaited[0] + + @reactive.effect + def _(): + log.append("ran") + + c = RestoringClient(server) + try: + c.init() + assert await wait_until(lambda: log == ["restore start"]) + awaited[0].cancel() + assert await wait_until(lambda: log == ["restore start", "ran"]) + finally: + await c.close() + + +@pytest.mark.asyncio +async def test_module_effects_wait_for_restore(): + release = asyncio.Event() + log: list[str] = [] + + @module.server + def mod(input: Inputs, output: Outputs, session: Session) -> None: + @reactive.effect + def _(): + log.append("mod") + + def server(input: Inputs, output: Outputs, session: Session) -> None: + @session.bookmark.on_restore + async def _(state: RestoreState) -> None: + await release.wait() + log.append("restored") + + mod("mod") + + c = RestoringClient(server) + try: + c.init() + await asyncio.sleep(0.02) + assert log == [] + release.set() + assert await wait_until(lambda: log == ["restored", "mod"]) + finally: + release.set() + await c.close() + + +@pytest.mark.asyncio +async def test_destroying_an_effect_waiting_on_restore_leaves_the_others_waiting(): + release = asyncio.Event() + log: list[str] = [] + effects: dict[str, reactive.Effect_] = {} + + def server(input: Inputs, output: Outputs, session: Session) -> None: + @session.bookmark.on_restore + async def _(state: RestoreState) -> None: + await release.wait() + log.append("restored") + + @reactive.effect + def a(): + log.append("a") + + @reactive.effect + def b(): + log.append("b") + + effects.update(a=a, b=b) + + c = RestoringClient(server) + try: + c.init() + assert await wait_until( + lambda: len(effects) == 2 + and all(len(e._run_tasks) == 1 for e in effects.values()) + ) + effects["a"].destroy() + release.set() + assert await wait_until(lambda: log == ["restored", "b"]) + finally: + release.set() + await c.close() + + +@pytest.mark.asyncio +async def test_effect_waits_for_restore_at_the_priority_it_was_queued_at(): + release = asyncio.Event() + log: list[str] = [] + + def server(input: Inputs, output: Outputs, session: Session) -> None: + @session.bookmark.on_restore + async def _(state: RestoreState) -> None: + await release.wait() + log.append("restored") + + @reactive.effect + def e(): + log.append("e") + + # Takes effect on the next invalidation; this run is queued at priority 0. + e.set_priority(2_000_000) + + c = RestoringClient(server) + try: + c.init() + await asyncio.sleep(0.02) + assert log == [] + release.set() + assert await wait_until(lambda: log == ["restored", "e"]) + finally: + release.set() + await c.close() + + +@pytest.mark.asyncio +async def test_session_end_during_restore_cancels_waiting_effects(): + log: list[str] = [] + + def server(input: Inputs, output: Outputs, session: Session) -> None: + @session.bookmark.on_restore + async def _(state: RestoreState) -> None: + try: + log.append("restore start") + await asyncio.Event().wait() + finally: + log.append("restore finally") + + @reactive.effect + def _(): + log.append("effect ran") + + c = RestoringClient(server) + c.init() + assert await wait_until(lambda: log == ["restore start"]) + await c.close() + assert log == ["restore start", "restore finally"] + assert c.session._busy_count == 0 + + @pytest.mark.asyncio async def test_on_flushed_callback_that_updates_an_input_does_not_loop(): # The update registers an `on_flush` callback, which requests an output flush. diff --git a/tests/pytest/test_poll.py b/tests/pytest/test_poll.py index b2622e944..69906de88 100644 --- a/tests/pytest/test_poll.py +++ b/tests/pytest/test_poll.py @@ -55,6 +55,9 @@ def _decrement_busy_count(self) -> None: def _otel_reactive_update_span(self) -> None: return None + def _restore_gate(self, priority: int) -> None: + return None + async def __aenter__(self): self._session_context.__enter__()