Quick attempt at some fixing async dependencies

Author: Archmonger

src/reactpy/core/hooks.py

@@ -155,8 +155,6 @@ def decorator(func: _SyncEffectFunc) -> None:
             )
 
         async def effect(stop: asyncio.Event) -> None:
-            # Since the effect is asynchronous, we need to make sure we
-            # always clean up the previous effect's resources
             run_effect_cleanup(cleanup_func)
 
             # Execute the effect and store the clean-up function
@@ -219,17 +217,23 @@ def use_async_effect(
     dependencies = _try_to_infer_closure_values(function, dependencies)
     memoize = use_memo(dependencies=dependencies)
     cleanup_func: Ref[_EffectCleanFunc | None] = use_ref(None)
+    pending_task: Ref[asyncio.Task[_EffectCleanFunc | None] | None] = use_ref(None)
 
     def decorator(func: _AsyncEffectFunc) -> None:
         async def effect(stop: asyncio.Event) -> None:
-            # Since the effect is asynchronous, we need to make sure we
-            # always clean up the previous effect's resources
+            # Make sure we always clean up the previous effect's resources
+            if pending_task.current:
+                pending_task.current.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await pending_task.current
+
             run_effect_cleanup(cleanup_func)
 
             # Execute the effect and store the clean-up function.
             # We run this in a task so it can be cancelled if the stop signal
             # is set before the effect completes.
             task = asyncio.create_task(func())
+            pending_task.current = task
 
             # Wait for either the effect to complete or the stop signal
             stop_task = asyncio.create_task(stop.wait())
@@ -240,15 +244,17 @@ async def effect(stop: asyncio.Event) -> None:
 
             # If the effect completed first, store the cleanup function
             if task in done:
-                cleanup_func.current = task.result()
+                pending_task.current = None
+                with contextlib.suppress(asyncio.CancelledError):
+                    cleanup_func.current = task.result()
                 # Cancel the stop task since we don't need it anymore
                 stop_task.cancel()
                 with contextlib.suppress(asyncio.CancelledError):
                     await stop_task
                 # Now wait for the stop signal to run cleanup
                 await stop.wait()
+            # Stop signal came first - cancel the effect task
             else:
-                # Stop signal came first - cancel the effect task
                 task.cancel()
                 with contextlib.suppress(asyncio.CancelledError):
                     await task

tests/test_core/test_hooks.py

@@ -600,6 +600,41 @@ async def effect():
     event_that_never_occurs.set()
 
 
+async def test_async_effect_sleep_is_cancelled_on_re_render():
+    """Test that async effects waiting on asyncio.sleep are properly cancelled."""
+    component_hook = HookCatcher()
+    effect_ran = asyncio.Event()
+    effect_was_cancelled = asyncio.Event()
+
+    @reactpy.component
+    @component_hook.capture
+    def ComponentWithSleepEffect():
+        @reactpy.hooks.use_async_effect(dependencies=None)
+        async def effect():
+            effect_ran.set()
+            try:
+                await asyncio.sleep(1000)
+            except asyncio.CancelledError:
+                effect_was_cancelled.set()
+                raise
+
+        return reactpy.html.div()
+
+    async with Layout(ComponentWithSleepEffect()) as layout:
+        await layout.render()
+
+        # Wait for the effect to start
+        await effect_ran.wait()
+
+        # Trigger a re-render which should cancel the previous effect
+        component_hook.latest.schedule_render()
+        await layout.render()
+
+        # Verify the previous effect was cancelled
+        await asyncio.wait_for(effect_was_cancelled.wait(), 1)
+
+
+
 async def test_error_in_effect_is_gracefully_handled():
     @reactpy.component
     def ComponentWithEffect():
@@ -1349,3 +1384,40 @@ async def run_test():
                 await layout.render()
 
         asyncio.run(run_test())
+
+
+async def test_async_effect_cancelled_on_dependency_change():
+    """Test that async effects are cancelled when dependencies change."""
+    set_state = reactpy.Ref()
+    effect_ran = asyncio.Event()
+    effect_was_cancelled = asyncio.Event()
+
+    @reactpy.component
+    def ComponentWithDependentEffect():
+        state, set_state.current = reactpy.hooks.use_state(0)
+
+        @reactpy.hooks.use_async_effect(dependencies=[state])
+        async def effect():
+            effect_ran.set()
+            try:
+                await asyncio.sleep(1000)
+            except asyncio.CancelledError:
+                effect_was_cancelled.set()
+                raise
+
+        return reactpy.html.div()
+
+    async with Layout(ComponentWithDependentEffect()) as layout:
+        await layout.render()
+
+        # Wait for the effect to start
+        await effect_ran.wait()
+        effect_ran.clear()
+
+        # Change state to trigger effect cleanup/re-run
+        set_state.current(1)
+        await layout.render()
+
+        # Verify the previous effect was cancelled
+        await asyncio.wait_for(effect_was_cancelled.wait(), 1)
+