diff --git a/packages/core/src/global_context.rs b/packages/core/src/global_context.rs index 738bf64c3..d84f138be 100644 --- a/packages/core/src/global_context.rs +++ b/packages/core/src/global_context.rs @@ -50,9 +50,9 @@ pub fn provide_root_context(value: T) -> T { .expect("to be in a dioxus runtime") } -/// Suspends the current component -pub fn suspend() -> Option { - Runtime::with_current_scope(|cx| cx.suspend()); +/// Suspended the current component on a specific task and then return None +pub fn suspend(task: Task) -> Element { + Runtime::with_current_scope(|cx| cx.suspend(task)); None } diff --git a/packages/core/src/runtime.rs b/packages/core/src/runtime.rs index 355e4e9ac..dbfe125b2 100644 --- a/packages/core/src/runtime.rs +++ b/packages/core/src/runtime.rs @@ -1,3 +1,5 @@ +use rustc_hash::FxHashSet; + use crate::{ innerlude::{LocalTask, SchedulerMsg}, render_signal::RenderSignal, @@ -24,11 +26,14 @@ pub struct Runtime { // We use this to track the current task pub(crate) current_task: Cell>, - pub(crate) rendering: Cell, - /// Tasks created with cx.spawn pub(crate) tasks: RefCell>>, + // Currently suspended tasks + pub(crate) suspended_tasks: RefCell>, + + pub(crate) rendering: Cell, + pub(crate) sender: futures_channel::mpsc::UnboundedSender, // Synchronous tasks need to be run after the next render. The virtual dom stores a list of those tasks to send a signal to them when the next render is done. @@ -45,6 +50,7 @@ impl Runtime { scope_stack: Default::default(), current_task: Default::default(), tasks: Default::default(), + suspended_tasks: Default::default(), }) } diff --git a/packages/core/src/scope_arena.rs b/packages/core/src/scope_arena.rs index 5f5424143..46de9c184 100644 --- a/packages/core/src/scope_arena.rs +++ b/packages/core/src/scope_arena.rs @@ -41,7 +41,6 @@ impl VirtualDom { let new_nodes = { let context = scope.state(); - context.suspended.set(false); context.hook_index.set(0); // Run all pre-render hooks @@ -70,12 +69,11 @@ impl VirtualDom { self.dirty_scopes .remove(&ScopeOrder::new(context.height, scope_id)); - if context.suspended.get() { + if let Some(task) = context.last_suspendable_task.take() { if matches!(new_nodes, RenderReturn::Aborted(_)) { - self.suspended_scopes.insert(context.id); + tracing::trace!("Suspending {:?} on {:?}", scope_id, task); + self.runtime.suspended_tasks.borrow_mut().insert(task); } - } else if !self.suspended_scopes.is_empty() { - _ = self.suspended_scopes.remove(&context.id); } self.runtime.scope_stack.borrow_mut().pop(); diff --git a/packages/core/src/scope_context.rs b/packages/core/src/scope_context.rs index 9fa54e71c..0acde5055 100644 --- a/packages/core/src/scope_context.rs +++ b/packages/core/src/scope_context.rs @@ -16,13 +16,14 @@ pub(crate) struct Scope { pub(crate) parent_id: Option, pub(crate) height: u32, pub(crate) render_count: Cell, - pub(crate) suspended: Cell, // Note: the order of the hook and context fields is important. The hooks field must be dropped before the contexts field in case a hook drop implementation tries to access a context. pub(crate) hooks: RefCell>>, pub(crate) hook_index: Cell, pub(crate) shared_contexts: RefCell>>, pub(crate) spawned_tasks: RefCell>, + /// The task that was last spawned that may suspend. We use this task to check what task to suspend in the event of an early None return from a component + pub(crate) last_suspendable_task: Cell>, pub(crate) before_render: RefCell>>, pub(crate) after_render: RefCell>>, } @@ -40,9 +41,9 @@ impl Scope { parent_id, height, render_count: Cell::new(0), - suspended: Cell::new(false), shared_contexts: RefCell::new(vec![]), spawned_tasks: RefCell::new(FxHashSet::default()), + last_suspendable_task: Cell::new(None), hooks: RefCell::new(vec![]), hook_index: Cell::new(0), before_render: RefCell::new(vec![]), @@ -241,9 +242,9 @@ impl Scope { Runtime::with(|rt| rt.spawn(self.id, fut)).expect("Runtime to exist") } - /// Mark this component as suspended and then return None - pub fn suspend(&self) -> Option { - self.suspended.set(true); + /// Mark this component as suspended on a specific task and then return None + pub fn suspend(&self, task: Task) -> Option { + self.last_suspendable_task.set(Some(task)); None } @@ -340,10 +341,10 @@ impl ScopeId { .expect("to be in a dioxus runtime") } - /// Suspends the current component - pub fn suspend(self) -> Option { + /// Suspended a component on a specific task and then return None + pub fn suspend(self, task: Task) -> Option { Runtime::with_scope(self, |cx| { - cx.suspend(); + cx.suspend(task); }); None } diff --git a/packages/core/src/tasks.rs b/packages/core/src/tasks.rs index 39f353a4c..2335afac3 100644 --- a/packages/core/src/tasks.rs +++ b/packages/core/src/tasks.rs @@ -173,6 +173,9 @@ impl Runtime { // Remove it from the scheduler self.tasks.borrow_mut().try_remove(id.0); + + // Remove it from the suspended tasks + self.suspended_tasks.borrow_mut().remove(&id); } // Remove the scope from the stack @@ -187,6 +190,7 @@ impl Runtime { /// /// 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> { + self.suspended_tasks.borrow_mut().remove(&id); self.tasks.borrow_mut().try_remove(id.0) } } diff --git a/packages/core/src/virtual_dom.rs b/packages/core/src/virtual_dom.rs index db7cc349c..dec5b261a 100644 --- a/packages/core/src/virtual_dom.rs +++ b/packages/core/src/virtual_dom.rs @@ -18,7 +18,7 @@ use crate::{ AttributeValue, ComponentFunction, Element, Event, Mutations, }; use futures_util::StreamExt; -use rustc_hash::{FxHashMap, FxHashSet}; +use rustc_hash::FxHashMap; use slab::Slab; use std::{any::Any, rc::Rc}; use tracing::instrument; @@ -201,9 +201,6 @@ pub struct VirtualDom { pub(crate) runtime: Rc, - // Currently suspended scopes - pub(crate) suspended_scopes: FxHashSet, - rx: futures_channel::mpsc::UnboundedReceiver, } @@ -319,7 +316,6 @@ impl VirtualDom { queued_templates: Default::default(), elements: Default::default(), mounts: Default::default(), - suspended_scopes: Default::default(), }; let root = dom.new_scope(Box::new(root), "app"); @@ -448,13 +444,6 @@ impl VirtualDom { /// ``` #[instrument(skip(self), level = "trace", name = "VirtualDom::wait_for_work")] pub async fn wait_for_work(&mut self) { - // And then poll the futures - self.poll_tasks().await; - } - - /// Poll the scheduler for any work - #[instrument(skip(self), level = "trace", name = "VirtualDom::poll_tasks")] - async fn poll_tasks(&mut self) { loop { // Process all events - Scopes are marked dirty, etc // Sometimes when wakers fire we get a slew of updates at once, so its important that we drain this completely @@ -469,17 +458,22 @@ impl VirtualDom { let _runtime = RuntimeGuard::new(self.runtime.clone()); // There isn't any more work we can do synchronously. Wait for any new work to be ready - match self.rx.next().await.expect("channel should never close") { - SchedulerMsg::Immediate(id) => self.mark_dirty(id), - SchedulerMsg::TaskNotified(id) => { - // Instead of running the task immediately, we insert it into the runtime's task queue. - // The task may be marked dirty at the same time as the scope that owns the task is dropped. - self.mark_task_dirty(id); - } - }; + self.wait_for_event().await; } } + /// Wait for the next event to trigger and add it to the queue + async fn wait_for_event(&mut self) { + match self.rx.next().await.expect("channel should never close") { + SchedulerMsg::Immediate(id) => self.mark_dirty(id), + SchedulerMsg::TaskNotified(id) => { + // Instead of running the task immediately, we insert it into the runtime's task queue. + // The task may be marked dirty at the same time as the scope that owns the task is dropped. + self.mark_task_dirty(id); + } + }; + } + /// Queue any pending events fn queue_events(&mut self) { // Prevent a task from deadlocking the runtime by repeatedly queueing itself @@ -494,7 +488,6 @@ impl VirtualDom { /// Process all events in the queue until there are no more left #[instrument(skip(self), level = "trace", name = "VirtualDom::process_events")] pub fn process_events(&mut self) { - let _runtime = RuntimeGuard::new(self.runtime.clone()); self.queue_events(); // Now that we have collected all queued work, we should check if we have any dirty scopes. If there are not, then we can poll any queued futures @@ -502,6 +495,14 @@ impl VirtualDom { return; } + self.poll_tasks() + } + + /// Poll any queued tasks + #[instrument(skip(self), level = "trace", name = "VirtualDom::poll_tasks")] + fn poll_tasks(&mut self) { + // Make sure we set the runtime since we're running user code + let _runtime = RuntimeGuard::new(self.runtime.clone()); // Next, run any queued tasks // We choose not to poll the deadline since we complete pretty quickly anyways while let Some(task) = self.dirty_scopes.pop_task() { @@ -617,7 +618,6 @@ impl VirtualDom { { let _runtime = RuntimeGuard::new(self.runtime.clone()); // Then, poll any tasks that might be pending in the scope - // This will run effects, so this **must** be done after the scope is diffed for task in work.tasks { let _ = self.runtime.handle_task_wakeup(task); } @@ -649,15 +649,75 @@ impl VirtualDom { #[instrument(skip(self), level = "trace", name = "VirtualDom::wait_for_suspense")] pub async fn wait_for_suspense(&mut self) { loop { - if self.suspended_scopes.is_empty() { + if self.runtime.suspended_tasks.borrow().is_empty() { break; } // Wait for a work to be ready (IE new suspense leaves to pop up) - self.poll_tasks().await; + 'wait_for_work: loop { + // Process all events - Scopes are marked dirty, etc + // Sometimes when wakers fire we get a slew of updates at once, so its important that we drain this completely + self.queue_events(); + + // Now that we have collected all queued work, we should check if we have any dirty scopes. If there are not, then we can poll any queued futures + if self.dirty_scopes.has_dirty_scopes() { + println!("dirty scopes"); + break; + } + + { + // Make sure we set the runtime since we're running user code + let _runtime = RuntimeGuard::new(self.runtime.clone()); + // Next, run any queued tasks + // We choose not to poll the deadline since we complete pretty quickly anyways + while let Some(task) = self.dirty_scopes.pop_task() { + // If the scope doesn't exist for whatever reason, then we should skip it + if !self.scopes.contains(task.order.id.0) { + continue; + } + + // Then poll any tasks that might be pending + let tasks = task.tasks_queued.into_inner(); + for task in tasks { + if self.runtime.suspended_tasks.borrow().contains(&task) { + let _ = self.runtime.handle_task_wakeup(task); + // Running that task, may mark a scope higher up as dirty. If it does, return from the function early + self.queue_events(); + if self.dirty_scopes.has_dirty_scopes() { + break 'wait_for_work; + } + } + } + } + } + + self.wait_for_event().await; + } // Render whatever work needs to be rendered, unlocking new futures and suspense leaves - self.render_immediate(&mut NoOpMutations); + while let Some(work) = self.dirty_scopes.pop_work() { + // If the scope doesn't exist for whatever reason, then we should skip it + if !self.scopes.contains(work.scope.id.0) { + continue; + } + + { + let _runtime = RuntimeGuard::new(self.runtime.clone()); + // Then, poll any tasks that might be pending in the scope + for task in work.tasks { + // During suspense, we only want to run tasks that are suspended + if self.runtime.suspended_tasks.borrow().contains(&task) { + let _ = self.runtime.handle_task_wakeup(task); + } + } + // If the scope is dirty, run the scope and get the mutations + if work.rerun_scope { + let new_nodes = self.run_scope(work.scope.id); + + self.diff_scope(&mut NoOpMutations, work.scope.id, new_nodes); + } + } + } } } diff --git a/packages/core/tests/suspense.rs b/packages/core/tests/suspense.rs index 2957bc0a8..56de18c92 100644 --- a/packages/core/tests/suspense.rs +++ b/packages/core/tests/suspense.rs @@ -1,9 +1,10 @@ use dioxus::prelude::*; +use std::future::poll_fn; +use std::task::Poll; #[test] fn suspense_resolves() { // wait just a moment, not enough time for the boundary to resolve - tokio::runtime::Builder::new_current_thread() .build() .unwrap() @@ -31,11 +32,35 @@ fn app() -> Element { fn suspended_child() -> Element { let mut val = use_signal(|| 0); + // Tasks that are not suspended should never be polled + spawn(async move { + panic!("Non-suspended task was polled"); + }); + + // // Memos should still work like normal + // let memo = use_memo(move || val * 2); + // assert_eq!(memo, val * 2); + if val() < 3 { - spawn(async move { + let task = spawn(async move { + // Poll each task 3 times + let mut count = 0; + poll_fn(|cx| { + println!("polling... {}", count); + if count < 3 { + count += 1; + cx.waker().wake_by_ref(); + Poll::Pending + } else { + Poll::Ready(()) + } + }) + .await; + + println!("waiting... {}", val); val += 1; }); - suspend()?; + suspend(task)?; } rsx!("child")