Wait for the suspense boundary above the router to resolve before sending the first streaming chunk (#3891)

* Change the suspense streaming context after suspense is resolved in the router

* Wait for the suspense context before streaming the first chunk in ssr

* forward more response parts

* warn

* extract the hydration format into a crate

* Fix hydration when no data is sent

* forward web and server features

* Fix suspense that starts on the client

* fix use_server_future on the client

* Add documentation for fullstack hooks and protocol

* downgrade rust edition

* fix clippy

* fix tests

* fix clippy

* Fix pure web compile

* undo changes to the fullstack router example

* fix formatting

packages/fullstack/src/render.rs

@@ -1,11 +1,13 @@
 //! A shared pool of renderers for efficient server side rendering.
 use crate::document::ServerDocument;
-use crate::html_storage::serialize::SerializedHydrationData;
 use crate::streaming::{Mount, StreamingRenderer};
 use dioxus_cli_config::base_path;
+use dioxus_fullstack_hooks::{StreamingContext, StreamingStatus};
+use dioxus_fullstack_protocol::{HydrationContext, SerializedHydrationData};
 use dioxus_interpreter_js::INITIALIZE_STREAMING_JS;
 use dioxus_isrg::{CachedRender, IncrementalRendererError, RenderFreshness};
 use dioxus_lib::document::Document;
+use dioxus_lib::prelude::dioxus_core::DynamicNode;
 use dioxus_router::prelude::ParseRouteError;
 use dioxus_ssr::Renderer;
 use futures_channel::mpsc::Sender;
@@ -51,6 +53,10 @@ where
     }
 }
 
+fn in_root_scope<T>(virtual_dom: &VirtualDom, f: impl FnOnce() -> T) -> T {
+    virtual_dom.in_runtime(|| ScopeId::ROOT.in_runtime(f))
+}
+
 /// Errors that can occur during server side rendering before the initial chunk is sent down
 pub enum SSRError {
     /// An error from the incremental renderer. This should result in a 500 code
@@ -180,7 +186,7 @@ impl SsrRendererPool {
         let myself = self.clone();
         let streaming_mode = cfg.streaming_mode;
 
-        let join_handle = spawn_platform(move || async move {
+        let create_render_future = move || async move {
             let mut virtual_dom = virtual_dom_factory();
             let document = std::rc::Rc::new(crate::document::server::ServerDocument::default());
             virtual_dom.provide_root_context(document.clone());
@@ -196,27 +202,47 @@ impl SsrRendererPool {
             } else {
                 history = dioxus_history::MemoryHistory::with_initial_path(&route);
             }
+            let streaming_context = in_root_scope(&virtual_dom, StreamingContext::new);
             virtual_dom.provide_root_context(Rc::new(history) as Rc<dyn dioxus_history::History>);
             virtual_dom.provide_root_context(document.clone() as std::rc::Rc<dyn Document>);
+            virtual_dom.provide_root_context(streaming_context);
 
-            // rebuild the virtual dom, which may call server_context()
-            with_server_context(server_context.clone(), || virtual_dom.rebuild_in_place());
+            // rebuild the virtual dom
+            virtual_dom.rebuild_in_place();
 
             // If streaming is disabled, wait for the virtual dom to finish all suspense work
             // before rendering anything
             if streaming_mode == StreamingMode::Disabled {
-                ProvideServerContext::new(virtual_dom.wait_for_suspense(), server_context.clone())
-                    .await
+                virtual_dom.wait_for_suspense().await;
             }
+            // Otherwise, just wait for the streaming context to signal the initial chunk is ready
+            else {
+                loop {
+                    // Check if the router has finished and set the streaming context to finished
+                    let streaming_context_finished =
+                        in_root_scope(&virtual_dom, || streaming_context.current_status())
+                            == StreamingStatus::InitialChunkCommitted;
+                    // Or if this app isn't using the router and has finished suspense
+                    let suspense_finished = !virtual_dom.suspended_tasks_remaining();
+                    if streaming_context_finished || suspense_finished {
+                        break;
+                    }
+
+                    // Wait for new async work that runs during suspense (mainly use_server_futures)
+                    virtual_dom.wait_for_suspense_work().await;
+
+                    // Do that async work
+                    virtual_dom.render_suspense_immediate().await;
+                }
+            }
+
             // check if there are any errors
-            let errors = with_server_context(server_context.clone(), || {
-                virtual_dom.in_runtime(|| {
-                    let error_context: ErrorContext = ScopeId::APP
-                        .consume_context()
-                        .expect("The root should be under an error boundary");
-                    let errors = error_context.errors();
-                    errors.to_vec()
-                })
+            let errors = virtual_dom.in_runtime(|| {
+                let error_context: ErrorContext = ScopeId::APP
+                    .consume_context()
+                    .expect("The root should be under an error boundary");
+                let errors = error_context.errors();
+                errors.to_vec()
             });
             if errors.is_empty() {
                 // If routing was successful, we can return a 200 status and render into the stream
@@ -282,16 +308,8 @@ impl SsrRendererPool {
 
             // After the initial render, we need to resolve suspense
             while virtual_dom.suspended_tasks_remaining() {
-                ProvideServerContext::new(
-                    virtual_dom.wait_for_suspense_work(),
-                    server_context.clone(),
-                )
-                .await;
-                let resolved_suspense_nodes = ProvideServerContext::new(
-                    virtual_dom.render_suspense_immediate(),
-                    server_context.clone(),
-                )
-                .await;
+                virtual_dom.wait_for_suspense_work().await;
+                let resolved_suspense_nodes = virtual_dom.render_suspense_immediate().await;
 
                 // Just rerender the resolved nodes
                 for scope in resolved_suspense_nodes {
@@ -369,6 +387,10 @@ impl SsrRendererPool {
 
             renderer.reset_render_components();
             myself.renderers.write().unwrap().push(renderer);
+        };
+
+        let join_handle = spawn_platform(move || {
+            ProvideServerContext::new(create_render_future(), server_context)
         });
 
         // Wait for the initial result which determines the status code
@@ -464,13 +486,83 @@ fn start_capturing_errors(suspense_scope: ScopeId) {
 fn serialize_server_data(virtual_dom: &VirtualDom, scope: ScopeId) -> SerializedHydrationData {
     // After we replace the placeholder in the dom with javascript, we need to send down the resolved data so that the client can hydrate the node
     // Extract any data we serialized for hydration (from server futures)
-    let html_data =
-        crate::html_storage::HTMLData::extract_from_suspense_boundary(virtual_dom, scope);
+    let html_data = extract_from_suspense_boundary(virtual_dom, scope);
 
     // serialize the server state into a base64 string
     html_data.serialized()
 }
 
+/// Walks through the suspense boundary in a depth first order and extracts the data from the context API.
+/// We use depth first order instead of relying on the order the hooks are called in because during suspense on the server, the order that futures are run in may be non deterministic.
+pub(crate) fn extract_from_suspense_boundary(
+    vdom: &VirtualDom,
+    scope: ScopeId,
+) -> HydrationContext {
+    let data = HydrationContext::default();
+    serialize_errors(&data, vdom, scope);
+    take_from_scope(&data, vdom, scope);
+    data
+}
+
+/// Get the errors from the suspense boundary
+fn serialize_errors(context: &HydrationContext, vdom: &VirtualDom, scope: ScopeId) {
+    // If there is an error boundary on the suspense boundary, grab the error from the context API
+    // and throw it on the client so that it bubbles up to the nearest error boundary
+    let error = vdom.in_runtime(|| {
+        scope
+            .consume_context::<ErrorContext>()
+            .and_then(|error_context| error_context.errors().first().cloned())
+    });
+    context
+        .error_entry()
+        .insert(&error, std::panic::Location::caller());
+}
+
+fn take_from_scope(context: &HydrationContext, vdom: &VirtualDom, scope: ScopeId) {
+    vdom.in_runtime(|| {
+        scope.in_runtime(|| {
+            // Grab any serializable server context from this scope
+            let other: Option<HydrationContext> = has_context();
+            if let Some(other) = other {
+                context.extend(&other);
+            }
+        });
+    });
+
+    // then continue to any children
+    if let Some(scope) = vdom.get_scope(scope) {
+        // If this is a suspense boundary, move into the children first (even if they are suspended) because that will be run first on the client
+        if let Some(suspense_boundary) =
+            SuspenseContext::downcast_suspense_boundary_from_scope(&vdom.runtime(), scope.id())
+        {
+            if let Some(node) = suspense_boundary.suspended_nodes() {
+                take_from_vnode(context, vdom, &node);
+            }
+        }
+        if let Some(node) = scope.try_root_node() {
+            take_from_vnode(context, vdom, node);
+        }
+    }
+}
+
+fn take_from_vnode(context: &HydrationContext, vdom: &VirtualDom, vnode: &VNode) {
+    for (dynamic_node_index, dyn_node) in vnode.dynamic_nodes.iter().enumerate() {
+        match dyn_node {
+            DynamicNode::Component(comp) => {
+                if let Some(scope) = comp.mounted_scope(dynamic_node_index, vnode, vdom) {
+                    take_from_scope(context, vdom, scope.id());
+                }
+            }
+            DynamicNode::Fragment(nodes) => {
+                for node in nodes {
+                    take_from_vnode(context, vdom, node);
+                }
+            }
+            _ => {}
+        }
+    }
+}
+
 /// State used in server side rendering. This utilizes a pool of [`dioxus_ssr::Renderer`]s to cache static templates between renders.
 #[derive(Clone)]
 pub struct SSRState {