From 326ef09dc4c0091caaf1d0dc3e1a31925258d87f Mon Sep 17 00:00:00 2001
From: Ryan Morshead <ryan.morshead@gmail.com>
Date: Tue, 4 Jul 2023 17:01:15 -0600
Subject: [PATCH] effects with cancel scopes

---
 src/py/reactpy/reactpy/core/hooks.py         | 19 +++++----
 src/py/reactpy/tests/test_core/test_hooks.py | 44 ++++++++++++++++++++
 2 files changed, 56 insertions(+), 7 deletions(-)

diff --git a/src/py/reactpy/reactpy/core/hooks.py b/src/py/reactpy/reactpy/core/hooks.py
index a8334458b..5c6575be0 100644
--- a/src/py/reactpy/reactpy/core/hooks.py
+++ b/src/py/reactpy/reactpy/core/hooks.py
@@ -16,6 +16,7 @@
     overload,
 )
 
+from anyio import create_task_group
 from typing_extensions import TypeAlias
 
 from reactpy.config import REACTPY_DEBUG_MODE
@@ -147,15 +148,19 @@ def add_effect(function: _EffectApplyFunc) -> None:
             async_function = cast(_AsyncEffectFunc, function)
 
             def sync_function() -> _EffectCleanFunc | None:
-                future = asyncio.ensure_future(async_function())
+                tg = create_task_group()
 
-                def clean_future() -> None:
-                    if not future.cancel():
-                        clean = future.result()
-                        if clean is not None:
-                            clean()
+                async def wrapper() -> None:
+                    async with tg:
+                        tg.start_soon(async_function)
 
-                return clean_future
+                # TODO: Component unmount should block until all effect tasks exit
+                # Right now we don't have a good way to do this, so we just cancel
+                # the task and hope the user's effect doesn't misbehave and ignore
+                # cancellation.
+                asyncio.create_task(wrapper())  # noqa: RUF006
+
+                return tg.cancel_scope.cancel  # type: ignore
 
         def effect() -> None:
             if last_clean_callback.current is not None:
diff --git a/src/py/reactpy/tests/test_core/test_hooks.py b/src/py/reactpy/tests/test_core/test_hooks.py
index 453d07c99..cee2a860d 100644
--- a/src/py/reactpy/tests/test_core/test_hooks.py
+++ b/src/py/reactpy/tests/test_core/test_hooks.py
@@ -1,6 +1,7 @@
 import asyncio
 
 import pytest
+from anyio import CancelScope
 
 import reactpy
 from reactpy import html
@@ -10,6 +11,7 @@
     LifeCycleHook,
     current_hook,
     strictly_equal,
+    use_effect,
 )
 from reactpy.core.layout import Layout
 from reactpy.testing import DisplayFixture, HookCatcher, assert_reactpy_did_log, poll
@@ -1257,3 +1259,45 @@ def bad_effect():
             await layout.render()
             component_hook.latest.schedule_render()
             await layout.render()  # no error
+
+
+async def test_async_effect_can_be_protected_from_cancellation_with_cancel_scope():
+    component_hook = HookCatcher()
+
+    cancel_initiated = asyncio.Event()
+    cancel_completed = asyncio.Event()
+    cancel_scope_complete = asyncio.Event()
+
+    @reactpy.component
+    @component_hook.capture
+    def ComponentWithEffect():
+        @use_effect(dependencies=None)  # run on every render
+        async def effect_with_shield():
+            async with CancelScope(shield=True):
+                await cancel_initiated.wait()
+                cancel_scope_complete.set()
+            try:
+                await asyncio.sleep(0)
+            except asyncio.CancelledError:
+                cancel_completed.set()
+                raise
+
+        return reactpy.html.div()
+
+    async with reactpy.Layout(ComponentWithEffect()) as layout:
+        await layout.render()
+
+        # sanity check that nothing has happened yet
+        assert not cancel_scope_complete.is_set()
+        assert not cancel_completed.is_set()
+
+        # schedule render that will trigger effect cancellation
+        component_hook.latest.schedule_render()
+
+        # wait for render to ensure cancellation has been initiated
+        await layout.render()
+        cancel_initiated.set()
+
+        # wait for effect to complete
+        await asyncio.wait_for(cancel_scope_complete.wait(), timeout=1)
+        await asyncio.wait_for(cancel_completed.wait(), timeout=1)