Merge pull request #1650 from ealmloff/effect-cleanup

Add an optional cleanup closure to the use effect hook

packages/hooks/src/use_effect.rs

@@ -1,5 +1,10 @@
 use dioxus_core::{ScopeState, TaskId};
-use std::{any::Any, cell::Cell, future::Future};
+use std::{
+    any::Any,
+    cell::{Cell, RefCell},
+    future::Future,
+    rc::Rc,
+};
 
 use crate::UseFutureDep;
 
@@ -14,7 +19,7 @@ use crate::UseFutureDep;
 /// ## Arguments
 ///
 /// - `dependencies`: a tuple of references to values that are `PartialEq` + `Clone`.
-/// - `future`: a closure that takes the `dependencies` as arguments and returns a `'static` future.
+/// - `future`: a closure that takes the `dependencies` as arguments and returns a `'static` future. That future may return nothing or a closure that will be executed when the dependencies change to clean up the effect.
 ///
 /// ## Examples
 ///
@@ -33,6 +38,16 @@ use crate::UseFutureDep;
 ///         }
 ///     });
 ///
+///     // Only fetch the user data when the id changes.
+///     use_effect(cx, (id,), |(id,)| {
+///         to_owned![name];
+///         async move {
+///             let user = fetch_user(id).await;
+///             name.set(user.name);
+///             move || println!("Cleaning up from {}", id)
+///         }
+///     });
+///
 ///     let name = name.get().clone().unwrap_or("Loading...".to_string());
 ///
 ///     render!(
@@ -45,34 +60,80 @@ use crate::UseFutureDep;
 ///     render!(Profile { id: 0 })
 /// }
 /// ```
-pub fn use_effect<T, F, D>(cx: &ScopeState, dependencies: D, future: impl FnOnce(D::Out) -> F)
+pub fn use_effect<T, R, D>(cx: &ScopeState, dependencies: D, future: impl FnOnce(D::Out) -> R)
 where
-    T: 'static,
-    F: Future<Output = T> + 'static,
     D: UseFutureDep,
+    R: UseEffectReturn<T>,
 {
     struct UseEffect {
         needs_regen: bool,
         task: Cell<Option<TaskId>>,
         dependencies: Vec<Box<dyn Any>>,
+        cleanup: UseEffectCleanup,
+    }
+
+    impl Drop for UseEffect {
+        fn drop(&mut self) {
+            if let Some(cleanup) = self.cleanup.borrow_mut().take() {
+                cleanup();
+            }
+        }
     }
 
     let state = cx.use_hook(move || UseEffect {
         needs_regen: true,
         task: Cell::new(None),
         dependencies: Vec::new(),
+        cleanup: Rc::new(RefCell::new(None)),
     });
 
     if dependencies.clone().apply(&mut state.dependencies) || state.needs_regen {
+        // Call the cleanup function if it exists
+        if let Some(cleanup) = state.cleanup.borrow_mut().take() {
+            cleanup();
+        }
+
         // We don't need regen anymore
         state.needs_regen = false;
 
         // Create the new future
-        let fut = future(dependencies.out());
+        let return_value = future(dependencies.out());
 
-        state.task.set(Some(cx.push_future(async move {
-            fut.await;
-        })));
+        if let Some(task) = return_value.apply(state.cleanup.clone(), cx) {
+            state.task.set(Some(task));
+        }
+    }
+}
+
+type UseEffectCleanup = Rc<RefCell<Option<Box<dyn FnOnce()>>>>;
+
+/// Something that can be returned from a `use_effect` hook.
+pub trait UseEffectReturn<T> {
+    fn apply(self, oncleanup: UseEffectCleanup, cx: &ScopeState) -> Option<TaskId>;
+}
+
+impl<T> UseEffectReturn<()> for T
+where
+    T: Future<Output = ()> + 'static,
+{
+    fn apply(self, _: UseEffectCleanup, cx: &ScopeState) -> Option<TaskId> {
+        Some(cx.push_future(self))
+    }
+}
+
+#[doc(hidden)]
+pub struct CleanupFutureMarker;
+impl<T, F> UseEffectReturn<CleanupFutureMarker> for T
+where
+    T: Future<Output = F> + 'static,
+    F: FnOnce() + 'static,
+{
+    fn apply(self, oncleanup: UseEffectCleanup, cx: &ScopeState) -> Option<TaskId> {
+        let task = cx.push_future(async move {
+            let cleanup = self.await;
+            *oncleanup.borrow_mut() = Some(Box::new(cleanup) as Box<dyn FnOnce()>);
+        });
+        Some(task)
     }
 }