Fix event handler memory leak (#2298)

* fix event handler memory leak and double drops

* Prevent double dropping generational boxes

* recycle instead of dropping

packages/core-macro/tests/values_memoize_in_place.rs

@@ -2,12 +2,41 @@ use dioxus::prelude::*;
 use dioxus_core::ElementId;
 use std::rc::Rc;
 
-thread_local! {
-    static DROP_COUNT: std::cell::RefCell<usize> = const { std::cell::RefCell::new(0) };
-}
-
 #[test]
 fn values_memoize_in_place() {
+    thread_local! {
+        static DROP_COUNT: std::cell::RefCell<usize> = const { std::cell::RefCell::new(0) };
+    }
+
+    struct CountsDrop;
+
+    impl Drop for CountsDrop {
+        fn drop(&mut self) {
+            DROP_COUNT.with(|c| *c.borrow_mut() += 1);
+        }
+    }
+
+    fn app() -> Element {
+        let mut count = use_signal(|| 0);
+        let x = CountsDrop;
+
+        if generation() < 15 {
+            count += 1;
+        }
+
+        rsx! {
+            TakesEventHandler {
+                click: move |num| {
+                    // Force the closure to own the drop counter
+                    let _ = &x;
+                    println!("num is {num}");
+                },
+                children: count() / 2
+            }
+            TakesSignal { sig: count(), children: count() / 2 }
+        }
+    }
+
     set_event_converter(Box::new(dioxus::html::SerializedHtmlEventConverter));
     let mut dom = VirtualDom::new(app);
 
@@ -29,33 +58,42 @@ fn values_memoize_in_place() {
     assert_eq!(drop_count, 15);
 }
 
-struct CountsDrop;
+// We move over event handlers in place. Make sure we do that in a way that doesn't destroy the original event handler
+#[test]
+fn cloning_event_handler_components_work() {
+    fn app() -> Element {
+        let rsx_with_event_handler_component = rsx! {
+            TakesEventHandler {
+                click: move |evt| {
+                    println!("Clicked {evt:?}!");
+                }
+            }
+        };
 
-impl Drop for CountsDrop {
-    fn drop(&mut self) {
-        DROP_COUNT.with(|c| *c.borrow_mut() += 1);
+        rsx! {
+            {rsx_with_event_handler_component.clone()}
+            {rsx_with_event_handler_component.clone()}
+            {rsx_with_event_handler_component.clone()}
+            {rsx_with_event_handler_component}
+        }
     }
-}
-
-fn app() -> Element {
-    let mut count = use_signal(|| 0);
-    let x = CountsDrop;
 
-    if generation() < 15 {
-        count += 1;
-    }
+    set_event_converter(Box::new(dioxus::html::SerializedHtmlEventConverter));
+    let mut dom = VirtualDom::new(app);
 
-    rsx! {
-        TakesEventHandler {
-            click: move |num| {
-                // Force the closure to own the drop counter
-                let _ = &x;
-                println!("num is {num}");
-            },
-            children: count() / 2
-        }
-        TakesSignal { sig: count(), children: count() / 2 }
+    let mutations = dom.rebuild_to_vec();
+    println!("{:#?}", mutations);
+    dom.mark_dirty(ScopeId::ROOT);
+    for _ in 0..20 {
+        dom.handle_event(
+            "click",
+            Rc::new(PlatformEventData::new(Box::<SerializedMouseData>::default())),
+            ElementId(1),
+            true,
+        );
+        dom.render_immediate(&mut dioxus_core::NoOpMutations);
     }
+    dom.render_immediate(&mut dioxus_core::NoOpMutations);
 }
 
 #[component]

packages/core/src/events.rs

@@ -1,6 +1,9 @@
 use crate::{global_context::current_scope_id, Runtime, ScopeId};
 use generational_box::GenerationalBox;
-use std::{cell::Cell, rc::Rc};
+use std::{
+    cell::{Cell, RefCell},
+    rc::Rc,
+};
 
 /// A wrapper around some generic data that handles the event's state
 ///
@@ -163,9 +166,34 @@ impl<T: std::fmt::Debug> std::fmt::Debug for Event<T> {
 /// ```
 pub struct EventHandler<T = ()> {
     pub(crate) origin: ScopeId,
+    // During diffing components with EventHandler, we move the EventHandler over in place instead of rerunning the child component.
+    // ```rust
+    // #[component]
+    // fn Child(onclick: EventHandler<MouseEvent>) -> Element {
+    //     rsx!{
+    //         button {
+    //             // Diffing Child will not rerun this component, it will just update the EventHandler in place so that if this callback is called, it will run the latest version of the callback
+    //             onclick: move |evt| cx.onclick.call(evt),
+    //         }
+    //     }
+    /// }
+    /// ```
+    // This is both more efficient and allows us to avoid out of date EventHandlers.
+    //
+    // We double box here because we want the data to be copy (GenerationalBox) and still update in place (ExternalListenerCallback)
+    // This isn't an ideal solution for performance, but it is non-breaking and fixes the issues described in https://github.com/DioxusLabs/dioxus/pull/2298
     pub(super) callback: GenerationalBox<Option<ExternalListenerCallback<T>>>,
 }
 
+impl<T> std::fmt::Debug for EventHandler<T> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("EventHandler")
+            .field("origin", &self.origin)
+            .field("callback", &self.callback)
+            .finish()
+    }
+}
+
 impl<T: 'static> Default for EventHandler<T> {
     fn default() -> Self {
         EventHandler::new(|_| {})
@@ -192,15 +220,29 @@ impl<T: 'static> PartialEq for EventHandler<T> {
     }
 }
 
-type ExternalListenerCallback<T> = Box<dyn FnMut(T)>;
+type ExternalListenerCallback<T> = Rc<RefCell<dyn FnMut(T)>>;
 
 impl<T: 'static> EventHandler<T> {
-    /// Create a new [`EventHandler`] from an [`FnMut`]
+    /// Create a new [`EventHandler`] from an [`FnMut`]. The callback is owned by the current scope and will be dropped when the scope is dropped.
+    /// This should not be called directly in the body of a component because it will not be dropped until the component is dropped.
     #[track_caller]
     pub fn new(mut f: impl FnMut(T) + 'static) -> EventHandler<T> {
-        let callback = GenerationalBox::leak(Some(Box::new(move |event: T| {
+        let owner = crate::innerlude::current_owner::<generational_box::UnsyncStorage>();
+        let callback = owner.insert(Some(Rc::new(RefCell::new(move |event: T| {
+            f(event);
+        })) as Rc<RefCell<dyn FnMut(T)>>));
+        EventHandler {
+            callback,
+            origin: current_scope_id().expect("to be in a dioxus runtime"),
+        }
+    }
+
+    /// Leak a new [`EventHandler`] that will not be dropped unless it is manually dropped.
+    #[track_caller]
+    pub fn leak(mut f: impl FnMut(T) + 'static) -> EventHandler<T> {
+        let callback = GenerationalBox::leak(Some(Rc::new(RefCell::new(move |event: T| {
             f(event);
-        }) as Box<dyn FnMut(T)>));
+        })) as Rc<RefCell<dyn FnMut(T)>>));
         EventHandler {
             callback,
             origin: current_scope_id().expect("to be in a dioxus runtime"),
@@ -211,9 +253,12 @@ impl<T: 'static> EventHandler<T> {
     ///
     /// This borrows the event using a RefCell. Recursively calling a listener will cause a panic.
     pub fn call(&self, event: T) {
-        if let Some(callback) = self.callback.write().as_mut() {
+        if let Some(callback) = self.callback.read().as_ref() {
             Runtime::with(|rt| rt.scope_stack.borrow_mut().push(self.origin));
-            callback(event);
+            {
+                let mut callback = callback.borrow_mut();
+                callback(event);
+            }
             Runtime::with(|rt| rt.scope_stack.borrow_mut().pop());
         }
     }
@@ -227,16 +272,16 @@ impl<T: 'static> EventHandler<T> {
 
     #[doc(hidden)]
     /// This should only be used by the `rsx!` macro.
-    pub fn __set(&mut self, value: impl FnMut(T) + 'static) {
-        self.callback.set(Some(Box::new(value)));
+    pub fn __set(&mut self, value: ExternalListenerCallback<T>) {
+        self.callback.set(Some(value));
     }
 
     #[doc(hidden)]
     /// This should only be used by the `rsx!` macro.
     pub fn __take(&self) -> ExternalListenerCallback<T> {
         self.callback
-            .manually_drop()
-            .expect("Signal has already been dropped")
+            .read()
+            .clone()
             .expect("EventHandler was manually dropped")
     }
 }

packages/core/src/nodes.rs

@@ -737,7 +737,7 @@ impl AttributeValue {
     pub fn listener<T: 'static>(mut callback: impl FnMut(Event<T>) + 'static) -> AttributeValue {
         // TODO: maybe don't use the copy-variant of EventHandler here?
         // Maybe, create an Owned variant so we are less likely to run into leaks
-        AttributeValue::Listener(EventHandler::new(move |event: Event<dyn Any>| {
+        AttributeValue::Listener(EventHandler::leak(move |event: Event<dyn Any>| {
             let data = event.data.downcast::<T>().unwrap();
             callback(Event {
                 propagates: event.propagates,
@@ -761,7 +761,7 @@ impl std::fmt::Debug for AttributeValue {
             Self::Float(arg0) => f.debug_tuple("Float").field(arg0).finish(),
             Self::Int(arg0) => f.debug_tuple("Int").field(arg0).finish(),
             Self::Bool(arg0) => f.debug_tuple("Bool").field(arg0).finish(),
-            Self::Listener(_) => f.debug_tuple("Listener").finish(),
+            Self::Listener(listener) => f.debug_tuple("Listener").field(listener).finish(),
             Self::Any(_) => f.debug_tuple("Any").finish(),
             Self::None => write!(f, "None"),
         }

packages/generational-box/src/lib.rs

@@ -200,15 +200,19 @@ impl<T, S: Storage<T>> GenerationalBox<T, S> {
 
     /// Recycle the generationalbox, dropping the value.
     pub fn recycle(&self) {
-        _ = Storage::take(&self.raw.0.data);
-        <S as AnyStorage>::recycle(&self.raw);
+        if self.validate() {
+            <S as AnyStorage>::recycle(&self.raw);
+        }
     }
 
     /// Drop the value out of the generational box and invalidate the generational box.
     /// This will return the value if the value was taken.
     pub fn manually_drop(&self) -> Option<T> {
         if self.validate() {
-            Storage::take(&self.raw.0.data)
+            // TODO: Next breaking release we should change the take method to automatically recycle the box
+            let value = Storage::take(&self.raw.0.data)?;
+            <S as AnyStorage>::recycle(&self.raw);
+            Some(value)
         } else {
             None
         }
@@ -371,9 +375,10 @@ impl<S> MemoryLocation<S> {
     where
         S: AnyStorage,
     {
-        let old = self.0.data.manually_drop();
+        self.0.data.manually_drop();
+
         #[cfg(any(debug_assertions, feature = "check_generation"))]
-        if old {
+        {
             let new_generation = self.0.generation.load(std::sync::atomic::Ordering::Relaxed) + 1;
             self.0
                 .generation
@@ -402,14 +407,52 @@ impl<S> MemoryLocation<S> {
     }
 }
 
+// We track the generation along with the memory location so that when generational boxes are dropped early, we don't end up dropping the new occupant of the slot
+struct LocationKey<S: 'static> {
+    #[cfg(any(debug_assertions, feature = "check_generation"))]
+    generation: u32,
+    location: MemoryLocation<S>,
+}
+
+impl<S: AnyStorage> LocationKey<S> {
+    fn exists(&self) -> bool {
+        #[cfg(any(debug_assertions, feature = "check_generation"))]
+        return self.generation
+            == self
+                .location
+                .0
+                .generation
+                .load(std::sync::atomic::Ordering::Relaxed);
+        #[cfg(not(any(debug_assertions, feature = "check_generation")))]
+        return true;
+    }
+
+    fn drop(self) {
+        // If this is the same box we own, we can just drop it
+        if self.exists() {
+            S::recycle(&self.location);
+        }
+    }
+}
+
+impl<T, S: AnyStorage> From<GenerationalBox<T, S>> for LocationKey<S> {
+    fn from(value: GenerationalBox<T, S>) -> Self {
+        Self {
+            #[cfg(any(debug_assertions, feature = "check_generation"))]
+            generation: value.generation,
+            location: value.raw,
+        }
+    }
+}
+
 struct OwnerInner<S: AnyStorage + 'static> {
-    owned: Vec<MemoryLocation<S>>,
+    owned: Vec<LocationKey<S>>,
 }
 
 impl<S: AnyStorage> Drop for OwnerInner<S> {
     fn drop(&mut self) {
-        for location in &self.owned {
-            S::recycle(location)
+        for key in self.owned.drain(..) {
+            key.drop();
         }
     }
 }
@@ -459,7 +502,7 @@ impl<S: AnyStorage> Owner<S> {
             #[cfg(any(debug_assertions, feature = "debug_borrows"))]
             caller,
         );
-        self.0.lock().owned.push(location);
+        self.0.lock().owned.push(key.into());
         key
     }
 
@@ -477,7 +520,7 @@ impl<S: AnyStorage> Owner<S> {
             created_at: std::panic::Location::caller(),
             _marker: PhantomData,
         };
-        self.0.lock().owned.push(location);
+        self.0.lock().owned.push(generational_box.into());
         generational_box
     }
 }

packages/generational-box/src/sync.rs

@@ -100,8 +100,9 @@ impl AnyStorage for SyncStorage {
     }
 
     fn recycle(location: &MemoryLocation<Self>) {
+        let location = *location;
         location.drop();
-        sync_runtime().lock().push(*location);
+        sync_runtime().lock().push(location);
     }
 }
 

packages/generational-box/src/unsync.rs

@@ -163,7 +163,8 @@ impl AnyStorage for UnsyncStorage {
     }
 
     fn recycle(location: &MemoryLocation<Self>) {
+        let location = *location;
         location.drop();
-        UNSYNC_RUNTIME.with(|runtime| runtime.borrow_mut().push(*location));
+        UNSYNC_RUNTIME.with(|runtime| runtime.borrow_mut().push(location));
     }
 }

packages/router-macro/src/segment.rs

@@ -135,7 +135,7 @@ pub fn static_segment_idx(idx: usize) -> Ident {
 
 pub fn parse_route_segments<'a>(
     route_span: Span,
-    mut fields: impl Iterator<Item = (&'a Ident, &'a Type)> + Clone,
+    fields: impl Iterator<Item = (&'a Ident, &'a Type)> + Clone,
     route: &str,
 ) -> syn::Result<(
     Vec<RouteSegment>,