simplify concurrent render process

Author: rmorshea

src/py/reactpy/reactpy/core/_life_cycle_hook.py

@@ -101,12 +101,12 @@ async def stop_effect():
         "_context_providers",
         "_current_state_index",
         "_effect_funcs",
-        "_effect_tasks",
         "_effect_stops",
+        "_effect_tasks",
         "_render_access",
         "_rendered_atleast_once",
         "_schedule_render_callback",
-        "_schedule_render_later",
+        "_scheduled_render",
         "_state",
         "component",
     )
@@ -119,7 +119,7 @@ def __init__(
     ) -> None:
         self._context_providers: dict[Context[Any], ContextProviderType[Any]] = {}
         self._schedule_render_callback = schedule_render
-        self._schedule_render_later = False
+        self._scheduled_render = False
         self._rendered_atleast_once = False
         self._current_state_index = 0
         self._state: tuple[Any, ...] = ()
@@ -129,10 +129,15 @@ def __init__(
         self._render_access = Semaphore(1)  # ensure only one render at a time
 
     def schedule_render(self) -> None:
-        if self._is_rendering():
-            self._schedule_render_later = True
+        if self._scheduled_render:
+            return None
+        try:
+            self._schedule_render_callback()
+        except Exception:
+            msg = f"Failed to schedule render via {self._schedule_render_callback}"
+            logger.exception(msg)
         else:
-            self._schedule_render()
+            self._scheduled_render = True
 
     def use_state(self, function: Callable[[], T]) -> T:
         if not self._rendered_atleast_once:
@@ -166,6 +171,7 @@ def get_context_provider(
     async def affect_component_will_render(self, component: ComponentType) -> None:
         """The component is about to render"""
         await self._render_access.acquire()
+        self._scheduled_render = False
         self.component = component
         self.set_current()
 
@@ -183,9 +189,6 @@ async def affect_layout_did_render(self) -> None:
         self._effect_stops.append(stop)
         self._effect_tasks.extend(create_task(e(stop)) for e in self._effect_funcs)
         self._effect_funcs.clear()
-        if self._schedule_render_later:
-            self._schedule_render()
-        self._schedule_render_later = False
 
     async def affect_component_will_unmount(self) -> None:
         """The component is about to be removed from the layout"""
@@ -215,14 +218,3 @@ def unset_current(self) -> None:
         """Unset this hook as the active hook in this thread"""
         if _HOOK_STATE.get().pop() is not self:
             raise RuntimeError("Hook stack is in an invalid state")  # nocov
-
-    def _is_rendering(self) -> bool:
-        return self._render_access.value == 0
-
-    def _schedule_render(self) -> None:
-        try:
-            self._schedule_render_callback()
-        except Exception:
-            logger.exception(
-                f"Failed to schedule render via {self._schedule_render_callback}"
-            )

src/py/reactpy/reactpy/core/layout.py

@@ -3,11 +3,11 @@
 import abc
 from asyncio import (
     FIRST_COMPLETED,
+    CancelledError,
     Event,
     Queue,
     Task,
     create_task,
-    gather,
     get_running_loop,
     wait,
 )
@@ -27,6 +27,8 @@
 from uuid import uuid4
 from weakref import ref as weakref
 
+from anyio import Semaphore
+
 from reactpy.config import (
     REACTPY_CHECK_VDOM_SPEC,
     REACTPY_DEBUG_MODE,
@@ -55,6 +57,7 @@ class Layout:
         "_event_handlers",
         "_rendering_queue",
         "_render_tasks",
+        "_render_tasks_ready",
         "_root_life_cycle_state_id",
         "_model_states_by_life_cycle_state_id",
     )
@@ -73,21 +76,28 @@ async def __aenter__(self) -> Layout:
         # create attributes here to avoid access before entering context manager
         self._event_handlers: EventHandlerDict = {}
         self._render_tasks: set[Task[LayoutUpdateMessage]] = set()
+        self._render_tasks_ready: Semaphore = Semaphore(0)
 
         self._rendering_queue: _ThreadSafeQueue[_LifeCycleStateId] = _ThreadSafeQueue()
-        root_model_state = _new_root_model_state(self.root, self._rendering_queue.put)
+        root_model_state = _new_root_model_state(self.root, self._schedule_render_task)
 
         self._root_life_cycle_state_id = root_id = root_model_state.life_cycle_state.id
-        self._rendering_queue.put(root_id)
-
         self._model_states_by_life_cycle_state_id = {root_id: root_model_state}
+        self._schedule_render_task(root_id)
 
         return self
 
     async def __aexit__(self, *exc: Any) -> None:
         root_csid = self._root_life_cycle_state_id
         root_model_state = self._model_states_by_life_cycle_state_id[root_csid]
-        await gather(*self._render_tasks, return_exceptions=True)
+
+        for t in self._render_tasks:
+            t.cancel()
+            try:
+                await t
+            except CancelledError:
+                pass
+
         await self._unmount_model_states([root_model_state])
 
         # delete attributes here to avoid access after exiting context manager
@@ -137,40 +147,11 @@ async def _serial_render(self) -> LayoutUpdateMessage:  # nocov
 
     async def _concurrent_render(self) -> LayoutUpdateMessage:
         """Await the next available render. This will block until a component is updated"""
-        while True:
-            render_completed = (
-                create_task(wait(self._render_tasks, return_when=FIRST_COMPLETED))
-                if self._render_tasks
-                else get_running_loop().create_future()
-            )
-            queue_ready = create_task(self._rendering_queue.ready())
-            try:
-                await wait((queue_ready, render_completed), return_when=FIRST_COMPLETED)
-            finally:
-                # Ensure we delete this task to avoid warnings that
-                # task was deleted without being awaited.
-                queue_ready.cancel()
-
-            if render_completed.done():
-                done, _ = await render_completed
-                update_task: Task[LayoutUpdateMessage] = done.pop()
-                self._render_tasks.remove(update_task)
-                return update_task.result()
-            else:
-                model_state_id = await self._rendering_queue.get()
-                try:
-                    model_state = self._model_states_by_life_cycle_state_id[
-                        model_state_id
-                    ]
-                except KeyError:
-                    logger.debug(
-                        "Did not render component with model state ID "
-                        f"{model_state_id!r} - component already unmounted"
-                    )
-                else:
-                    self._render_tasks.add(
-                        create_task(self._create_layout_update(model_state))
-                    )
+        await self._render_tasks_ready.acquire()
+        done, _ = await wait(self._render_tasks, return_when=FIRST_COMPLETED)
+        update_task: Task[LayoutUpdateMessage] = done.pop()
+        self._render_tasks.remove(update_task)
+        return update_task.result()
 
     async def _create_layout_update(
         self, old_state: _ModelState
@@ -403,7 +384,7 @@ async def _render_model_children(
                         index,
                         key,
                         child,
-                        self._rendering_queue.put,
+                        self._schedule_render_task,
                     )
                 elif old_child_state.is_component_state and (
                     old_child_state.life_cycle_state.component.type != child.type
@@ -415,15 +396,15 @@ async def _render_model_children(
                         index,
                         key,
                         child,
-                        self._rendering_queue.put,
+                        self._schedule_render_task,
                     )
                 else:
                     new_child_state = _update_component_model_state(
                         old_child_state,
                         new_state,
                         index,
                         child,
-                        self._rendering_queue.put,
+                        self._schedule_render_task,
                     )
                 await self._render_component(
                     exit_stack, old_child_state, new_child_state, child
@@ -458,7 +439,7 @@ async def _render_model_children_without_old_state(
                 new_state.children_by_key[key] = child_state
             elif child_type is _COMPONENT_TYPE:
                 child_state = _make_component_model_state(
-                    new_state, index, key, child, self._rendering_queue.put
+                    new_state, index, key, child, self._schedule_render_task
                 )
                 await self._render_component(exit_stack, None, child_state, child)
             else:
@@ -479,6 +460,21 @@ async def _unmount_model_states(self, old_states: list[_ModelState]) -> None:
 
             to_unmount.extend(model_state.children_by_key.values())
 
+    def _schedule_render_task(self, lcs_id: _LifeCycleStateId) -> None:
+        if not REACTPY_FEATURE_CONCURRENT_RENDERING.current:
+            self._rendering_queue.put(lcs_id)
+            return None
+        try:
+            model_state = self._model_states_by_life_cycle_state_id[lcs_id]
+        except KeyError:
+            logger.debug(
+                "Did not render component with model state ID "
+                f"{lcs_id!r} - component already unmounted"
+            )
+        else:
+            self._render_tasks.add(create_task(self._create_layout_update(model_state)))
+            self._render_tasks_ready.release()
+
     def __repr__(self) -> str:
         return f"{type(self).__name__}({self.root})"
 

src/py/reactpy/tests/test_core/test_hooks.py

@@ -28,10 +28,15 @@ def SimpleComponentWithHook():
 
 
 async def test_simple_stateful_component():
+    index = 0
+
+    def set_index(x):
+        return None
+
     @reactpy.component
     def SimpleStatefulComponent():
+        nonlocal index, set_index
         index, set_index = reactpy.hooks.use_state(0)
-        set_index(index + 1)
         return reactpy.html.div(index)
 
     sse = SimpleStatefulComponent()
@@ -45,6 +50,7 @@ def SimpleStatefulComponent():
                 "children": [{"tagName": "div", "children": ["0"]}],
             },
         )
+        set_index(index + 1)
 
         update_2 = await layout.render()
         assert update_2 == update_message(
@@ -54,6 +60,7 @@ def SimpleStatefulComponent():
                 "children": [{"tagName": "div", "children": ["1"]}],
             },
         )
+        set_index(index + 1)
 
         update_3 = await layout.render()
         assert update_3 == update_message(
@@ -1026,13 +1033,13 @@ def SetStateDuringRender():
 
     async with Layout(SetStateDuringRender()) as layout:
         await layout.render()
-        assert render_count.current == 1
-        await layout.render()
-        assert render_count.current == 2
 
-        # there should be no more renders to perform
+        # we expect a second render to be triggered in the background
+        await poll(lambda: render_count.current).until_equals(2)
+
+        # there should be no more renders that happen
         with pytest.raises(asyncio.TimeoutError):
-            await asyncio.wait_for(layout.render(), timeout=0.1)
+            await poll(lambda: render_count.current).until_equals(3, timeout=0.1)
 
 
 @pytest.mark.skipif(not REACTPY_DEBUG_MODE.current, reason="only logs in debug mode")

src/py/reactpy/tests/test_core/test_layout.py

@@ -22,6 +22,7 @@
 )
 from reactpy.utils import Ref
 from tests.tooling import select
+from tests.tooling.aio import Event
 from tests.tooling.common import event_message, update_message
 from tests.tooling.hooks import use_force_render, use_toggle
 from tests.tooling.layout import layout_runner
@@ -1250,3 +1251,49 @@ def App():
         c, c_info = find_element(tree, select.id_equals("C"))
         assert c_info.path == (0, 1, 0)
         assert c["attributes"]["color"] == "blue"
+
+
+async def test_concurrent_renders():
+    child_1_hook = HookCatcher()
+    child_2_hook = HookCatcher()
+    child_1_rendered = Event()
+    child_2_rendered = Event()
+    child_1_render_count = Ref(0)
+    child_2_render_count = Ref(0)
+
+    @component
+    def outer():
+        return html._(child_1(), child_2())
+
+    @component
+    @child_1_hook.capture
+    def child_1():
+        child_1_rendered.set()
+        child_1_render_count.current += 1
+
+    @component
+    @child_2_hook.capture
+    def child_2():
+        child_2_rendered.set()
+        child_2_render_count.current += 1
+
+    async with Layout(outer()) as layout:
+        await layout.render()
+
+        # clear render events and counts
+        child_1_rendered.clear()
+        child_2_rendered.clear()
+        child_1_render_count.current = 0
+        child_2_render_count.current = 0
+
+        # we schedule two renders but expect only one
+        child_1_hook.latest.schedule_render()
+        child_1_hook.latest.schedule_render()
+        child_2_hook.latest.schedule_render()
+        child_2_hook.latest.schedule_render()
+
+        await child_1_rendered.wait()
+        await child_2_rendered.wait()
+
+        assert child_1_render_count.current == 1
+        assert child_2_render_count.current == 1