Fix suspense diffing (#2789)

* Fix suspense placeholders after creating a node
* add nested suspense client side core test
* fix resolve suspense

packages/core/src/runtime.rs

@@ -243,7 +243,7 @@ impl Runtime {
         // If this is not a suspended scope, and we are under a frozen context, then we should
         let scopes = self.scope_states.borrow();
         let scope = &scopes[scope_id.0].as_ref().unwrap();
-        !matches!(scope.suspense_location(), SuspenseLocation::UnderSuspense(suspense) if suspense.has_suspended_tasks())
+        !matches!(scope.suspense_location(), SuspenseLocation::UnderSuspense(suspense) if suspense.is_suspended())
     }
 }
 

packages/core/src/suspense/component.rs

@@ -421,7 +421,8 @@ impl SuspenseBoundaryProps {
             let suspense_context =
                 SuspenseContext::downcast_suspense_boundary_from_scope(&dom.runtime, scope_id)
                     .unwrap();
-            let suspended = suspense_context.suspended_nodes();
+            // Take the suspended nodes out of the suspense boundary so the children know that the boundary is not suspended while diffing
+            let suspended = suspense_context.take_suspended_nodes();
             if let Some(node) = suspended {
                 node.remove_node(&mut *dom, None::<&mut M>, None);
             }
@@ -444,10 +445,6 @@ impl SuspenseBoundaryProps {
             let props = Self::downcast_from_props(&mut *scope_state.props).unwrap();
             props.children = children.clone().node;
             scope_state.last_rendered_node = Some(children);
-            let suspense_context =
-                SuspenseContext::downcast_suspense_boundary_from_scope(&dom.runtime, scope_id)
-                    .unwrap();
-            suspense_context.take_suspended_nodes();
         })
     }
 
@@ -555,8 +552,11 @@ impl SuspenseBoundaryProps {
                     suspense_context.set_suspended_nodes(new_children);
 
                     un_resolve_suspense(dom, scope_id);
-                } // We have suspended nodes, but we just got out of suspense. Move the suspended nodes to the foreground
+                }
-                (Some(old_suspended_nodes), false) => {
+                // We have suspended nodes, but we just got out of suspense. Move the suspended nodes to the foreground
+                (Some(_), false) => {
+                    // Take the suspended nodes out of the suspense boundary so the children know that the boundary is not suspended while diffing
+                    let old_suspended_nodes = suspense_context.take_suspended_nodes().unwrap();
                     let old_placeholder = last_rendered_node;
                     let new_children = RenderReturn { node: children };
 
@@ -579,13 +579,6 @@ impl SuspenseBoundaryProps {
                     // Set the last rendered node to the new children
                     dom.scopes[scope_id.0].last_rendered_node = Some(new_children);
 
-                    let suspense_context = SuspenseContext::downcast_suspense_boundary_from_scope(
-                        &dom.runtime,
-                        scope_id,
-                    )
-                    .unwrap();
-                    suspense_context.take_suspended_nodes();
-
                     mark_suspense_resolved(dom, scope_id);
                 }
             }

packages/core/src/suspense/mod.rs

@@ -148,6 +148,11 @@ impl SuspenseContext {
         !self.inner.suspended_tasks.borrow().is_empty()
     }
 
+    /// Check if the suspense boundary is currently rendered as suspended
+    pub fn is_suspended(&self) -> bool {
+        self.inner.suspended_nodes.borrow().is_some()
+    }
+
     /// Add a suspended task
     pub(crate) fn add_suspended_task(&self, task: SuspendedFuture) {
         self.inner.suspended_tasks.borrow_mut().push(task);

packages/core/tests/suspense.rs

@@ -1,4 +1,6 @@
 use dioxus::prelude::*;
+use dioxus_core::{AttributeValue, ElementId, Mutation};
+use pretty_assertions::assert_eq;
 use std::future::poll_fn;
 use std::task::Poll;
 
@@ -19,15 +21,16 @@ async fn poll_three_times() {
 }
 
 #[test]
-fn suspense_resolves() {
+fn suspense_resolves_ssr() {
     // wait just a moment, not enough time for the boundary to resolve
     tokio::runtime::Builder::new_current_thread()
         .build()
         .unwrap()
         .block_on(async {
             let mut dom = VirtualDom::new(app);
-            dom.rebuild(&mut dioxus_core::NoOpMutations);
+            dom.rebuild_in_place();
             dom.wait_for_suspense().await;
+            dom.render_immediate(&mut dioxus_core::NoOpMutations);
             let out = dioxus_ssr::render(&dom);
 
             assert_eq!(out, "<div>Waiting for... child</div>");
@@ -255,10 +258,6 @@ fn suspended_nodes_dont_trigger_effects() {
 /// Make sure we keep any state of components when we switch from a resolved future to a suspended future
 #[test]
 fn resolved_to_suspended() {
-    tracing_subscriber::fmt::SubscriberBuilder::default()
-        .with_max_level(tracing::Level::INFO)
-        .init();
-
     static SUSPENDED: GlobalSignal<bool> = Signal::global(|| false);
 
     tokio::runtime::Builder::new_current_thread()
@@ -372,3 +371,588 @@ fn suspense_tracks_resolved() {
         }
     }
 }
+
+// Regression test for https://github.com/DioxusLabs/dioxus/issues/2783
+#[test]
+fn toggle_suspense() {
+    use dioxus::prelude::*;
+
+    fn app() -> Element {
+        rsx! {
+            SuspenseBoundary {
+                fallback: |_| rsx! { "fallback" },
+                if generation() % 2 == 0 {
+                    Page {}
+                } else {
+                    Home {}
+                }
+            }
+        }
+    }
+
+    #[component]
+    pub fn Home() -> Element {
+        let _calculation = use_resource(|| async move {
+            tokio::time::sleep(std::time::Duration::from_secs(1)).await;
+            1 + 1
+        })
+        .suspend()?;
+        rsx! {
+            "hello world"
+        }
+    }
+
+    #[component]
+    pub fn Page() -> Element {
+        rsx! {
+            "goodbye world"
+        }
+    }
+
+    tokio::runtime::Builder::new_current_thread()
+        .enable_time()
+        .build()
+        .unwrap()
+        .block_on(async {
+            let mut dom = VirtualDom::new(app);
+            let mutations = dom.rebuild_to_vec().sanitize();
+
+            // First create goodbye world
+            println!("{:#?}", mutations);
+            assert_eq!(
+                mutations.edits,
+                [
+                    Mutation::LoadTemplate { name: "template", index: 0, id: ElementId(1) },
+                    Mutation::AppendChildren { id: ElementId(0), m: 1 }
+                ]
+            );
+
+            dom.mark_dirty(ScopeId::APP);
+            let mutations = dom.render_immediate_to_vec().sanitize();
+
+            // Then replace that with nothing
+            println!("{:#?}", mutations);
+            assert_eq!(
+                mutations.edits,
+                [
+                    Mutation::CreatePlaceholder { id: ElementId(2) },
+                    Mutation::ReplaceWith { id: ElementId(1), m: 1 },
+                ]
+            );
+
+            dom.wait_for_work().await;
+            let mutations = dom.render_immediate_to_vec().sanitize();
+
+            // Then replace it with a placeholder
+            println!("{:#?}", mutations);
+            assert_eq!(
+                mutations.edits,
+                [
+                    Mutation::LoadTemplate { name: "template", index: 0, id: ElementId(1) },
+                    Mutation::ReplaceWith { id: ElementId(2), m: 1 },
+                ]
+            );
+
+            dom.wait_for_work().await;
+            let mutations = dom.render_immediate_to_vec().sanitize();
+
+            // Then replace it with the resolved node
+            println!("{:#?}", mutations);
+            assert_eq!(
+                mutations.edits,
+                [
+                    Mutation::CreatePlaceholder { id: ElementId(2,) },
+                    Mutation::ReplaceWith { id: ElementId(1,), m: 1 },
+                    Mutation::LoadTemplate { name: "template", index: 0, id: ElementId(1) },
+                    Mutation::ReplaceWith { id: ElementId(2), m: 1 },
+                ]
+            );
+        });
+}
+
+#[test]
+fn nested_suspense_resolves_client() {
+    use Mutation::*;
+
+    async fn poll_three_times() {
+        // 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;
+    }
+
+    fn app() -> Element {
+        rsx! {
+            SuspenseBoundary {
+                fallback: move |_| rsx! {},
+                LoadTitle {}
+            }
+            MessageWithLoader { id: 0 }
+        }
+    }
+
+    #[component]
+    fn MessageWithLoader(id: usize) -> Element {
+        rsx! {
+            SuspenseBoundary {
+                fallback: move |_| rsx! {
+                    "Loading {id}..."
+                },
+                Message { id }
+            }
+        }
+    }
+
+    #[component]
+    fn LoadTitle() -> Element {
+        let title = use_resource(move || async_content(0)).suspend()?();
+
+        rsx! {
+            Title { "{title.title}" }
+        }
+    }
+
+    #[component]
+    fn Message(id: usize) -> Element {
+        let message = use_resource(move || async_content(id)).suspend()?();
+
+        rsx! {
+            h2 {
+                id: "title-{id}",
+                "{message.title}"
+            }
+            p {
+                id: "body-{id}",
+                "{message.body}"
+            }
+            div {
+                id: "children-{id}",
+                padding: "10px",
+                for child in message.children {
+                    MessageWithLoader { id: child }
+                }
+            }
+        }
+    }
+
+    #[derive(Clone)]
+    pub struct Content {
+        title: String,
+        body: String,
+        children: Vec<usize>,
+    }
+
+    async fn async_content(id: usize) -> Content {
+        let content_tree = [
+            Content {
+                title: "The robot says hello world".to_string(),
+                body: "The robot becomes sentient and says hello world".to_string(),
+                children: vec![1, 2],
+            },
+            Content {
+                title: "The world says hello back".to_string(),
+                body: "In a stunning turn of events, the world collectively unites and says hello back"
+                    .to_string(),
+                children: vec![],
+            },
+            Content {
+                title: "Goodbye Robot".to_string(),
+                body: "The robot says goodbye".to_string(),
+                children: vec![3],
+            },
+            Content {
+                title: "Goodbye Robot again".to_string(),
+                body: "The robot says goodbye again".to_string(),
+                children: vec![],
+            },
+        ];
+        poll_three_times().await;
+        content_tree[id].clone()
+    }
+
+    // wait just a moment, not enough time for the boundary to resolve
+    tokio::runtime::Builder::new_current_thread()
+        .build()
+        .unwrap()
+        .block_on(async {
+            let mut dom = VirtualDom::new(app);
+            let mutations = dom.rebuild_to_vec();
+            // Initial loading message and loading title
+            assert_eq!(
+                mutations.edits,
+                vec![
+                    CreatePlaceholder { id: ElementId(1,) },
+                    CreateTextNode { value: "Loading 0...".to_string(), id: ElementId(2,) },
+                    AppendChildren { id: ElementId(0,), m: 2 },
+                ]
+            );
+
+            dom.wait_for_work().await;
+            // DOM STATE:
+            // placeholder // ID: 1
+            // "Loading 0..." // ID: 2
+            let mutations = dom.render_immediate_to_vec().sanitize();
+            // Fill in the contents of the initial message and start loading the nested suspense
+            // The title also finishes loading
+            assert_eq!(
+                mutations.edits,
+                vec![
+                    // Creating and swapping these placeholders doesn't do anything
+                    // It is just extra work that we are forced to do because mutations are not 
+                    // reversible. We start rendering the children and then realize it is suspended.
+                    // Then we need to replace what we just rendered with the suspense placeholder
+                    CreatePlaceholder { id: ElementId(3,) },
+                    ReplaceWith { id: ElementId(1,), m: 1 },
+
+                    // Replace the pending placeholder with the title placeholder
+                    CreatePlaceholder { id: ElementId(1,) },
+                    ReplaceWith { id: ElementId(3,), m: 1 },
+
+                    // Replace loading... with a placeholder for us to fill in later
+                    CreatePlaceholder { id: ElementId(3,) },
+                    ReplaceWith { id: ElementId(2,), m: 1 },
+
+                    // Load the title
+                    LoadTemplate { name: "template", index: 0, id: ElementId(2,) },
+                    HydrateText {
+                        path: &[0,],
+                        value: "The robot says hello world".to_string(),
+                        id: ElementId(4,),
+                    },
+                    SetAttribute {
+                        name: "id",
+                        ns: None,
+                        value: AttributeValue::Text("title-0".to_string()),
+                        id: ElementId(2,),
+                    },
+
+                    // Then load the body
+                    LoadTemplate { name: "template", index: 1, id: ElementId(5,) },
+                    HydrateText {
+                        path: &[0,],
+                        value: "The robot becomes sentient and says hello world".to_string(),
+                        id: ElementId(6,),
+                    },
+                    SetAttribute {
+                        name: "id",
+                        ns: None,
+                        value: AttributeValue::Text("body-0".to_string()),
+                        id: ElementId(5,),
+                    },
+
+                    // Then load the suspended children
+                    LoadTemplate { name: "template", index: 2, id: ElementId(7,) },
+                    CreateTextNode { value: "Loading 1...".to_string(), id: ElementId(8,) },
+                    CreateTextNode { value: "Loading 2...".to_string(), id: ElementId(9,) },
+                    ReplacePlaceholder { path: &[0,], m: 2 },
+                    SetAttribute {
+                        name: "id",
+                        ns: None,
+                        value: AttributeValue::Text("children-0".to_string()),
+                        id: ElementId(7,),
+                    },
+
+                    // Finally replace the loading placeholder in the body with the resolved children
+                    ReplaceWith { id: ElementId(3,), m: 3 },
+                ]
+            );
+
+            dom.wait_for_work().await;
+            // DOM STATE:
+            // placeholder // ID: 1
+            // h2 // ID: 2
+            // p // ID: 5
+            // div // ID: 7
+            //   "Loading 1..." // ID: 8
+            //   "Loading 2..." // ID: 9
+            let mutations = dom.render_immediate_to_vec().sanitize();
+            assert_eq!(
+                mutations.edits,
+                vec![
+                    // Replace the first loading placeholder with a placeholder for us to fill in later
+                    CreatePlaceholder {
+                        id: ElementId(
+                            3,
+                        ),
+                    },
+                    ReplaceWith {
+                        id: ElementId(
+                            8,
+                        ),
+                        m: 1,
+                    },
+
+                    // Load the nested suspense
+                    LoadTemplate {
+                        name: "template",
+                        index: 0,
+                        id: ElementId(
+                            8,
+                        ),
+                    },
+                    HydrateText {
+                        path: &[
+                            0,
+                        ],
+                        value: "The world says hello back".to_string(),
+                        id: ElementId(
+                            10,
+                        ),
+                    },
+                    SetAttribute {
+                        name: "id",
+                        ns: None,
+                        value: AttributeValue::Text("title-1".to_string()),
+                        id: ElementId(
+                            8,
+                        ),
+                    },
+                    LoadTemplate {
+                        name: "template",
+                        index: 1,
+                        id: ElementId(
+                            11,
+                        ),
+                    },
+                    HydrateText {
+                        path: &[
+                            0,
+                        ],
+                        value: "In a stunning turn of events, the world collectively unites and says hello back".to_string(),
+                        id: ElementId(
+                            12,
+                        ),
+                    },
+                    SetAttribute {
+                        name: "id",
+                        ns: None,
+                        value: AttributeValue::Text("body-1".to_string()),
+                        id: ElementId(
+                            11,
+                        ),
+                    },
+                    LoadTemplate {
+                        name: "template",
+                        index: 2,
+                        id: ElementId(
+                            13,
+                        ),
+                    },
+                    AssignId {
+                        path: &[
+                            0,
+                        ],
+                        id: ElementId(
+                            14,
+                        ),
+                    },
+                    SetAttribute {
+                        name: "id",
+                        ns: None,
+                        value: AttributeValue::Text("children-1".to_string()),
+                        id: ElementId(
+                            13,
+                        ),
+                    },
+                    ReplaceWith {
+                        id: ElementId(
+                            3,
+                        ),
+                        m: 3,
+                    },
+
+                    // Replace the second loading placeholder with a placeholder for us to fill in later
+                    CreatePlaceholder {
+                        id: ElementId(
+                            3,
+                        ),
+                    },
+                    ReplaceWith {
+                        id: ElementId(
+                            9,
+                        ),
+                        m: 1,
+                    },
+                    LoadTemplate {
+                        name: "template",
+                        index: 0,
+                        id: ElementId(
+                            9,
+                        ),
+                    },
+                    HydrateText {
+                        path: &[
+                            0,
+                        ],
+                        value: "Goodbye Robot".to_string(),
+                        id: ElementId(
+                            15,
+                        ),
+                    },
+                    SetAttribute {
+                        name: "id",
+                        ns: None,
+                        value: AttributeValue::Text("title-2".to_string()),
+                        id: ElementId(
+                            9,
+                        ),
+                    },
+                    LoadTemplate {
+                        name: "template",
+                        index: 1,
+                        id: ElementId(
+                            16,
+                        ),
+                    },
+                    HydrateText {
+                        path: &[
+                            0,
+                        ],
+                        value: "The robot says goodbye".to_string(),
+                        id: ElementId(
+                            17,
+                        ),
+                    },
+                    SetAttribute {
+                        name: "id",
+                        ns: None,
+                        value: AttributeValue::Text("body-2".to_string()),
+                        id: ElementId(
+                            16,
+                        ),
+                    },
+                    LoadTemplate {
+                        name: "template",
+                        index: 2,
+                        id: ElementId(
+                            18,
+                        ),
+                    },
+                    // Create a placeholder for the resolved children
+                    CreateTextNode { value: "Loading 3...".to_string(), id: ElementId(19,) },
+                    ReplacePlaceholder { path: &[0,], m: 1 },
+                    SetAttribute {
+                        name: "id",
+                        ns: None,
+                        value: AttributeValue::Text("children-2".to_string()),
+                        id: ElementId(
+                            18,
+                        ),
+                    },
+
+                    // Replace the loading placeholder with the resolved children
+                    ReplaceWith {
+                        id: ElementId(
+                            3,
+                        ),
+                        m: 3,
+                    },
+                ]
+            );
+
+            dom.wait_for_work().await;
+            let mutations = dom.render_immediate_to_vec().sanitize();
+            assert_eq!(
+                mutations.edits,
+                vec![
+                    CreatePlaceholder {
+                        id: ElementId(
+                            3,
+                        ),
+                    },
+                    ReplaceWith {
+                        id: ElementId(
+                            19,
+                        ),
+                        m: 1,
+                    },
+                    LoadTemplate {
+                        name: "template",
+                        index: 0,
+                        id: ElementId(
+                            19,
+                        ),
+                    },
+                    HydrateText {
+                        path: &[
+                            0,
+                        ],
+                        value: "Goodbye Robot again".to_string(),
+                        id: ElementId(
+                            20,
+                        ),
+                    },
+                    SetAttribute {
+                        name: "id",
+                        ns: None,
+                        value: AttributeValue::Text("title-3".to_string()),
+                        id: ElementId(
+                            19,
+                        ),
+                    },
+                    LoadTemplate {
+                        name: "template",
+                        index: 1,
+                        id: ElementId(
+                            21,
+                        ),
+                    },
+                    HydrateText {
+                        path: &[
+                            0,
+                        ],
+                        value: "The robot says goodbye again".to_string(),
+                        id: ElementId(
+                            22,
+                        ),
+                    },
+                    SetAttribute {
+                        name: "id",
+                        ns: None,
+                        value: AttributeValue::Text("body-3".to_string()),
+                        id: ElementId(
+                            21,
+                        ),
+                    },
+                    LoadTemplate {
+                        name: "template",
+                        index: 2,
+                        id: ElementId(
+                            23,
+                        ),
+                    },
+                    AssignId {
+                        path: &[
+                            0,
+                        ],
+                        id: ElementId(
+                            24,
+                        ),
+                    },
+                    SetAttribute {
+                        name: "id",
+                        ns: None,
+                        value: AttributeValue::Text("children-3".to_string()),
+                        id: ElementId(
+                            23,
+                        ),
+                    },
+                    ReplaceWith {
+                        id: ElementId(
+                            3,
+                        ),
+                        m: 3,
+                    },
+                ]
+            )
+        });
+}

packages/web/src/lib.rs

@@ -24,7 +24,7 @@ use std::{panic, rc::Rc};
 
 pub use crate::cfg::Config;
 use crate::hydration::SuspenseMessage;
-use dioxus_core::{ScopeId, VirtualDom};
+use dioxus_core::VirtualDom;
 use futures_util::{pin_mut, select, FutureExt, StreamExt};
 
 mod cfg;
@@ -100,7 +100,7 @@ pub async fn run(virtual_dom: VirtualDom, web_config: Config) -> ! {
             let server_data = HTMLDataCursor::from_serialized(&hydration_data);
             // If the server serialized an error into the root suspense boundary, throw it into the root scope
             if let Some(error) = server_data.error() {
-                dom.in_runtime(|| ScopeId::APP.throw_error(error));
+                dom.in_runtime(|| dioxus_core::ScopeId::APP.throw_error(error));
             }
             with_server_data(server_data, || {
                 dom.rebuild(&mut websys_dom);