From d28221bdc2a9aefe910b9bc7289b52314ca32f19 Mon Sep 17 00:00:00 2001
From: Evan Almloff <evanalmloff@gmail.com>
Date: Wed, 17 Jan 2024 08:36:39 -0600
Subject: [PATCH] unqueue tasks when a scope is dropped

---
 packages/core/src/arena.rs       |  9 ++-------
 packages/core/src/runtime.rs     |  5 +++++
 packages/core/src/tasks.rs       | 12 +++++++++++-
 packages/core/src/virtual_dom.rs | 13 +++----------
 4 files changed, 21 insertions(+), 18 deletions(-)

diff --git a/packages/core/src/arena.rs b/packages/core/src/arena.rs
index ec2412efb..ac680299f 100644
--- a/packages/core/src/arena.rs
+++ b/packages/core/src/arena.rs
@@ -68,18 +68,13 @@ impl VirtualDom {
     //
     // Note: This will not remove any ids from the arena
     pub(crate) fn drop_scope(&mut self, id: ScopeId) {
-        let (height, spawned_tasks) = {
+        let height = {
             let scope = self.scopes.remove(id.0);
             let context = scope.context();
-            let spawned_tasks = context.spawned_tasks.borrow();
-            let spawned_tasks: Vec<_> = spawned_tasks.iter().copied().collect();
-            (context.height, spawned_tasks)
+            context.height
         };
 
         self.dirty_scopes.remove(&DirtyScope { height, id });
-        for task in spawned_tasks {
-            self.runtime.remove_task(task);
-        }
     }
 }
 
diff --git a/packages/core/src/runtime.rs b/packages/core/src/runtime.rs
index 5cb8dea48..764d54f57 100644
--- a/packages/core/src/runtime.rs
+++ b/packages/core/src/runtime.rs
@@ -6,6 +6,7 @@ use crate::{
 };
 use std::{
     cell::{Cell, Ref, RefCell},
+    collections::VecDeque,
     rc::Rc,
 };
 
@@ -28,6 +29,9 @@ pub struct Runtime {
     /// Tasks created with cx.spawn
     pub(crate) tasks: RefCell<slab::Slab<LocalTask>>,
 
+    /// Queued tasks that are waiting to be polled
+    pub(crate) queued_tasks: Rc<RefCell<VecDeque<Task>>>,
+
     pub(crate) sender: futures_channel::mpsc::UnboundedSender<SchedulerMsg>,
 }
 
@@ -40,6 +44,7 @@ impl Runtime {
             scope_stack: Default::default(),
             current_task: Default::default(),
             tasks: Default::default(),
+            queued_tasks: Rc::new(RefCell::new(VecDeque::new())),
         })
     }
 
diff --git a/packages/core/src/tasks.rs b/packages/core/src/tasks.rs
index 950f9fc19..09ef4d6d6 100644
--- a/packages/core/src/tasks.rs
+++ b/packages/core/src/tasks.rs
@@ -112,11 +112,21 @@ impl Runtime {
         self.current_task.set(None);
     }
 
+    /// Take a queued task from the scheduler
+    pub(crate) fn take_queued_task(&self) -> Option<Task> {
+        self.queued_tasks.borrow_mut().pop_front()
+    }
+
     /// Drop the future with the given TaskId
     ///
     /// This does not abort the task, so you'll want to wrap it in an abort handle if that's important to you
     pub(crate) fn remove_task(&self, id: Task) -> Option<LocalTask> {
-        self.tasks.borrow_mut().try_remove(id.0)
+        let task = self.tasks.borrow_mut().try_remove(id.0);
+
+        // Remove the task from the queued tasks so we don't poll a different task with the same id
+        self.queued_tasks.borrow_mut().retain(|t| *t != id);
+
+        task
     }
 
     /// Get the currently running task
diff --git a/packages/core/src/virtual_dom.rs b/packages/core/src/virtual_dom.rs
index b660fd281..243375a36 100644
--- a/packages/core/src/virtual_dom.rs
+++ b/packages/core/src/virtual_dom.rs
@@ -19,12 +19,7 @@ use crate::{
 use futures_util::{pin_mut, StreamExt};
 use rustc_hash::{FxHashMap, FxHashSet};
 use slab::Slab;
-use std::{
-    any::Any,
-    collections::{BTreeSet, VecDeque},
-    future::Future,
-    rc::Rc,
-};
+use std::{any::Any, collections::BTreeSet, future::Future, rc::Rc};
 
 /// A virtual node system that progresses user events and diffs UI trees.
 ///
@@ -189,7 +184,6 @@ pub struct VirtualDom {
     pub(crate) scopes: Slab<ScopeState>,
 
     pub(crate) dirty_scopes: BTreeSet<DirtyScope>,
-    pub(crate) dirty_tasks: VecDeque<Task>,
 
     // Maps a template path to a map of byte indexes to templates
     pub(crate) templates: FxHashMap<TemplateId, FxHashMap<usize, Template>>,
@@ -284,7 +278,6 @@ impl VirtualDom {
             runtime: Runtime::new(tx),
             scopes: Default::default(),
             dirty_scopes: Default::default(),
-            dirty_tasks: Default::default(),
             templates: Default::default(),
             queued_templates: Default::default(),
             elements: Default::default(),
@@ -420,7 +413,7 @@ impl VirtualDom {
 
                     if !has_dirty_scopes {
                         // If we have no dirty scopes, then we should poll any tasks that have been notified
-                        while let Some(task) = self.dirty_tasks.pop_front() {
+                        while let Some(task) = self.runtime.take_queued_task() {
                             self.handle_task_wakeup(task);
                         }
                     }
@@ -626,7 +619,7 @@ impl VirtualDom {
 
     /// Queue a task to be polled after all dirty scopes have been rendered
     fn queue_task_wakeup(&mut self, id: Task) {
-        self.dirty_tasks.push_back(id);
+        self.runtime.queued_tasks.borrow_mut().push_back(id);
     }
 
     /// Handle notifications by tasks inside the scheduler