feat: tasks are dropped when components are unmounted

packages/core/src/scopes.rs

@@ -68,7 +68,12 @@ impl ScopeArena {
             heuristics: RefCell::new(FxHashMap::default()),
             free_scopes: RefCell::new(Vec::new()),
             nodes: RefCell::new(nodes),
-            tasks: TaskQueue::new(sender),
+            tasks: Rc::new(TaskQueue {
+                tasks: RefCell::new(FxHashMap::default()),
+                task_map: RefCell::new(FxHashMap::default()),
+                gen: Cell::new(0),
+                sender,
+            }),
         }
     }
 
@@ -179,8 +184,14 @@ impl ScopeArena {
         log::trace!("removing scope {:?}", id);
         self.ensure_drop_safety(id);
 
-        //
-        let tasks = self.tasks.tasks.borrow_mut();
+        // Dispose of any ongoing tasks
+        let mut tasks = self.tasks.tasks.borrow_mut();
+        let mut task_map = self.tasks.task_map.borrow_mut();
+        if let Some(cur_tasks) = task_map.remove(&id) {
+            for task in cur_tasks {
+                tasks.remove(&task);
+            }
+        }
 
         // Safety:
         // - ensure_drop_safety ensures that no references to this scope are in use
@@ -438,7 +449,13 @@ pub struct ScopeId(pub usize);
 /// once a Task has been completed.
 #[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
 #[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
-pub struct TaskId(pub usize);
+pub struct TaskId {
+    /// The global ID of the task
+    pub id: usize,
+
+    /// The original scope that this task was scheduled in
+    pub scope: ScopeId,
+}
 
 /// Every component in Dioxus is represented by a `ScopeState`.
 ///
@@ -748,7 +765,7 @@ impl ScopeState {
             .unbounded_send(SchedulerMsg::NewTask(self.our_arena_idx))
             .unwrap();
 
-        self.tasks.push_fut(fut)
+        self.tasks.spawn(self.our_arena_idx, fut)
     }
 
     /// Spawns the future but does not return the TaskId
@@ -759,7 +776,7 @@ impl ScopeState {
     /// Informs the scheduler that this task is no longer needed and should be removed
     /// on next poll.
     pub fn remove_future(&self, id: TaskId) {
-        self.tasks.remove_fut(id);
+        self.tasks.remove(id);
     }
 
     /// Take a lazy VNode structure and actually build it with the context of the VDom's efficient VNode allocator.
@@ -950,28 +967,36 @@ pub(crate) struct TaskQueue {
 
 pub(crate) type InnerTask = Pin<Box<dyn Future<Output = ()>>>;
 impl TaskQueue {
-    fn new(sender: UnboundedSender<SchedulerMsg>) -> Rc<Self> {
-        Rc::new(Self {
-            tasks: RefCell::new(FxHashMap::default()),
-            task_map: RefCell::new(FxHashMap::default()),
-            gen: Cell::new(0),
-            sender,
-        })
-    }
-    fn push_fut(&self, task: impl Future<Output = ()> + 'static) -> TaskId {
+    fn spawn(&self, scope: ScopeId, task: impl Future<Output = ()> + 'static) -> TaskId {
         let pinned = Box::pin(task);
         let id = self.gen.get();
         self.gen.set(id + 1);
-        let tid = TaskId(id);
+        let tid = TaskId { id, scope };
 
         self.tasks.borrow_mut().insert(tid, pinned);
+
+        // also add to the task map
+        // when the component is unmounted we know to remove it from the map
+        self.task_map
+            .borrow_mut()
+            .entry(scope)
+            .or_default()
+            .insert(tid);
+
         tid
     }
-    fn remove_fut(&self, id: TaskId) {
+
+    fn remove(&self, id: TaskId) {
         if let Ok(mut tasks) = self.tasks.try_borrow_mut() {
             let _ = tasks.remove(&id);
         }
+
+        // the task map is still around, but it'll be removed when the scope is unmounted
+        if let Some(task_map) = self.task_map.borrow_mut().get_mut(&id.scope) {
+            task_map.remove(&id);
+        }
     }
+
     pub(crate) fn has_tasks(&self) -> bool {
         !self.tasks.borrow().is_empty()
     }