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/core/src/global_context.rs

@@ -1,3 +1,4 @@
+use crate::prelude::SuspenseContext;
 use crate::runtime::RuntimeError;
 use crate::{innerlude::SuspendedFuture, runtime::Runtime, CapturedError, Element, ScopeId, Task};
 use std::future::Future;
@@ -40,6 +41,13 @@ pub fn throw_error(error: impl Into<CapturedError> + 'static) {
         .throw_error(error)
 }
 
+/// Get the suspense context the current scope is in
+pub fn suspense_context() -> Option<SuspenseContext> {
+    current_scope_id()
+        .unwrap_or_else(|e| panic!("{}", e))
+        .suspense_context()
+}
+
 /// Consume context from the current scope
 pub fn try_consume_context<T: 'static + Clone>() -> Option<T> {
     Runtime::with_current_scope(|cx| cx.consume_context::<T>())

packages/core/src/lib.rs

@@ -96,14 +96,14 @@ pub mod prelude {
         fc_to_builder, generation, has_context, needs_update, needs_update_any, parent_scope,
         provide_context, provide_error_boundary, provide_root_context, queue_effect, remove_future,
         schedule_update, schedule_update_any, spawn, spawn_forever, spawn_isomorphic, suspend,
-        throw_error, try_consume_context, use_after_render, use_before_render, use_drop, use_hook,
-        use_hook_with_cleanup, with_owner, AnyValue, Attribute, Callback, Component,
-        ComponentFunction, Context, Element, ErrorBoundary, ErrorContext, Event, EventHandler,
-        Fragment, HasAttributes, IntoAttributeValue, IntoDynNode, OptionStringFromMarker,
-        Properties, ReactiveContext, RenderError, Runtime, RuntimeGuard, ScopeId, ScopeState,
-        SuperFrom, SuperInto, SuspendedFuture, SuspenseBoundary, SuspenseBoundaryProps,
-        SuspenseContext, SuspenseExtension, Task, Template, TemplateAttribute, TemplateNode, VNode,
-        VNodeInner, VirtualDom,
+        suspense_context, throw_error, try_consume_context, use_after_render, use_before_render,
+        use_drop, use_hook, use_hook_with_cleanup, with_owner, AnyValue, Attribute, Callback,
+        Component, ComponentFunction, Context, Element, ErrorBoundary, ErrorContext, Event,
+        EventHandler, Fragment, HasAttributes, IntoAttributeValue, IntoDynNode,
+        OptionStringFromMarker, Properties, ReactiveContext, RenderError, Runtime, RuntimeGuard,
+        ScopeId, ScopeState, SuperFrom, SuperInto, SuspendedFuture, SuspenseBoundary,
+        SuspenseBoundaryProps, SuspenseContext, SuspenseExtension, Task, Template,
+        TemplateAttribute, TemplateNode, VNode, VNodeInner, VirtualDom,
     };
 }
 

packages/core/src/runtime.rs

@@ -116,6 +116,15 @@ impl Runtime {
         result
     }
 
+    /// Run a closure with the rendering flag set to false
+    pub(crate) fn while_not_rendering<T>(&self, f: impl FnOnce() -> T) -> T {
+        let previous = self.rendering.get();
+        self.rendering.set(false);
+        let result = f();
+        self.rendering.set(previous);
+        result
+    }
+
     /// Create a scope context. This slab is synchronized with the scope slab.
     pub(crate) fn create_scope(&self, context: Scope) {
         let id = context.id;

packages/core/src/scope_context.rs

@@ -447,7 +447,12 @@ impl Scope {
         let mut hooks = self.hooks.try_borrow_mut().expect("The hook list is already borrowed: This error is likely caused by trying to use a hook inside a hook which violates the rules of hooks.");
 
         if cur_hook >= hooks.len() {
-            hooks.push(Box::new(initializer()));
+            Runtime::with(|rt| {
+                rt.while_not_rendering(|| {
+                    hooks.push(Box::new(initializer()));
+                });
+            })
+            .unwrap()
         }
 
         self.use_hook_inner::<State>(hooks, cur_hook)
@@ -620,4 +625,9 @@ impl ScopeId {
     pub fn throw_error(self, error: impl Into<CapturedError> + 'static) {
         throw_into(error, self)
     }
+
+    /// Get the suspense context the current scope is in
+    pub fn suspense_context(&self) -> Option<SuspenseContext> {
+        Runtime::with_scope(*self, |cx| cx.suspense_boundary.suspense_context().cloned()).unwrap()
+    }
 }

packages/core/src/suspense/component.rs

@@ -360,7 +360,7 @@ impl SuspenseBoundaryProps {
                     SuspenseContext::downcast_suspense_boundary_from_scope(&dom.runtime, scope_id)
                         .unwrap();
                 suspense_context.take_suspended_nodes();
-                mark_suspense_resolved(dom, scope_id);
+                mark_suspense_resolved(&suspense_context, dom, scope_id);
 
                 nodes_created
             };
@@ -437,6 +437,9 @@ impl SuspenseBoundaryProps {
             let props = Self::downcast_from_props(&mut *scope_state.props).unwrap();
             props.children.clone_from(&children);
             scope_state.last_rendered_node = Some(children);
+
+            // Run any closures that were waiting for the suspense to resolve
+            suspense_context.run_resolved_closures(&dom.runtime);
         })
     }
 
@@ -571,7 +574,7 @@ impl SuspenseBoundaryProps {
                     // Set the last rendered node to the new children
                     dom.scopes[scope_id.0].last_rendered_node = Some(new_children);
 
-                    mark_suspense_resolved(dom, scope_id);
+                    mark_suspense_resolved(&suspense_context, dom, scope_id);
                 }
             }
         })
@@ -579,8 +582,14 @@ impl SuspenseBoundaryProps {
 }
 
 /// Move to a resolved suspense state
-fn mark_suspense_resolved(dom: &mut VirtualDom, scope_id: ScopeId) {
+fn mark_suspense_resolved(
+    suspense_context: &SuspenseContext,
+    dom: &mut VirtualDom,
+    scope_id: ScopeId,
+) {
     dom.resolved_scopes.push(scope_id);
+    // Run any closures that were waiting for the suspense to resolve
+    suspense_context.run_resolved_closures(&dom.runtime);
 }
 
 /// Move from a resolved suspense state to an suspended state
@@ -590,12 +599,12 @@ fn un_resolve_suspense(dom: &mut VirtualDom, scope_id: ScopeId) {
 
 impl SuspenseContext {
     /// Run a closure under a suspense boundary
-    pub fn under_suspense_boundary<O>(&self, runtime: &Runtime, f: impl FnOnce() -> O) -> O {
+    pub(crate) fn under_suspense_boundary<O>(&self, runtime: &Runtime, f: impl FnOnce() -> O) -> O {
         runtime.with_suspense_location(SuspenseLocation::UnderSuspense(self.clone()), f)
     }
 
     /// Run a closure under a suspense placeholder
-    pub fn in_suspense_placeholder<O>(&self, runtime: &Runtime, f: impl FnOnce() -> O) -> O {
+    pub(crate) fn in_suspense_placeholder<O>(&self, runtime: &Runtime, f: impl FnOnce() -> O) -> O {
         runtime.with_suspense_location(SuspenseLocation::InSuspensePlaceholder(self.clone()), f)
     }
 

packages/core/src/suspense/mod.rs

@@ -102,6 +102,7 @@ impl SuspenseContext {
                 id: Cell::new(ScopeId::ROOT),
                 suspended_nodes: Default::default(),
                 frozen: Default::default(),
+                after_suspense_resolved: Default::default(),
             }),
         }
     }
@@ -184,10 +185,26 @@ impl SuspenseContext {
             .find_map(|task| task.suspense_placeholder())
             .map(std::result::Result::Ok)
     }
+
+    /// Run a closure after suspense is resolved
+    pub fn after_suspense_resolved(&self, callback: impl FnOnce() + 'static) {
+        let mut closures = self.inner.after_suspense_resolved.borrow_mut();
+        closures.push(Box::new(callback));
+    }
+
+    /// Run all closures that were queued to run after suspense is resolved
+    pub(crate) fn run_resolved_closures(&self, runtime: &Runtime) {
+        runtime.while_not_rendering(|| {
+            self.inner
+                .after_suspense_resolved
+                .borrow_mut()
+                .drain(..)
+                .for_each(|f| f());
+        })
+    }
 }
 
 /// A boundary that will capture any errors from child components
-#[derive(Debug)]
 pub struct SuspenseBoundaryInner {
     suspended_tasks: RefCell<Vec<SuspendedFuture>>,
     id: Cell<ScopeId>,
@@ -195,6 +212,19 @@ pub struct SuspenseBoundaryInner {
     suspended_nodes: RefCell<Option<VNode>>,
     /// On the server, you can only resolve a suspense boundary once. This is used to track if the suspense boundary has been resolved and if it should be frozen
     frozen: Cell<bool>,
+    /// Closures queued to run after the suspense boundary is resolved
+    after_suspense_resolved: RefCell<Vec<Box<dyn FnOnce()>>>,
+}
+
+impl Debug for SuspenseBoundaryInner {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("SuspenseBoundaryInner")
+            .field("suspended_tasks", &self.suspended_tasks)
+            .field("id", &self.id)
+            .field("suspended_nodes", &self.suspended_nodes)
+            .field("frozen", &self.frozen)
+            .finish()
+    }
 }
 
 /// Provides context methods to [`Result<T, RenderError>`] to show loading indicators for suspended results