Merge pull request #1710 from Exotik850/cx-spawn-poll

Make cx.spawn poll the task before returning `TaskId`

examples/compose.rs

@@ -27,13 +27,7 @@ fn app(cx: Scope) -> Element {
 
             button {
                 onclick: move |_| {
-                    let dom = VirtualDom::new_with_props(compose, ComposeProps {
-                        app_tx: tx.clone()
-                    });
-
-                    // this returns a weak reference to the other window
-                    // Be careful not to keep a strong reference to the other window or it will never be dropped
-                    // and the window will never close.
+                    let dom = VirtualDom::new_with_props(compose, ComposeProps { app_tx: tx.clone() });
                     window.new_window(dom, Default::default());
                 },
                 "Click to compose a new email"

packages/core/src/scheduler/task.rs

@@ -19,7 +19,7 @@ pub struct TaskId(pub usize);
 /// the task itself is the waker
 pub(crate) struct LocalTask {
     pub scope: ScopeId,
-    pub(super) task: RefCell<Pin<Box<dyn Future<Output = ()> + 'static>>>,
+    pub task: RefCell<Pin<Box<dyn Future<Output = ()> + 'static>>>,
     pub waker: Waker,
 }
 
@@ -48,11 +48,15 @@ impl Scheduler {
             })),
         };
 
-        entry.insert(task);
+        let mut cx = std::task::Context::from_waker(&task.waker);
 
-        self.sender
-            .unbounded_send(SchedulerMsg::TaskNotified(task_id))
-            .expect("Scheduler should exist");
+        if !task.task.borrow_mut().as_mut().poll(&mut cx).is_ready() {
+            self.sender
+                .unbounded_send(SchedulerMsg::TaskNotified(task_id))
+                .expect("Scheduler should exist");
+        }
+
+        entry.insert(task);
 
         task_id
     }
@@ -60,8 +64,8 @@ impl Scheduler {
     /// Drop the future with the given TaskId
     ///
     /// This does not abort the task, so you'll want to wrap it in an aborthandle if that's important to you
-    pub fn remove(&self, id: TaskId) {
-        self.tasks.borrow_mut().try_remove(id.0);
+    pub fn remove(&self, id: TaskId) -> Option<LocalTask> {
+        self.tasks.borrow_mut().try_remove(id.0)
     }
 }
 

packages/hooks/src/use_effect.rs

@@ -99,9 +99,8 @@ where
         // Create the new future
         let return_value = future(dependencies.out());
 
-        if let Some(task) = return_value.apply(state.cleanup.clone(), cx) {
-            state.task.set(Some(task));
-        }
+        let task = return_value.apply(state.cleanup.clone(), cx);
+        state.task.set(Some(task));
     }
 }
 
@@ -109,15 +108,15 @@ 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>;
+    fn apply(self, oncleanup: UseEffectCleanup, cx: &ScopeState) -> TaskId;
 }
 
 impl<T> UseEffectReturn<()> for T
 where
     T: Future<Output = ()> + 'static,
 {
-    fn apply(self, _: UseEffectCleanup, cx: &ScopeState) -> Option<TaskId> {
-        Some(cx.push_future(self))
+    fn apply(self, _: UseEffectCleanup, cx: &ScopeState) -> TaskId {
+        cx.push_future(self)
     }
 }
 
@@ -128,12 +127,11 @@ 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 {
+    fn apply(self, oncleanup: UseEffectCleanup, cx: &ScopeState) -> TaskId {
+        cx.push_future(async move {
             let cleanup = self.await;
             *oncleanup.borrow_mut() = Some(Box::new(cleanup) as Box<dyn FnOnce()>);
-        });
-        Some(task)
+        })
     }
 }