Merge pull request #2639 from Giovanni-Tably/dispose-fix

fix: ensure everything is disposed of consistently

leptos_reactive/src/node.rs

@@ -13,10 +13,7 @@ pub struct Disposer(pub(crate) NodeId);
 impl Drop for Disposer {
     fn drop(&mut self) {
         let id = self.0;
-        _ = with_runtime(|runtime| {
-            runtime.cleanup_node(id);
-            runtime.dispose_node(id);
-        });
+        _ = with_runtime(|runtime| runtime.dispose_node(id));
     }
 }
 

leptos_reactive/src/runtime.rs

@@ -190,27 +190,6 @@ impl Runtime {
         self.mark_clean(node_id);
     }
 
-    pub(crate) fn cleanup_node(&self, node_id: NodeId) {
-        // first, run our cleanups, if any
-        let c = { self.on_cleanups.borrow_mut().remove(node_id) };
-        // untrack around all cleanups
-        let prev_observer = self.observer.take();
-        if let Some(cleanups) = c {
-            for cleanup in cleanups {
-                cleanup();
-            }
-        }
-        self.observer.set(prev_observer);
-
-        // dispose of any of our properties
-        let properties = { self.node_properties.borrow_mut().remove(node_id) };
-        if let Some(properties) = properties {
-            for property in properties {
-                self.cleanup_property(property);
-            }
-        }
-    }
-
     pub(crate) fn update(&self, node_id: NodeId) {
         let node = {
             let nodes = self.nodes.borrow();
@@ -254,44 +233,45 @@ impl Runtime {
         }
     }
 
-    pub(crate) fn cleanup_property(&self, property: ScopeProperty) {
+    pub(crate) fn dispose_node(&self, node_id: NodeId) {
+        self.cleanup_node(node_id);
+
+        // each of the subs needs to remove the node from its dependencies
+        // so that it doesn't try to read the (now disposed) signal
+        let subs = self.node_subscribers.borrow_mut().remove(node_id);
+        if let Some(subs) = subs {
+            let source_map = self.node_sources.borrow();
+            for effect in subs.borrow().iter() {
+                if let Some(effect_sources) = source_map.get(*effect) {
+                    effect_sources.borrow_mut().swap_remove(&node_id);
+                }
+            }
+        }
+
+        self.node_sources.borrow_mut().remove(node_id);
+        let node = { self.nodes.borrow_mut().remove(node_id) };
+        drop(node);
+    }
+    fn cleanup_node(&self, node_id: NodeId) {
+        self.run_on_cleanups(node_id);
+        self.dispose_children(node_id);
+    }
+    /// Dispose of all of the children of the node recursively and completely.
+    fn dispose_children(&self, node_id: NodeId) {
+        let properties = { self.node_properties.borrow_mut().remove(node_id) };
+        if let Some(properties) = properties {
+            for property in properties {
+                self.dispose_property(property);
+            }
+        }
+    }
+    fn dispose_property(&self, property: ScopeProperty) {
         // for signals, triggers, memos, effects, shared node cleanup
         match property {
             ScopeProperty::Signal(node)
             | ScopeProperty::Trigger(node)
             | ScopeProperty::Effect(node) => {
-                // run all cleanups for this node
-                let cleanups = { self.on_cleanups.borrow_mut().remove(node) };
-                for cleanup in cleanups.into_iter().flatten() {
-                    cleanup();
-                }
-
-                // clean up all children
-                let properties =
-                    { self.node_properties.borrow_mut().remove(node) };
-                for property in properties.into_iter().flatten() {
-                    self.cleanup_property(property);
-                }
-
-                // each of the subs needs to remove the node from its dependencies
-                // so that it doesn't try to read the (now disposed) signal
-                let subs = self.node_subscribers.borrow_mut().remove(node);
-
-                if let Some(subs) = subs {
-                    let source_map = self.node_sources.borrow();
-                    for effect in subs.borrow().iter() {
-                        if let Some(effect_sources) = source_map.get(*effect) {
-                            effect_sources.borrow_mut().swap_remove(&node);
-                        }
-                    }
-                }
-
-                // no longer needs to track its sources
-                self.node_sources.borrow_mut().remove(node);
-
-                // remove the node from the graph
-                let node = { self.nodes.borrow_mut().remove(node) };
-                drop(node);
+                self.dispose_node(node);
             }
             ScopeProperty::Resource(id) => {
                 let value = self.resources.borrow_mut().remove(id);
@@ -303,6 +283,16 @@ impl Runtime {
             }
         }
     }
+    fn run_on_cleanups(&self, node_id: NodeId) {
+        let c = { self.on_cleanups.borrow_mut().remove(node_id) };
+        let prev_observer = self.observer.take(); // untrack around all cleanups
+        if let Some(cleanups) = c {
+            for cleanup in cleanups {
+                cleanup();
+            }
+        }
+        self.observer.set(prev_observer);
+    }
 
     pub(crate) fn cleanup_sources(&self, node_id: NodeId) {
         let sources = self.node_sources.borrow();
@@ -502,12 +492,6 @@ impl Runtime {
         }
     }
 
-    pub(crate) fn dispose_node(&self, node: NodeId) {
-        self.node_sources.borrow_mut().remove(node);
-        self.node_subscribers.borrow_mut().remove(node);
-        self.nodes.borrow_mut().remove(node);
-    }
-
     #[track_caller]
     pub(crate) fn register_property(
         &self,
@@ -1191,11 +1175,7 @@ impl RuntimeId {
         );
 
         (id, move || {
-            with_runtime(|runtime| {
-                runtime.nodes.borrow_mut().remove(id);
-                runtime.node_sources.borrow_mut().remove(id);
-            })
-            .expect(
+            with_runtime(|runtime| runtime.dispose_node(id)).expect(
                 "tried to stop a watch in a runtime that has been disposed",
             );
         })

leptos_reactive/tests/cleanup.rs

@@ -43,3 +43,66 @@ fn cleanup() {
 
     runtime.dispose();
 }
+
+#[test]
+fn cleanup_on_dispose() {
+    use leptos_reactive::{
+        create_memo, create_runtime, create_trigger, on_cleanup, SignalDispose,
+        SignalGetUntracked,
+    };
+
+    struct ExecuteOnDrop(Option<Box<dyn FnOnce()>>);
+    impl ExecuteOnDrop {
+        fn new(f: impl FnOnce() + 'static) -> Self {
+            Self(Some(Box::new(f)))
+        }
+    }
+    impl Drop for ExecuteOnDrop {
+        fn drop(&mut self) {
+            self.0.take().unwrap()();
+        }
+    }
+
+    let runtime = create_runtime();
+
+    let trigger = create_trigger();
+
+    println!("STARTING");
+
+    let memo = create_memo(move |_| {
+        trigger.track();
+
+        // An example of why you might want to do this is that
+        // when something goes out of reactive scope you want it to be cleaned up.
+        // The cleaning up might have side effects, and those side effects might cause
+        // re-renders where new `on_cleanup` are registered.
+        let on_drop = ExecuteOnDrop::new(|| {
+            on_cleanup(|| println!("Nested cleanup in progress."))
+        });
+
+        on_cleanup(move || {
+            println!("Cleanup in progress.");
+            drop(on_drop)
+        });
+    });
+    println!("Memo 1: {:?}", memo);
+    let _ = memo.get_untracked(); // First cleanup registered.
+
+    memo.dispose(); // Cleanup not run here.
+
+    println!("Cleanup should have been executed.");
+
+    let memo = create_memo(move |_| {
+        // New cleanup registered. It'll panic here.
+        on_cleanup(move || println!("Test passed."));
+    });
+    println!("Memo 2: {:?}", memo);
+    println!("^ Note how the memos have the same key (different versions).");
+    let _ = memo.get_untracked(); // First cleanup registered.
+
+    println!("Test passed.");
+
+    memo.dispose();
+
+    runtime.dispose();
+}

leptos_reactive/tests/memo.rs

@@ -293,3 +293,29 @@ fn owning_memo_slice() {
 
     runtime.dispose();
 }
+
+#[test]
+fn leak_on_dispose() {
+    use std::rc::Rc;
+
+    let runtime = create_runtime();
+
+    let trigger = create_trigger();
+
+    let value = Rc::new(());
+    let weak = Rc::downgrade(&value);
+
+    let memo = create_memo(move |_| {
+        trigger.track();
+
+        create_rw_signal(value.clone());
+    });
+
+    memo.get_untracked();
+
+    memo.dispose();
+
+    assert!(weak.upgrade().is_none()); // Should have been dropped.
+
+    runtime.dispose();
+}