From 1332b82dc8b556b1fb593ba8acedac7a20f7b2eb Mon Sep 17 00:00:00 2001 From: Jonathan Kelley Date: Mon, 15 Jan 2024 17:32:49 -0800 Subject: [PATCH] Expose current task --- packages/core/src/runtime.rs | 28 ++++++++++++++++++------ packages/core/src/scheduler/mod.rs | 18 ---------------- packages/core/src/scheduler/task.rs | 13 +++++++++--- packages/core/src/scheduler/wait.rs | 2 +- packages/core/src/scope_arena.rs | 3 +-- packages/core/src/scope_context.rs | 33 ++++++++++++++++------------- packages/core/src/virtual_dom.rs | 8 +++---- 7 files changed, 55 insertions(+), 50 deletions(-) diff --git a/packages/core/src/runtime.rs b/packages/core/src/runtime.rs index 69d83c6fc..ecff4b199 100644 --- a/packages/core/src/runtime.rs +++ b/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>>, - pub(crate) scheduler: Rc, // We use this to track the current scope pub(crate) scope_stack: RefCell>, + + // We use this to track the current task + pub(crate) current_task: Cell>, + pub(crate) rendering: Cell, + + /// Tasks created with cx.spawn + pub tasks: RefCell>, + + pub sender: futures_channel::mpsc::UnboundedSender, } impl Runtime { - pub(crate) fn new(scheduler: Rc) -> Rc { + pub(crate) fn new(tx: futures_channel::mpsc::UnboundedSender) -> Rc { 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(), }) } diff --git a/packages/core/src/scheduler/mod.rs b/packages/core/src/scheduler/mod.rs index 0446ae91b..c6cd0681d 100644 --- a/packages/core/src/scheduler/mod.rs +++ b/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, - - /// Tasks created with cx.spawn - pub tasks: RefCell>, -} - -impl Scheduler { - pub fn new(sender: futures_channel::mpsc::UnboundedSender) -> Rc { - Rc::new(Scheduler { - sender, - tasks: RefCell::new(Slab::new()), - }) - } -} diff --git a/packages/core/src/scheduler/task.rs b/packages/core/src/scheduler/task.rs index c6223e5ea..6cddf8043 100644 --- a/packages/core/src/scheduler/task.rs +++ b/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, pub task: RefCell + '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 { self.tasks.borrow_mut().try_remove(id.0) } + + /// Get the currently running task + pub fn current_task(&self) -> Option { + self.current_task.get() + } } pub struct LocalTaskHandle { diff --git a/packages/core/src/scheduler/wait.rs b/packages/core/src/scheduler/wait.rs index 9d6969847..45b9c3791 100644 --- a/packages/core/src/scheduler/wait.rs +++ b/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, diff --git a/packages/core/src/scope_arena.rs b/packages/core/src/scope_arena.rs index 2647c230e..fe66adb8a 100644 --- a/packages/core/src/scope_arena.rs +++ b/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 diff --git a/packages/core/src/scope_context.rs b/packages/core/src/scope_context.rs index 56e1cdbd6..659a08323 100644 --- a/packages/core/src/scope_context.rs +++ b/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>>, pub(crate) hook_index: Cell, - pub(crate) tasks: Rc, pub(crate) spawned_tasks: RefCell>, } @@ -41,7 +40,6 @@ impl ScopeContext { id: ScopeId, parent_id: Option, height: u32, - tasks: Rc, ) -> 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 { + 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 { - 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 { - 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 + '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 + '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") } } diff --git a/packages/core/src/virtual_dom.rs b/packages/core/src/virtual_dom.rs index 53b5528e4..21dffbefa 100644 --- a/packages/core/src/virtual_dom.rs +++ b/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(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(),