queue backpressure

Author: Archmonger

src/reactpy/core/layout.py

@@ -99,6 +99,7 @@ async def __aexit__(
                 await t
 
         await self._unmount_model_states([root_model_state])
+        await self._rendering_queue.close()
 
         # delete attributes here to avoid access after exiting context manager
         del self._event_handlers
@@ -773,17 +774,39 @@ def __init__(self) -> None:
         self._loop = get_running_loop()
         self._queue: Queue[_Type] = Queue(REACTPY_MAX_QUEUE_SIZE.current)
         self._pending: set[_Type] = set()
+        self._put_tasks: dict[_Type, Task[None]] = {}
 
     def put(self, value: _Type) -> None:
         if value not in self._pending:
             self._pending.add(value)
-            self._loop.call_soon_threadsafe(self._queue.put_nowait, value)
+            self._loop.call_soon_threadsafe(self._schedule_put, value)
+
+    def _schedule_put(self, value: _Type) -> None:
+        self._put_tasks[value] = create_task(self._put_with_backpressure(value))
+
+    async def _put_with_backpressure(self, value: _Type) -> None:
+        try:
+            await self._queue.put(value)
+        except BaseException:
+            self._pending.discard(value)
+            raise
+        finally:
+            self._put_tasks.pop(value, None)
 
     async def get(self) -> _Type:
         value = await self._queue.get()
         self._pending.remove(value)
         return value
 
+    async def close(self) -> None:
+        for task in list(self._put_tasks.values()):
+            task.cancel()
+        for task in list(self._put_tasks.values()):
+            with suppress(CancelledError):
+                await task
+        self._put_tasks.clear()
+        self._pending.clear()
+
 
 def _get_children_info(
     children: list[VdomChild],

tests/test_core/test_layout.py

@@ -11,11 +11,15 @@
 
 import reactpy
 from reactpy import html
-from reactpy.config import REACTPY_ASYNC_RENDERING, REACTPY_DEBUG
+from reactpy.config import (
+    REACTPY_ASYNC_RENDERING,
+    REACTPY_DEBUG,
+    REACTPY_MAX_QUEUE_SIZE,
+)
 from reactpy.core.component import component
 from reactpy.core.events import EventHandler
 from reactpy.core.hooks import use_async_effect, use_effect, use_state
-from reactpy.core.layout import Layout
+from reactpy.core.layout import Layout, _ThreadSafeQueue
 from reactpy.testing import (
     HookCatcher,
     StaticEventHandler,
@@ -102,6 +106,22 @@ def SimpleComponent():
         )
 
 
+async def test_thread_safe_queue_applies_backpressure():
+    with patch.object(REACTPY_MAX_QUEUE_SIZE, "current", 1):
+        queue = _ThreadSafeQueue[int]()
+
+        queue.put(1)
+        queue.put(2)
+
+        await asyncio.sleep(0)
+        assert await asyncio.wait_for(queue.get(), 1) == 1
+
+        await asyncio.sleep(0)
+        assert await asyncio.wait_for(queue.get(), 1) == 2
+
+        await queue.close()
+
+
 async def test_nested_component_layout():
     parent_set_state = reactpy.Ref(None)
     child_set_state = reactpy.Ref(None)