Expose current task

packages/core/src/runtime.rs

@@ -1,6 +1,13 @@
 use std::cell::{Cell, Ref, RefCell};
 
-use crate::{innerlude::Scheduler, scope_context::ScopeContext, scopes::ScopeId};
+use slab::Slab;
+
+use crate::{
+    innerlude::{LocalTask, SchedulerMsg},
+    scope_context::ScopeContext,
+    scopes::ScopeId,
+    Task,
+};
 use std::rc::Rc;
 
 thread_local! {
@@ -52,23 +59,30 @@ where
 /// A global runtime that is shared across all scopes that provides the async runtime and context API
 pub struct Runtime {
     pub(crate) scope_contexts: RefCell<Vec<Option<ScopeContext>>>,
-    pub(crate) scheduler: Rc<Scheduler>,
 
     // We use this to track the current scope
     pub(crate) scope_stack: RefCell<Vec<ScopeId>>,
+
+    // We use this to track the current task
+    pub(crate) current_task: Cell<Option<Task>>,
+
     pub(crate) rendering: Cell<bool>,
+
+    /// Tasks created with cx.spawn
+    pub tasks: RefCell<Slab<LocalTask>>,
+
+    pub sender: futures_channel::mpsc::UnboundedSender<SchedulerMsg>,
 }
 
 impl Runtime {
-    pub(crate) fn new(scheduler: Rc<Scheduler>) -> Rc<Self> {
+    pub(crate) fn new(tx: futures_channel::mpsc::UnboundedSender<SchedulerMsg>) -> Rc<Self> {
         Rc::new(Self {
-            scheduler,
-
             scope_contexts: Default::default(),
-
             scope_stack: Default::default(),
-
+            current_task: Default::default(),
             rendering: Cell::new(true),
+            sender: tx,
+            tasks: Default::default(),
         })
     }
 

packages/core/src/scheduler/mod.rs

@@ -17,21 +17,3 @@ pub(crate) enum SchedulerMsg {
     /// A task has woken and needs to be progressed
     TaskNotified(Task),
 }
-
-use std::{cell::RefCell, rc::Rc};
-
-pub(crate) struct Scheduler {
-    pub sender: futures_channel::mpsc::UnboundedSender<SchedulerMsg>,
-
-    /// Tasks created with cx.spawn
-    pub tasks: RefCell<Slab<LocalTask>>,
-}
-
-impl Scheduler {
-    pub fn new(sender: futures_channel::mpsc::UnboundedSender<SchedulerMsg>) -> Rc<Self> {
-        Rc::new(Scheduler {
-            sender,
-            tasks: RefCell::new(Slab::new()),
-        })
-    }
-}

packages/core/src/scheduler/task.rs

@@ -1,7 +1,7 @@
 use futures_util::task::ArcWake;
 
-use super::{Scheduler, SchedulerMsg};
-use crate::innerlude::{push_future, remove_future};
+use super::SchedulerMsg;
+use crate::innerlude::{push_future, remove_future, Runtime};
 use crate::ScopeId;
 use std::cell::RefCell;
 use std::future::Future;
@@ -42,11 +42,12 @@ impl Task {
 /// the task itself is the waker
 pub(crate) struct LocalTask {
     pub scope: ScopeId,
+    pub parent: Option<Task>,
     pub task: RefCell<Pin<Box<dyn Future<Output = ()> + 'static>>>,
     pub waker: Waker,
 }
 
-impl Scheduler {
+impl Runtime {
     /// Start a new future on the same thread as the rest of the VirtualDom.
     ///
     /// This future will not contribute to suspense resolving, so you should primarily use this for reacting to changes
@@ -63,6 +64,7 @@ impl Scheduler {
         let task_id = Task(entry.key());
 
         let task = LocalTask {
+            parent: self.current_task(),
             task: RefCell::new(Box::pin(task)),
             scope,
             waker: futures_util::task::waker(Arc::new(LocalTaskHandle {
@@ -90,6 +92,11 @@ impl Scheduler {
     pub fn remove(&self, id: Task) -> Option<LocalTask> {
         self.tasks.borrow_mut().try_remove(id.0)
     }
+
+    /// Get the currently running task
+    pub fn current_task(&self) -> Option<Task> {
+        self.current_task.get()
+    }
 }
 
 pub struct LocalTaskHandle {

packages/core/src/scheduler/wait.rs

@@ -8,7 +8,7 @@ impl VirtualDom {
     /// queue
     pub(crate) fn handle_task_wakeup(&mut self, id: Task) {
         let _runtime = RuntimeGuard::new(self.runtime.clone());
-        let mut tasks = self.runtime.scheduler.tasks.borrow_mut();
+        let mut tasks = self.runtime.tasks.borrow_mut();
 
         let task = match tasks.get(id.0) {
             Some(task) => task,

packages/core/src/scope_arena.rs

@@ -24,8 +24,7 @@ impl VirtualDom {
             last_rendered_node: Default::default(),
         });
 
-        let context =
-            ScopeContext::new(name, id, parent_id, height, self.runtime.scheduler.clone());
+        let context = ScopeContext::new(name, id, parent_id, height);
         self.runtime.create_context_at(id, context);
 
         scope

packages/core/src/scope_context.rs

@@ -1,5 +1,5 @@
 use crate::{
-    innerlude::{Scheduler, SchedulerMsg},
+    innerlude::SchedulerMsg,
     runtime::{with_runtime, with_scope},
     Element, ScopeId, Task,
 };
@@ -31,7 +31,6 @@ pub(crate) struct ScopeContext {
     pub(crate) hooks: RefCell<Vec<Box<dyn Any>>>,
     pub(crate) hook_index: Cell<usize>,
 
-    pub(crate) tasks: Rc<Scheduler>,
     pub(crate) spawned_tasks: RefCell<FxHashSet<Task>>,
 }
 
@@ -41,7 +40,6 @@ impl ScopeContext {
         id: ScopeId,
         parent_id: Option<ScopeId>,
         height: u32,
-        tasks: Rc<Scheduler>,
     ) -> Self {
         Self {
             name,
@@ -51,7 +49,6 @@ impl ScopeContext {
             render_count: Cell::new(0),
             suspended: Cell::new(false),
             shared_contexts: RefCell::new(vec![]),
-            tasks,
             spawned_tasks: RefCell::new(FxHashSet::default()),
             hooks: RefCell::new(vec![]),
             hook_index: Cell::new(0),
@@ -62,10 +59,13 @@ impl ScopeContext {
         self.parent_id
     }
 
+    fn sender(&self) -> futures_channel::mpsc::UnboundedSender<SchedulerMsg> {
+        with_runtime(|rt| rt.sender.clone()).unwrap()
+    }
+
     /// Mark this scope as dirty, and schedule a render for it.
     pub fn needs_update(&self) {
-        self.tasks
-            .sender
+        self.sender()
             .unbounded_send(SchedulerMsg::Immediate(self.id))
             .expect("Scheduler to exist if scope exists");
     }
@@ -74,7 +74,7 @@ impl ScopeContext {
     ///
     /// ## Notice: you should prefer using [`Self::schedule_update_any`] and [`Self::scope_id`]
     pub fn schedule_update(&self) -> Arc<dyn Fn() + Send + Sync + 'static> {
-        let (chan, id) = (self.tasks.sender.clone(), self.id);
+        let (chan, id) = (self.sender(), self.id);
         Arc::new(move || drop(chan.unbounded_send(SchedulerMsg::Immediate(id))))
     }
 
@@ -84,7 +84,7 @@ impl ScopeContext {
     ///
     /// This method should be used when you want to schedule an update for a component
     pub fn schedule_update_any(&self) -> Arc<dyn Fn(ScopeId) + Send + Sync> {
-        let chan = self.tasks.sender.clone();
+        let chan = self.sender();
         Arc::new(move |id| {
             chan.unbounded_send(SchedulerMsg::Immediate(id)).unwrap();
         })
@@ -228,7 +228,7 @@ impl ScopeContext {
 
     /// Pushes the future onto the poll queue to be polled after the component renders.
     pub fn push_future(&self, fut: impl Future<Output = ()> + 'static) -> Task {
-        let id = self.tasks.spawn(self.id, fut);
+        let id = with_runtime(|rt| rt.spawn(self.id, fut)).expect("Runtime to exist");
         self.spawned_tasks.borrow_mut().insert(id);
         id
     }
@@ -243,14 +243,14 @@ impl ScopeContext {
     /// This is good for tasks that need to be run after the component has been dropped.
     pub fn spawn_forever(&self, fut: impl Future<Output = ()> + 'static) -> Task {
         // The root scope will never be unmounted so we can just add the task at the top of the app
-        self.tasks.spawn(ScopeId::ROOT, fut)
+        with_runtime(|rt| rt.spawn(ScopeId::ROOT, fut)).expect("Runtime to exist")
     }
 
     /// Informs the scheduler that this task is no longer needed and should be removed.
     ///
     /// This drops the task immediately.
     pub fn remove_future(&self, id: Task) {
-        self.tasks.remove(id);
+        with_runtime(|rt| rt.remove(id)).expect("Runtime to exist");
     }
 
     /// Mark this component as suspended and then return None
@@ -314,10 +314,13 @@ impl ScopeContext {
 
 impl Drop for ScopeContext {
     fn drop(&mut self) {
-        // Drop all spawned tasks
-        for id in self.spawned_tasks.borrow().iter() {
-            self.tasks.remove(*id);
-        }
+        with_runtime(|rt| {
+            // Drop all spawned tasks
+            for id in self.spawned_tasks.borrow().iter() {
+                rt.remove(*id);
+            }
+        })
+        .expect("Runtime to exist")
     }
 }
 

packages/core/src/virtual_dom.rs

@@ -6,8 +6,8 @@ use crate::{
     any_props::{BoxedAnyProps, VProps},
     arena::ElementId,
     innerlude::{
-        DirtyScope, ElementRef, ErrorBoundary, NoOpMutations, Scheduler, SchedulerMsg, ScopeState,
-        VNodeMount, WriteMutations,
+        DirtyScope, ElementRef, ErrorBoundary, NoOpMutations, SchedulerMsg, ScopeState, VNodeMount,
+        WriteMutations,
     },
     nodes::RenderReturn,
     nodes::{Template, TemplateId},
@@ -269,10 +269,10 @@ impl VirtualDom {
     /// ```
     pub fn new_with_props<P: Clone + 'static>(root: fn(P) -> Element, root_props: P) -> Self {
         let (tx, rx) = futures_channel::mpsc::unbounded();
-        let scheduler = Scheduler::new(tx);
+
         let mut dom = Self {
             rx,
-            runtime: Runtime::new(scheduler),
+            runtime: Runtime::new(tx),
             scopes: Default::default(),
             dirty_scopes: Default::default(),
             templates: Default::default(),