Cleanup more of core

packages/core/src/tasks.rs

@@ -0,0 +1,165 @@
+use crate::innerlude::{remove_future, spawn, Runtime};
+use crate::ScopeId;
+use futures_util::task::ArcWake;
+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.
+///
+/// `Task` is a unique identifier for a task that has been spawned onto the runtime. It can be used to cancel the task
+#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
+#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
+pub struct Task(pub(crate) usize);
+
+impl Task {
+    /// 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
+    /// and long running tasks.
+    ///
+    /// Whenever the component that owns this future is dropped, the future will be dropped as well.
+    ///
+    /// 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 VirtualDom itself has been dropped.
+    pub fn new(task: impl Future<Output = ()> + 'static) -> Self {
+        spawn(task)
+    }
+
+    /// Drop the task immediately.
+    ///
+    /// This does not abort the task, so you'll want to wrap it in an abort handle if that's important to you
+    pub fn stop(self) {
+        remove_future(self);
+    }
+}
+
+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
+    /// and long running tasks.
+    ///
+    /// Whenever the component that owns this future is dropped, the future will be dropped as well.
+    ///
+    /// 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 VirtualDom itself has been dropped.
+    pub fn spawn(&self, scope: ScopeId, task: impl Future<Output = ()> + 'static) -> Task {
+        let mut tasks = self.tasks.borrow_mut();
+
+        let entry = tasks.vacant_entry();
+        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 {
+                id: task_id,
+                tx: self.sender.clone(),
+            })),
+        };
+
+        let mut cx = std::task::Context::from_waker(&task.waker);
+
+        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
+    }
+
+    pub(crate) fn handle_task_wakeup(&self, id: Task) {
+        let mut tasks = self.tasks.borrow_mut();
+
+        let task = match tasks.get(id.0) {
+            Some(task) => task,
+            // The task was removed from the scheduler, so we can just ignore it
+            None => return,
+        };
+
+        use std::task::Context;
+
+        let mut cx = Context::from_waker(&task.waker);
+
+        // update the scope stack
+        self.scope_stack.borrow_mut().push(task.scope);
+        self.rendering.set(false);
+        self.current_task.set(Some(id));
+
+        // 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.get_context(task.scope)
+                .unwrap()
+                .spawned_tasks
+                .borrow_mut()
+                .remove(&id);
+
+            // Remove it from the scheduler
+            tasks.try_remove(id.0);
+        }
+
+        // Remove the scope from the stack
+        self.scope_stack.borrow_mut().pop();
+        self.rendering.set(true);
+        self.current_task.set(None);
+    }
+
+    /// Drop the future with the given TaskId
+    ///
+    /// 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<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()
+    }
+
+    /// Get the parent task of the given task, if it exists
+    pub fn parent_task(&self, task: Task) -> Option<Task> {
+        self.tasks.borrow().get(task.0)?.parent
+    }
+}
+
+/// the task itself is the waker
+pub(crate) struct LocalTask {
+    scope: ScopeId,
+    parent: Option<Task>,
+    task: RefCell<Pin<Box<dyn Future<Output = ()> + 'static>>>,
+    waker: Waker,
+}
+
+/// The type of message that can be sent to the scheduler.
+///
+/// These messages control how the scheduler will process updates to the UI.
+#[derive(Debug)]
+pub(crate) enum SchedulerMsg {
+    /// Immediate updates from Components that mark them as dirty
+    Immediate(ScopeId),
+
+    /// A task has woken and needs to be progressed
+    TaskNotified(Task),
+}
+
+struct LocalTaskHandle {
+    id: Task,
+    tx: futures_channel::mpsc::UnboundedSender<SchedulerMsg>,
+}
+
+impl ArcWake for LocalTaskHandle {
+    fn wake_by_ref(arc_self: &Arc<Self>) {
+        // This can fail if the scheduler has been dropped while the application is shutting down
+        let _ = arc_self
+            .tx
+            .unbounded_send(SchedulerMsg::TaskNotified(arc_self.id));
+    }
+}