fix memos in children; introduce isomorphic spawn

Cargo.lock

@@ -2347,6 +2347,7 @@ dependencies = [
  "longest-increasing-subsequence",
  "pretty_assertions",
  "rand 0.8.5",
+ "reqwest",
  "rustc-hash",
  "serde",
  "slab",

packages/core/Cargo.toml

@@ -29,6 +29,7 @@ dioxus = { workspace = true }
 pretty_assertions = "1.3.0"
 rand = "0.8.5"
 dioxus-ssr = { workspace = true }
+reqwest.workspace = true
 
 [features]
 default = []

packages/core/src/global_context.rs

@@ -56,6 +56,35 @@ pub fn suspend(task: Task) -> Element {
     None
 }
 
+/// Start a new future on the same thread as the rest of the VirtualDom.
+///
+/// **You should generally use `spawn` instead of this method unless you specifically need to need to run a task during suspense**
+///
+/// This future will not contribute to suspense resolving but it will run during suspense.
+///
+/// Because this future runs during suspense, you need to be careful to work with hydration. It is not recommended to do any async IO work in this future, as it can easily cause hydration issues. However, you can use isomorphic tasks to do work that can be consistently replicated on the server and client like logging or responding to state changes.
+///
+/// ```rust, no_run
+/// # use dioxus::prelude::*;
+/// // ❌ Do not do requests in isomorphic tasks. It may resolve at a different time on the server and client, causing hydration issues.
+/// let mut state = use_signal(|| None);
+/// spawn_isomorphic(async move {
+///     state.set(Some(reqwest::get("https://api.example.com").await));
+/// });
+///
+/// // ✅ You may wait for a signal to change and then log it
+/// let mut state = use_signal(|| 0);
+/// spawn_isomorphic(async move {
+///     loop {
+///         tokio::time::sleep(std::time::Duration::from_secs(1)).await;
+///         println!("State is {state}");
+///     }
+/// });
+/// ```
+pub fn spawn_isomorphic(fut: impl Future<Output = ()> + 'static) -> Task {
+    Runtime::with_current_scope(|cx| cx.spawn_isomorphic(fut)).expect("to be in a dioxus runtime")
+}
+
 /// Spawns the future but does not return the [`TaskId`]
 pub fn spawn(fut: impl Future<Output = ()> + 'static) -> Task {
     Runtime::with_current_scope(|cx| cx.spawn(fut)).expect("to be in a dioxus runtime")

packages/core/src/lib.rs

@@ -91,11 +91,12 @@ pub mod prelude {
         consume_context, consume_context_from_scope, current_scope_id, fc_to_builder, generation,
         has_context, needs_update, needs_update_any, parent_scope, provide_context,
         provide_root_context, remove_future, schedule_update, schedule_update_any, spawn,
-        spawn_forever, suspend, try_consume_context, use_after_render, use_before_render, use_drop,
-        use_error_boundary, use_hook, use_hook_with_cleanup, wait_for_next_render, AnyValue,
-        Attribute, Component, ComponentFunction, Element, ErrorBoundary, Event, EventHandler,
-        Fragment, HasAttributes, IntoAttributeValue, IntoDynNode, OptionStringFromMarker,
-        Properties, Runtime, RuntimeGuard, ScopeId, ScopeState, SuperFrom, SuperInto, Task,
-        Template, TemplateAttribute, TemplateNode, Throw, VNode, VNodeInner, VirtualDom,
+        spawn_forever, spawn_isomorphic, suspend, try_consume_context, use_after_render,
+        use_before_render, use_drop, use_error_boundary, use_hook, use_hook_with_cleanup,
+        wait_for_next_render, AnyValue, Attribute, Component, ComponentFunction, Element,
+        ErrorBoundary, Event, EventHandler, Fragment, HasAttributes, IntoAttributeValue,
+        IntoDynNode, OptionStringFromMarker, Properties, Runtime, RuntimeGuard, ScopeId,
+        ScopeState, SuperFrom, SuperInto, Task, Template, TemplateAttribute, TemplateNode, Throw,
+        VNode, VNodeInner, VirtualDom,
     };
 }

packages/core/src/runtime.rs

@@ -1,5 +1,3 @@
-use rustc_hash::FxHashSet;
-
 use crate::{
     innerlude::{LocalTask, SchedulerMsg},
     render_signal::RenderSignal,
@@ -30,7 +28,7 @@ pub struct Runtime {
     pub(crate) tasks: RefCell<slab::Slab<Rc<LocalTask>>>,
 
     // Currently suspended tasks
-    pub(crate) suspended_tasks: RefCell<FxHashSet<Task>>,
+    pub(crate) suspended_tasks: Cell<usize>,
 
     pub(crate) rendering: Cell<bool>,
 

packages/core/src/scope_arena.rs

@@ -72,7 +72,10 @@ impl VirtualDom {
         if let Some(task) = context.last_suspendable_task.take() {
             if matches!(new_nodes, RenderReturn::Aborted(_)) {
                 tracing::trace!("Suspending {:?} on {:?}", scope_id, task);
-                self.runtime.suspended_tasks.borrow_mut().insert(task);
+                self.runtime.tasks.borrow().get(task.0).unwrap().suspend();
+                self.runtime
+                    .suspended_tasks
+                    .set(self.runtime.suspended_tasks.get() + 1);
             }
         }
 

packages/core/src/scope_context.rs

@@ -227,6 +227,37 @@ impl Scope {
         .expect("Runtime to exist")
     }
 
+    /// Start a new future on the same thread as the rest of the VirtualDom.
+    ///
+    /// **You should generally use `spawn` instead of this method unless you specifically need to need to run a task during suspense**
+    ///
+    /// This future will not contribute to suspense resolving but it will run during suspense.
+    ///
+    /// Because this future runs during suspense, you need to be careful to work with hydration. It is not recommended to do any async IO work in this future, as it can easily cause hydration issues. However, you can use isomorphic tasks to do work that can be consistently replicated on the server and client like logging or responding to state changes.
+    ///
+    /// ```rust, no_run
+    /// # use dioxus::prelude::*;
+    /// // ❌ Do not do requests in isomorphic tasks. It may resolve at a different time on the server and client, causing hydration issues.
+    /// let mut state = use_signal(|| None);
+    /// spawn_isomorphic(async move {
+    ///     state.set(Some(reqwest::get("https://api.example.com").await));
+    /// });
+    ///
+    /// // ✅ You may wait for a signal to change and then log it
+    /// let mut state = use_signal(|| 0);
+    /// spawn_isomorphic(async move {
+    ///     loop {
+    ///         tokio::time::sleep(std::time::Duration::from_secs(1)).await;
+    ///         println!("State is {state}");
+    ///     }
+    /// });
+    /// ```
+    pub fn spawn_isomorphic(&self, fut: impl Future<Output = ()> + 'static) -> Task {
+        let id = Runtime::with(|rt| rt.spawn_isomorphic(self.id, fut)).expect("Runtime to exist");
+        self.spawned_tasks.borrow_mut().insert(id);
+        id
+    }
+
     /// Spawns the future but does not return the [`TaskId`]
     pub fn spawn(&self, fut: impl Future<Output = ()> + 'static) -> Task {
         let id = Runtime::with(|rt| rt.spawn(self.id, fut)).expect("Runtime to exist");

packages/core/src/tasks.rs

@@ -81,6 +81,39 @@ impl Task {
 }
 
 impl Runtime {
+    /// Start a new future on the same thread as the rest of the VirtualDom.
+    ///
+    /// **You should generally use `spawn` instead of this method unless you specifically need to need to run a task during suspense**
+    ///
+    /// This future will not contribute to suspense resolving but it will run during suspense.
+    ///
+    /// Because this future runs during suspense, you need to be careful to work with hydration. It is not recommended to do any async IO work in this future, as it can easily cause hydration issues. However, you can use isomorphic tasks to do work that can be consistently replicated on the server and client like logging or responding to state changes.
+    ///
+    /// ```rust, no_run
+    /// # use dioxus::prelude::*;
+    /// // ❌ Do not do requests in isomorphic tasks. It may resolve at a different time on the server and client, causing hydration issues.
+    /// let mut state = use_signal(|| None);
+    /// spawn_isomorphic(async move {
+    ///     state.set(Some(reqwest::get("https://api.example.com").await));
+    /// });
+    ///
+    /// // ✅ You may wait for a signal to change and then log it
+    /// let mut state = use_signal(|| 0);
+    /// spawn_isomorphic(async move {
+    ///     loop {
+    ///         tokio::time::sleep(std::time::Duration::from_secs(1)).await;
+    ///         println!("State is {state}");
+    ///     }
+    /// });
+    /// ```
+    pub fn spawn_isomorphic(
+        &self,
+        scope: ScopeId,
+        task: impl Future<Output = ()> + 'static,
+    ) -> Task {
+        self.spawn_task_of_type(scope, task, TaskType::Isomorphic)
+    }
+
     /// 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
@@ -91,6 +124,15 @@ impl Runtime {
     /// 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 {
+        self.spawn_task_of_type(scope, task, TaskType::ClientOnly)
+    }
+
+    fn spawn_task_of_type(
+        &self,
+        scope: ScopeId,
+        task: impl Future<Output = ()> + 'static,
+        ty: TaskType,
+    ) -> Task {
         // Insert the task, temporarily holding a borrow on the tasks map
         let (task, task_id) = {
             let mut tasks = self.tasks.borrow_mut();
@@ -107,6 +149,7 @@ impl Runtime {
                     id: task_id,
                     tx: self.sender.clone(),
                 })),
+                ty: Cell::new(ty),
             });
 
             entry.insert(task.clone());
@@ -186,8 +229,20 @@ 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<Rc<LocalTask>> {
-        self.suspended_tasks.borrow_mut().remove(&id);
-        self.tasks.borrow_mut().try_remove(id.0)
+        let task = self.tasks.borrow_mut().try_remove(id.0);
+        if let Some(task) = &task {
+            if task.suspended() {
+                self.suspended_tasks.set(self.suspended_tasks.get() - 1);
+            }
+        }
+        task
+    }
+
+    /// Check if a task should be run during suspense
+    pub(crate) fn task_runs_during_suspense(&self, task: Task) -> bool {
+        let borrow = self.tasks.borrow();
+        let task: Option<&LocalTask> = borrow.get(task.0).map(|t| &**t);
+        matches!(task, Some(LocalTask { ty, .. }) if ty.get().runs_during_suspense())
     }
 }
 
@@ -197,9 +252,33 @@ pub(crate) struct LocalTask {
     parent: Option<Task>,
     task: RefCell<Pin<Box<dyn Future<Output = ()> + 'static>>>,
     waker: Waker,
+    ty: Cell<TaskType>,
     active: Cell<bool>,
 }
 
+impl LocalTask {
+    pub(crate) fn suspend(&self) {
+        self.ty.set(TaskType::Suspended);
+    }
+
+    pub(crate) fn suspended(&self) -> bool {
+        matches!(self.ty.get(), TaskType::Suspended)
+    }
+}
+
+#[derive(Clone, Copy)]
+enum TaskType {
+    ClientOnly,
+    Suspended,
+    Isomorphic,
+}
+
+impl TaskType {
+    fn runs_during_suspense(self) -> bool {
+        matches!(self, TaskType::Isomorphic | TaskType::Suspended)
+    }
+}
+
 /// The type of message that can be sent to the scheduler.
 ///
 /// These messages control how the scheduler will process updates to the UI.

packages/core/src/virtual_dom.rs

@@ -644,7 +644,7 @@ impl VirtualDom {
     #[instrument(skip(self), level = "trace", name = "VirtualDom::wait_for_suspense")]
     pub async fn wait_for_suspense(&mut self) {
         loop {
-            if self.runtime.suspended_tasks.borrow().is_empty() {
+            if self.runtime.suspended_tasks.get() == 0 {
                 break;
             }
 
@@ -668,7 +668,7 @@ impl VirtualDom {
                         // 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) {
+                            if self.runtime.task_runs_during_suspense(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();
@@ -689,7 +689,7 @@ impl VirtualDom {
                 // 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) {
+                    if self.runtime.task_runs_during_suspense(task) {
                         let _ = self.runtime.handle_task_wakeup(task);
                     }
                 }

packages/signals/src/memo.rs

@@ -14,31 +14,6 @@ use futures_util::StreamExt;
 use generational_box::UnsyncStorage;
 use once_cell::sync::OnceCell;
 
-/// A thread local that can only be read from the thread it was created on.
-pub struct ThreadLocal<T> {
-    value: T,
-    owner: std::thread::ThreadId,
-}
-
-impl<T> ThreadLocal<T> {
-    /// Create a new thread local.
-    pub fn new(value: T) -> Self {
-        ThreadLocal {
-            value,
-            owner: std::thread::current().id(),
-        }
-    }
-
-    /// Get the value of the thread local.
-    pub fn get(&self) -> Option<&T> {
-        (self.owner == std::thread::current().id()).then_some(&self.value)
-    }
-}
-
-// SAFETY: This is safe because the thread local can only be read from the thread it was created on.
-unsafe impl<T> Send for ThreadLocal<T> {}
-unsafe impl<T> Sync for ThreadLocal<T> {}
-
 struct UpdateInformation<T> {
     dirty: Arc<AtomicBool>,
     callback: RefCell<Box<dyn FnMut() -> T>>,
@@ -70,25 +45,12 @@ impl<T: 'static> Memo<T> {
         let (tx, mut rx) = futures_channel::mpsc::unbounded();
 
         let myself: Rc<OnceCell<Memo<T>>> = Rc::new(OnceCell::new());
-        let thread_local = ThreadLocal::new(myself.clone());
 
         let callback = {
             let dirty = dirty.clone();
-            move || match thread_local.get() {
-                Some(memo) => match memo.get() {
-                    Some(memo) => {
-                        memo.recompute();
-                    }
-                    None => {
-                        tracing::error!("Memo was not initialized in the same thread it was created in. This is likely a bug in dioxus");
-                        dirty.store(true, std::sync::atomic::Ordering::Relaxed);
-                        let _ = tx.unbounded_send(());
-                    }
-                },
-                None => {
-                    dirty.store(true, std::sync::atomic::Ordering::Relaxed);
-                    let _ = tx.unbounded_send(());
-                }
+            move || {
+                dirty.store(true, std::sync::atomic::Ordering::Relaxed);
+                let _ = tx.unbounded_send(());
             }
         };
         let rc = ReactiveContext::new_with_callback(
@@ -113,7 +75,7 @@ impl<T: 'static> Memo<T> {
         };
         let _ = myself.set(memo);
 
-        spawn(async move {
+        spawn_isomorphic(async move {
             while rx.next().await.is_some() {
                 // Remove any pending updates
                 while rx.try_next().is_ok() {}