Fast renderer: clean up removed components before new effects run

reacton/core.py

@@ -1205,6 +1205,11 @@ class ComponentContext:
     # the render phase skipped this whole subtree (nothing changed), so the
     # reconciliation phase can reuse the previous result without walking
     clean_subtree: bool = False
+    # Fast renderer only: stale component effect cleanups are run at the same
+    # point where the default renderer removes stale elements, while widget
+    # closing stays deferred to the fast renderer's stale sweep.
+    fast_stale_effect_keys: Optional[List[str]] = None
+    fast_stale_effects_cleaned: Set[str] = field(default_factory=set)
 
     # elements created in this context go there
     owns: Set[Element] = field(default_factory=set)
@@ -1263,6 +1268,8 @@ def _teardown_component_context(context: ComponentContext):
     context.exceptions_self = []
     context.exceptions_children = []
     context.context_managers = []
+    context.fast_stale_effect_keys = None
+    context.fast_stale_effects_cleaned = set()
 
 
 @dataclass
@@ -2523,6 +2530,8 @@ def _render(self, element: Element, default_key: str, parent_key: str):
             # the root element of a component determines which keys are in use,
             # everything else is stale and gets removed during reconciliation
             context.used_keys.clear()
+            context.fast_stale_effect_keys = None
+            context.fast_stale_effects_cleaned.clear()
 
         el = element
         key = el._key
@@ -2923,6 +2932,7 @@ def _reconsolidate(self, el: Element, default_key: str, parent_key: str):
                 self._shared_elements.add(el)
                 assert el in self._shared_elements_next
                 self._shared_elements_next.remove(el)
+            self._cleanup_stale_effects_for_context(context, parent_key)
 
     def _process_effects(self, child_context: "ComponentContext", context: "ComponentContext"):
         # NOTE: effect/cleanup exceptions are recorded on the context of the
@@ -2965,6 +2975,64 @@ def _store_widget(self, context: "ComponentContext", el: Element, key: str, widg
         else:
             context.widgets[key] = widget
 
+    def _cleanup_stale_effects_for_context(self, context: "ComponentContext", parent_key: str):
+        if context.fast_stale_effect_keys is None:
+            # reversed, so we can pop from the end and still go in sorted order
+            context.fast_stale_effect_keys = sorted(set(context.elements) - context.used_keys, reverse=True)
+        context_prev = self.context
+        try:
+            while context.fast_stale_effect_keys:
+                stale_key = context.fast_stale_effect_keys.pop()
+                if stale_key not in context.elements or stale_key in context.fast_stale_effects_cleaned:
+                    continue
+                self.context = context
+                self._cleanup_stale_effects(context.elements[stale_key], stale_key, parent_key)
+        finally:
+            self.context = context_prev
+
+    def _cleanup_stale_effects(self, el: Element, default_key: str, parent_key: str):
+        key = el._key
+        if key is None:
+            key = default_key
+        assert key is not None
+        context = self.context
+        assert context is not None
+        if key in context.fast_stale_effects_cleaned:
+            return
+        if el.is_shared and (el in self._shared_elements_next or el not in self._shared_elements):
+            return
+        context.fast_stale_effects_cleaned.add(key)
+
+        if isinstance(el.component, ComponentFunction):
+            if el.is_shared:
+                self._visit_children(el, key, parent_key, self._cleanup_stale_effects)
+            child_context = context.children.get(key)
+            if child_context is None:
+                return
+            try:
+                self.context = child_context
+                child_context.exceptions_self = []
+                child_context.exceptions_children = []
+                for effect in child_context.effects:
+                    try:
+                        if not effect._cleaned_up:
+                            effect.cleanup()
+                    except BaseException as e:
+                        effect._cleaned_up = True
+                        logger.exception("Effect cleanup %r raised exception %r", effect.callable, e)
+                        child_context.exceptions_self.append(e)
+                        self._set_rerender_needed("Exception ocurred during effect")
+                        _mark_needs_render_ancestors(child_context)
+                assert child_context.root_element is not None
+                self._cleanup_stale_effects(child_context.root_element, "/", parent_key=join_key(parent_key, key))
+            finally:
+                self.context = context
+            if child_context.exceptions_self or child_context.exceptions_children and not child_context.exception_handler:
+                context.exceptions_children.extend(child_context.exceptions_self)
+                context.exceptions_children.extend(child_context.exceptions_children)
+        else:
+            self._visit_children(el, key, parent_key, self._cleanup_stale_effects)
+
     def _remove_element(self, el: Element, default_key: str, parent_key):
         key = el._key
         if key is None:

reacton/core_test.py

@@ -4088,3 +4088,102 @@ def cleanup():
         assert box.children[0].value == "value 1"
     finally:
         rc.close()
+
+
+def test_stale_component_cleanup_runs_before_new_sibling_effect():
+    active = None
+    set_show_old = lambda x: None  # noqa
+
+    def owner(name):
+        def effect():
+            nonlocal active
+            active = name
+
+            def cleanup():
+                nonlocal active
+                active = None
+
+            return cleanup
+
+        return effect
+
+    @react.component
+    def Old():
+        react.use_effect(owner("old"), [])
+        return w.Label(value="old")
+
+    @react.component
+    def New():
+        react.use_effect(owner("new"), [])
+        return w.Label(value="new")
+
+    @react.component
+    def App():
+        nonlocal set_show_old
+        show_old, set_show_old = react.use_state(True)
+        children = [w.Label(value="first").key("first")]
+        if show_old:
+            children.append(Old().key("old"))
+        else:
+            children.append(New().key("new"))
+        return w.VBox(children=children)
+
+    box, rc = react.render(App(), handle_error=False)
+    assert active == "old"
+    set_show_old(False)
+    assert active == "new"
+    rc.close()
+
+
+def test_stale_component_cleanup_order_with_nested_removed_components():
+    log = []
+    set_show_old = lambda x: None  # noqa
+
+    def logger(name):
+        def effect():
+            log.append(f"{name} run")
+
+            def cleanup():
+                log.append(f"{name} cleanup")
+
+            return cleanup
+
+        return effect
+
+    @react.component
+    def Inner(name):
+        react.use_effect(logger(f"{name}-inner"), [])
+        return w.Label(value=f"{name}-inner")
+
+    @react.component
+    def Old(name):
+        react.use_effect(logger(f"{name}-outer"), [])
+        return w.VBox(children=[Inner(name=name)])
+
+    @react.component
+    def New():
+        react.use_effect(logger("new"), [])
+        return w.Label(value="new")
+
+    @react.component
+    def App():
+        nonlocal set_show_old
+        show_old, set_show_old = react.use_state(True)
+        children = [w.Label(value="first").key("first")]
+        if show_old:
+            children.extend([Old(name="z").key("z"), Old(name="a").key("a")])
+        else:
+            children.append(New().key("new"))
+        return w.VBox(children=children)
+
+    box, rc = react.render(App(), handle_error=False)
+    log.clear()
+    set_show_old(False)
+    assert log == [
+        "a-outer cleanup",
+        "a-inner cleanup",
+        "z-outer cleanup",
+        "z-inner cleanup",
+        "new run",
+    ]
+    rc.close()

reacton/fuzz_test.py

@@ -81,6 +81,12 @@ def effect():
             if state % 5 == 4:
                 set_state(state + 1)
 
+            def cleanup():
+                # both renderers must also run cleanups in the same order
+                log.append(("cleanup", id, state))
+
+            return cleanup
+
         react.use_effect(effect, [state])
         children: List[Any] = [w.Label(value=f"node {id} {state}")]
         for i in range(rnd.randint(0, 4)):