fix: use futuresunordered instead of bespoke waker system

This commit is contained in:
Jonathan Kelley committed 2022-12-29 19:53:43 -05:00
1 parent ff697d41e3
commit bd720e2268
12 files changed
+308 -272

No files matched your search

+4 -5
View File
@@ -1,14 +1,13 @@
use crate::ScopeId;
use futures_util::stream::FuturesUnordered;
use slab::Slab;
mod suspense;
mod task;
mod wait;
mod waker;
pub use suspense::*;
pub use task::*;
pub use waker::ArcWake;
/// The type of message that can be sent to the scheduler.
///
@@ -31,17 +30,17 @@ pub(crate) struct Scheduler {
pub sender: futures_channel::mpsc::UnboundedSender<SchedulerMsg>,
/// Tasks created with cx.spawn
pub tasks: RefCell<Slab<Arc<LocalTask>>>,
pub tasks: RefCell<FuturesUnordered<LocalTask>>,
/// Async components
pub leaves: RefCell<Slab<Arc<SuspenseLeaf>>>,
pub leaves: RefCell<Slab<SuspenseLeaf>>,
}
impl Scheduler {
pub fn new(sender: futures_channel::mpsc::UnboundedSender<SchedulerMsg>) -> Rc<Self> {
Rc::new(Scheduler {
sender,
tasks: RefCell::new(Slab::new()),
tasks: RefCell::new(FuturesUnordered::new()),
leaves: RefCell::new(Slab::new()),
})
}
+10 -3
View File
@@ -1,4 +1,6 @@
use super::{waker::ArcWake, SchedulerMsg};
use futures_util::task::ArcWake;
use super::SchedulerMsg;
use crate::ElementId;
use crate::{innerlude::Mutations, Element, ScopeId};
use std::future::Future;
@@ -40,11 +42,16 @@ pub(crate) struct SuspenseLeaf {
pub(crate) tx: futures_channel::mpsc::UnboundedSender<SchedulerMsg>,
pub(crate) notified: Cell<bool>,
pub(crate) task: *mut dyn Future<Output = Element<'static>>,
pub(crate) waker: Arc<SuspenseHandle>,
}
impl ArcWake for SuspenseLeaf {
pub struct SuspenseHandle {
pub(crate) id: SuspenseId,
pub(crate) tx: futures_channel::mpsc::UnboundedSender<SchedulerMsg>,
}
impl ArcWake for SuspenseHandle {
fn wake_by_ref(arc_self: &Arc<Self>) {
arc_self.notified.set(true);
_ = arc_self
.tx
.unbounded_send(SchedulerMsg::SuspenseNotified(arc_self.id));
+40 -10
View File
@@ -1,9 +1,13 @@
use super::{waker::ArcWake, Scheduler, SchedulerMsg};
use futures_util::task::ArcWake;
use futures_util::FutureExt;
use super::{Scheduler, SchedulerMsg};
use crate::ScopeId;
use std::cell::RefCell;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::task::Waker;
/// A task's unique identifier.
///
@@ -19,6 +23,17 @@ pub(crate) struct LocalTask {
pub(super) task: RefCell<Pin<Box<dyn Future<Output = ()> + 'static>>>,
id: TaskId,
tx: futures_channel::mpsc::UnboundedSender<SchedulerMsg>,
pub waker: Waker,
}
impl Future for LocalTask {
type Output = ();
fn poll(
self: Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Self::Output> {
self.task.borrow_mut().poll_unpin(cx)
}
}
impl Scheduler {
@@ -32,16 +47,23 @@ impl Scheduler {
/// Spawning a future onto the root scope will cause it to be dropped when the root component is dropped - which
/// will only occur when the VirtuaalDom itself has been dropped.
pub fn spawn(&self, scope: ScopeId, task: impl Future<Output = ()> + 'static) -> TaskId {
let mut tasks = self.tasks.borrow_mut();
let entry = tasks.vacant_entry();
let task_id = TaskId(entry.key());
// let entry = tasks.vacant_entry();
let task_id = TaskId(0);
entry.insert(Arc::new(LocalTask {
let task = LocalTask {
id: task_id,
tx: self.sender.clone(),
task: RefCell::new(Box::pin(task)),
scope,
}));
waker: futures_util::task::waker(Arc::new(LocalTaskHandle {
id: task_id,
tx: self.sender.clone(),
})),
};
self.tasks.borrow_mut().push(task);
// println!("Spawning task: {:?}", task_id);
self.sender
.unbounded_send(SchedulerMsg::TaskNotified(task_id))
@@ -54,14 +76,22 @@ impl Scheduler {
///
/// This does nto 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().remove(id.0);
// self.tasks.borrow_mut().remove(id.0);
}
}
impl ArcWake for LocalTask {
pub struct LocalTaskHandle {
id: TaskId,
tx: futures_channel::mpsc::UnboundedSender<SchedulerMsg>,
}
impl ArcWake for LocalTaskHandle {
fn wake_by_ref(arc_self: &Arc<Self>) {
_ = arc_self
println!("waking task: {:?} {:p}", arc_self.id, arc_self);
arc_self
.tx
.unbounded_send(SchedulerMsg::TaskNotified(arc_self.id));
.unbounded_send(SchedulerMsg::TaskNotified(arc_self.id))
.unwrap();
}
}
+70 -64
View File
@@ -1,3 +1,4 @@
use futures_util::task::ArcWake;
use futures_util::FutureExt;
use std::{
rc::Rc,
@@ -10,7 +11,7 @@ use crate::{
ScopeId, TaskId, VNode, VirtualDom,
};
use super::{waker::ArcWake, SuspenseId};
use super::SuspenseId;
impl VirtualDom {
/// Handle notifications by tasks inside the scheduler
@@ -20,22 +21,27 @@ impl VirtualDom {
pub(crate) fn handle_task_wakeup(&mut self, id: TaskId) {
let mut tasks = self.scheduler.tasks.borrow_mut();
let task = match tasks.get(id.0) {
Some(task) => task,
None => return,
};
// let task = match tasks.get(id.0) {
// Some(task) => task,
// None => {
// // The task was removed from the scheduler, so we can just ignore it
// return;
// }
// };
let waker = task.waker();
let mut cx = Context::from_waker(&waker);
// let mut cx = Context::from_waker(&task.waker);
// If the task completes...
if task.task.borrow_mut().as_mut().poll(&mut cx).is_ready() {
// Remove it from the scope so we dont try to double drop it when the scope dropes
self.scopes[task.scope.0].spawned_tasks.remove(&id);
// // If the task completes...
// println!("polling task");
// if task.task.borrow_mut().as_mut().poll(&mut cx).is_ready() {
// // Remove it from the scope so we dont try to double drop it when the scope dropes
// self.scopes[task.scope.0].spawned_tasks.remove(&id);
// Remove it from the scheduler
tasks.remove(id.0);
}
// // Remove it from the scheduler
// tasks.remove(id.0);
// } else {
// println!("task not ready yet, but we gave it a handle to the waker");
// }
}
pub(crate) fn acquire_suspense_boundary(&self, id: ScopeId) -> Rc<SuspenseContext> {
@@ -45,69 +51,69 @@ impl VirtualDom {
}
pub(crate) fn handle_suspense_wakeup(&mut self, id: SuspenseId) {
let leaf = self
.scheduler
.leaves
.borrow_mut()
.get(id.0)
.unwrap()
.clone();
// let leaf = self
// .scheduler
// .leaves
// .borrow_mut()
// .get(id.0)
// .unwrap()
// .clone();
let scope_id = leaf.scope_id;
// let scope_id = leaf.scope_id;
// todo: cache the waker
let waker = leaf.waker();
let mut cx = Context::from_waker(&waker);
// // todo: cache the waker
// let waker = leaf.waker();
// let mut cx = Context::from_waker(&waker);
// Safety: the future is always pinned to the bump arena
let mut pinned = unsafe { std::pin::Pin::new_unchecked(&mut *leaf.task) };
let as_pinned_mut = &mut pinned;
// // Safety: the future is always pinned to the bump arena
// let mut pinned = unsafe { std::pin::Pin::new_unchecked(&mut *leaf.task) };
// let as_pinned_mut = &mut pinned;
// the component finished rendering and gave us nodes
// we should attach them to that component and then render its children
// continue rendering the tree until we hit yet another suspended component
if let Poll::Ready(new_nodes) = as_pinned_mut.poll_unpin(&mut cx) {
let fiber = self.acquire_suspense_boundary(leaf.scope_id);
// // the component finished rendering and gave us nodes
// // we should attach them to that component and then render its children
// // continue rendering the tree until we hit yet another suspended component
// if let Poll::Ready(new_nodes) = as_pinned_mut.poll_unpin(&mut cx) {
// let fiber = self.acquire_suspense_boundary(leaf.scope_id);
let scope = &mut self.scopes[scope_id.0];
let arena = scope.current_frame();
// let scope = &mut self.scopes[scope_id.0];
// let arena = scope.current_frame();
let ret = arena.bump.alloc(match new_nodes {
Some(new) => RenderReturn::Ready(new),
None => RenderReturn::default(),
});
// let ret = arena.bump.alloc(match new_nodes {
// Some(new) => RenderReturn::Ready(new),
// None => RenderReturn::default(),
// });
arena.node.set(ret);
// arena.node.set(ret);
fiber.waiting_on.borrow_mut().remove(&id);
// fiber.waiting_on.borrow_mut().remove(&id);
if let RenderReturn::Ready(template) = ret {
let mutations_ref = &mut fiber.mutations.borrow_mut();
let mutations = &mut **mutations_ref;
let template: &VNode = unsafe { std::mem::transmute(template) };
let mutations: &mut Mutations = unsafe { std::mem::transmute(mutations) };
// if let RenderReturn::Ready(template) = ret {
// let mutations_ref = &mut fiber.mutations.borrow_mut();
// let mutations = &mut **mutations_ref;
// let template: &VNode = unsafe { std::mem::transmute(template) };
// let mutations: &mut Mutations = unsafe { std::mem::transmute(mutations) };
std::mem::swap(&mut self.mutations, mutations);
// std::mem::swap(&mut self.mutations, mutations);
let place_holder_id = scope.placeholder.get().unwrap();
self.scope_stack.push(scope_id);
let created = self.create(template);
self.scope_stack.pop();
mutations.push(Mutation::ReplaceWith {
id: place_holder_id,
m: created,
});
// let place_holder_id = scope.placeholder.get().unwrap();
// self.scope_stack.push(scope_id);
// let created = self.create(template);
// self.scope_stack.pop();
// mutations.push(Mutation::ReplaceWith {
// id: place_holder_id,
// m: created,
// });
for leaf in self.collected_leaves.drain(..) {
fiber.waiting_on.borrow_mut().insert(leaf);
}
// for leaf in self.collected_leaves.drain(..) {
// fiber.waiting_on.borrow_mut().insert(leaf);
// }
std::mem::swap(&mut self.mutations, mutations);
// std::mem::swap(&mut self.mutations, mutations);
if fiber.waiting_on.borrow().is_empty() {
self.finished_fibers.push(fiber.id);
}
}
}
// if fiber.waiting_on.borrow().is_empty() {
// self.finished_fibers.push(fiber.id);
// }
// }
// }
}
}
-37
View File
@@ -1,37 +0,0 @@
use std::mem;
use std::sync::Arc;
use std::task::{RawWaker, RawWakerVTable, Waker};
pub trait ArcWake: Sized {
/// Create a waker from this self-wakening object
fn waker(self: &Arc<Self>) -> Waker {
unsafe fn rc_vtable<T: ArcWake>() -> &'static RawWakerVTable {
&RawWakerVTable::new(
|data| {
let arc = mem::ManuallyDrop::new(Arc::<T>::from_raw(data.cast::<T>()));
let _rc_clone: mem::ManuallyDrop<_> = arc.clone();
RawWaker::new(data, rc_vtable::<T>())
},
|data| Arc::from_raw(data.cast::<T>()).wake(),
|data| {
let arc = mem::ManuallyDrop::new(Arc::<T>::from_raw(data.cast::<T>()));
ArcWake::wake_by_ref(&arc);
},
|data| drop(Arc::<T>::from_raw(data.cast::<T>())),
)
}
unsafe {
Waker::from_raw(RawWaker::new(
Arc::into_raw(self.clone()).cast(),
rc_vtable::<Self>(),
))
}
}
fn wake_by_ref(arc_self: &Arc<Self>);
fn wake(self: Arc<Self>) {
Self::wake_by_ref(&self)
}
}
+8 -5
View File
@@ -2,9 +2,8 @@ use crate::{
any_props::AnyProps,
bump_frame::BumpFrame,
innerlude::DirtyScope,
innerlude::{SuspenseId, SuspenseLeaf},
innerlude::{SuspenseHandle, SuspenseId, SuspenseLeaf},
nodes::RenderReturn,
scheduler::ArcWake,
scopes::{ScopeId, ScopeState},
virtual_dom::VirtualDom,
};
@@ -90,15 +89,19 @@ impl VirtualDom {
let entry = leaves.vacant_entry();
let suspense_id = SuspenseId(entry.key());
let leaf = Arc::new(SuspenseLeaf {
let leaf = SuspenseLeaf {
scope_id,
task: task.as_mut(),
id: suspense_id,
tx: self.scheduler.sender.clone(),
notified: Default::default(),
});
waker: Arc::new(SuspenseHandle {
id: suspense_id,
tx: self.scheduler.sender.clone(),
}),
};
let waker = leaf.waker();
let waker = futures_util::task::waker(leaf.waker.clone());
let mut cx = Context::from_waker(&waker);
// safety: the task is already pinned in the bump arena
+22 -5
View File
@@ -13,10 +13,21 @@ use crate::{
scopes::{ScopeId, ScopeState},
AttributeValue, Element, Event, Scope, SuspenseContext,
};
use futures_util::{pin_mut, StreamExt};
use futures_util::{
pin_mut,
stream::{futures_unordered, FuturesUnordered},
StreamExt,
};
use rustc_hash::FxHashMap;
use slab::Slab;
use std::{any::Any, borrow::BorrowMut, cell::Cell, collections::BTreeSet, future::Future, rc::Rc};
use std::{
any::Any,
borrow::BorrowMut,
cell::Cell,
collections::{BTreeSet, HashSet},
future::Future,
rc::Rc,
};
/// A virtual node system that progresses user events and diffs UI trees.
///
@@ -420,7 +431,7 @@ impl VirtualDom {
loop {
match some_msg.take() {
// If a bunch of messages are ready in a sequence, try to pop them off synchronously
Some(msg) => match msg {
Some(msg) => match dbg!(msg) {
SchedulerMsg::Immediate(id) => self.mark_dirty(id),
SchedulerMsg::TaskNotified(task) => self.handle_task_wakeup(task),
SchedulerMsg::SuspenseNotified(id) => self.handle_suspense_wakeup(id),
@@ -430,14 +441,20 @@ impl VirtualDom {
None => {
match self.rx.try_next() {
Ok(Some(val)) => some_msg = Some(val),
Ok(None) => return,
Ok(None) => panic!("Scheduler channel closed"),
Err(_) => {
// If we have any dirty scopes, or finished fiber trees then we should exit
if !self.dirty_scopes.is_empty() || !self.finished_fibers.is_empty() {
return;
}
some_msg = self.rx.next().await
println!("Waiting for next message...");
let mut tasks = self.scheduler.tasks.borrow_mut();
let takss = &mut *tasks;
use futures_util::StreamExt;
takss.next().await;
}
}
}