fix: correctly handle Suspense with local resources during hydration (

closes leptos-rs#1823) (leptos-rs#1824)

leptos/src/suspense_component.rs

@@ -1,6 +1,8 @@
 use leptos::ViewFn;
 use leptos_dom::{DynChild, HydrationCtx, IntoView};
 use leptos_macro::component;
+#[cfg(feature = "hydrate")]
+use leptos_reactive::SharedContext;
 #[cfg(any(feature = "csr", feature = "hydrate"))]
 use leptos_reactive::SignalGet;
 use leptos_reactive::{
@@ -75,15 +77,35 @@ where
     let owner =
         Owner::current().expect("<Suspense/> created with no reactive owner");
 
+    let current_id = HydrationCtx::next_component();
+
     // provide this SuspenseContext to any resources below it
     // run in a memo so the children are children of this parent
+    #[cfg(not(feature = "hydrate"))]
     let children = create_memo({
         let orig_children = Rc::clone(&orig_children);
         move |_| {
             provide_context(context);
             orig_children().into_view()
         }
     });
+    #[cfg(feature = "hydrate")]
+    let children = create_memo({
+        let orig_children = Rc::clone(&orig_children);
+        move |_| {
+            provide_context(context);
+            if SharedContext::fragment_has_local_resources(
+                &current_id.to_string(),
+            ) {
+                HydrationCtx::with_hydration_off({
+                    let orig_children = Rc::clone(&orig_children);
+                    move || orig_children().into_view()
+                })
+            } else {
+                orig_children().into_view()
+            }
+        }
+    });
 
     // likewise for the fallback
     let fallback = create_memo({
@@ -93,8 +115,6 @@ where
         }
     });
 
-    let current_id = HydrationCtx::next_component();
-
     #[cfg(any(feature = "csr", feature = "hydrate"))]
     let ready = context.ready();
 
@@ -126,6 +146,11 @@ where
                             DynChild::new(move || children_rendered.clone())
                                 .into_view()
                         })
+                    } else if context.has_any_local() {
+                        SharedContext::register_local_fragment(
+                            current_id.to_string(),
+                        );
+                        fallback.get_untracked()
                     }
                     // show the fallback, but also prepare to stream HTML
                     else {

leptos_dom/src/hydration.rs

@@ -250,6 +250,19 @@ impl HydrationCtx {
         value
     }
 
+    #[doc(hidden)]
+    #[cfg(feature = "hydrate")]
+    pub fn with_hydration_off<T>(f: impl FnOnce() -> T) -> T {
+        let prev = IS_HYDRATING.with(|is_hydrating| {
+            let prev = is_hydrating.get();
+            is_hydrating.set(false);
+            prev
+        });
+        let value = f();
+        IS_HYDRATING.with(|is_hydrating| is_hydrating.set(prev));
+        value
+    }
+
     /// Whether the UI is currently in the process of hydrating from the server-sent HTML.
     #[inline(always)]
     pub fn is_hydrating() -> bool {

leptos_dom/src/ssr.rs

@@ -203,6 +203,9 @@ pub fn render_to_stream_with_prefix_undisposed_with_context_and_block_replacemen
         .map(|nonce| format!(" nonce=\"{nonce}\""))
         .unwrap_or_default();
 
+    let local_only = SharedContext::fragments_with_local_resources();
+    let local_only = serde_json::to_string(&local_only).unwrap();
+
     let mut blocking_fragments = FuturesUnordered::new();
     let fragments = FuturesUnordered::new();
 
@@ -226,7 +229,8 @@ pub fn render_to_stream_with_prefix_undisposed_with_context_and_block_replacemen
                 let resolvers = format!(
                     "<script{nonce_str}>__LEPTOS_PENDING_RESOURCES = \
                      {pending_resources};__LEPTOS_RESOLVED_RESOURCES = new \
-                     Map();__LEPTOS_RESOURCE_RESOLVERS = new Map();</script>"
+                     Map();__LEPTOS_RESOURCE_RESOLVERS = new \
+                     Map();__LEPTOS_LOCAL_ONLY = {local_only};</script>"
                 );
 
                 if replace_blocks {

leptos_dom/src/ssr_in_order.rs

@@ -125,6 +125,9 @@ pub fn render_to_stream_in_order_with_prefix_undisposed_with_context(
         .map(|nonce| format!(" nonce=\"{nonce}\""))
         .unwrap_or_default();
 
+    let local_only = SharedContext::fragments_with_local_resources();
+    let local_only = serde_json::to_string(&local_only).unwrap();
+
     let stream = futures::stream::once({
         let nonce_str = nonce_str.clone();
         async move {
@@ -136,6 +139,7 @@ pub fn render_to_stream_in_order_with_prefix_undisposed_with_context(
             __LEPTOS_PENDING_RESOURCES = {pending_resources};
             __LEPTOS_RESOLVED_RESOURCES = new Map();
             __LEPTOS_RESOURCE_RESOLVERS = new Map();
+            __LEPTOS_LOCAL_ONLY = {local_only};
         </script>
       "#
             )

leptos_reactive/src/hydration.rs

@@ -20,6 +20,8 @@ pub struct SharedContext {
     pub resolved_resources: HashMap<ResourceId, String>,
     /// Suspended fragments that have not yet resolved.
     pub pending_fragments: HashMap<String, FragmentData>,
+    /// Suspense fragments that contain only local resources.
+    pub fragments_with_local_resources: HashSet<String>,
     #[cfg(feature = "experimental-islands")]
     pub no_hydrate: bool,
     #[cfg(all(feature = "hydrate", feature = "experimental-islands"))]
@@ -108,6 +110,7 @@ impl SharedContext {
                     is_ready: Some(Box::pin(async move {
                         rx3.next().await;
                     })),
+                    local_only: context.has_local_only(),
                 },
             );
         })
@@ -185,6 +188,44 @@ impl SharedContext {
             });
         }
     }
+
+    #[cfg_attr(
+        any(debug_assertions, features = "ssr"),
+        instrument(level = "trace", skip_all,)
+    )]
+    pub fn fragment_has_local_resources(fragment: &str) -> bool {
+        with_runtime(|runtime| {
+            let mut shared_context = runtime.shared_context.borrow_mut();
+            shared_context
+                .fragments_with_local_resources
+                .remove(fragment)
+        })
+        .unwrap_or_default()
+    }
+
+    #[cfg_attr(
+        any(debug_assertions, features = "ssr"),
+        instrument(level = "trace", skip_all,)
+    )]
+    pub fn fragments_with_local_resources() -> HashSet<String> {
+        with_runtime(|runtime| {
+            let mut shared_context = runtime.shared_context.borrow_mut();
+            std::mem::take(&mut shared_context.fragments_with_local_resources)
+        })
+        .unwrap_or_default()
+    }
+
+    #[cfg_attr(
+        any(debug_assertions, features = "ssr"),
+        instrument(level = "trace", skip_all,)
+    )]
+    pub fn register_local_fragment(key: String) {
+        with_runtime(|runtime| {
+            let mut shared_context = runtime.shared_context.borrow_mut();
+            shared_context.fragments_with_local_resources.insert(key);
+        })
+        .unwrap_or_default()
+    }
 }
 
 /// Represents its pending `<Suspense/>` fragment.
@@ -197,6 +238,8 @@ pub struct FragmentData {
     pub should_block: bool,
     /// Future that will resolve when the fragment is ready.
     pub is_ready: Option<PinnedFuture<()>>,
+    /// Whether the fragment contains only local resources.
+    pub local_only: bool,
 }
 
 impl std::fmt::Debug for SharedContext {
@@ -229,6 +272,17 @@ impl Default for SharedContext {
                     serde_wasm_bindgen::from_value(pr).map_err(|_| ())
                 })
                 .unwrap();
+            let fragments_with_local_resources = js_sys::Reflect::get(
+                &web_sys::window().unwrap(),
+                &wasm_bindgen::JsValue::from_str("__LEPTOS_LOCAL_ONLY"),
+            );
+            let fragments_with_local_resources: HashSet<String> =
+                fragments_with_local_resources
+                    .map_err(|_| ())
+                    .and_then(|pr| {
+                        serde_wasm_bindgen::from_value(pr).map_err(|_| ())
+                    })
+                    .unwrap_or_default();
 
             let resolved_resources = js_sys::Reflect::get(
                 &web_sys::window().unwrap(),
@@ -244,6 +298,7 @@ impl Default for SharedContext {
                 //events: Default::default(),
                 pending_resources,
                 resolved_resources,
+                fragments_with_local_resources,
                 pending_fragments: Default::default(),
                 #[cfg(feature = "experimental-islands")]
                 no_hydrate: true,
@@ -262,6 +317,7 @@ impl Default for SharedContext {
                 pending_resources: Default::default(),
                 resolved_resources: Default::default(),
                 pending_fragments: Default::default(),
+                fragments_with_local_resources: Default::default(),
                 #[cfg(feature = "experimental-islands")]
                 no_hydrate: true,
                 #[cfg(all(

leptos_reactive/src/resource.rs

@@ -5,10 +5,11 @@ use crate::SpecialNonReactiveZone;
 use crate::{
     create_isomorphic_effect, create_memo, create_render_effect, create_signal,
     queue_microtask, runtime::with_runtime, serialization::Serializable,
-    signal_prelude::format_signal_warning, spawn::spawn_local, use_context,
-    GlobalSuspenseContext, Memo, ReadSignal, ScopeProperty, Signal,
-    SignalDispose, SignalGet, SignalGetUntracked, SignalSet, SignalUpdate,
-    SignalWith, SuspenseContext, WriteSignal,
+    signal_prelude::format_signal_warning, spawn::spawn_local,
+    suspense::LocalStatus, use_context, GlobalSuspenseContext, Memo,
+    ReadSignal, ScopeProperty, Signal, SignalDispose, SignalGet,
+    SignalGetUntracked, SignalSet, SignalUpdate, SignalWith, SuspenseContext,
+    WriteSignal,
 };
 use std::{
     any::Any,
@@ -1178,7 +1179,16 @@ where
         let serializable = self.serializable;
         if let Some(suspense_cx) = &suspense_cx {
             if serializable != ResourceSerialization::Local {
-                suspense_cx.has_local_only.set_value(false);
+                suspense_cx.local_status.update_value(|status| {
+                    *status = Some(match status {
+                        None => LocalStatus::SerializableOnly,
+                        Some(LocalStatus::LocalOnly) => LocalStatus::LocalOnly,
+                        Some(LocalStatus::Mixed) => LocalStatus::Mixed,
+                        Some(LocalStatus::SerializableOnly) => {
+                            LocalStatus::SerializableOnly
+                        }
+                    });
+                });
             }
         } else {
             #[cfg(not(all(feature = "hydrate", debug_assertions)))]

leptos_reactive/src/suspense.rs

@@ -16,10 +16,17 @@ pub struct SuspenseContext {
     pub pending_resources: ReadSignal<usize>,
     set_pending_resources: WriteSignal<usize>,
     pub(crate) pending_serializable_resources: RwSignal<usize>,
-    pub(crate) has_local_only: StoredValue<bool>,
+    pub(crate) local_status: StoredValue<Option<LocalStatus>>,
     pub(crate) should_block: StoredValue<bool>,
 }
 
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+pub(crate) enum LocalStatus {
+    LocalOnly,
+    Mixed,
+    SerializableOnly,
+}
+
 /// A single, global suspense context that will be checked when resources
 /// are read. This won’t be “blocked” by lower suspense components. This is
 /// useful for e.g., holding route transitions.
@@ -54,7 +61,15 @@ impl SuspenseContext {
     /// Whether the suspense contains local resources at this moment,
     /// and therefore can't be serialized
     pub fn has_local_only(&self) -> bool {
-        self.has_local_only.get_value()
+        matches!(self.local_status.get_value(), Some(LocalStatus::LocalOnly))
+    }
+
+    /// Whether the suspense contains any local resources at this moment.
+    pub fn has_any_local(&self) -> bool {
+        matches!(
+            self.local_status.get_value(),
+            Some(LocalStatus::LocalOnly) | Some(LocalStatus::Mixed)
+        )
     }
 
     /// Whether any blocking resources are read under this suspense context,
@@ -102,13 +117,13 @@ impl SuspenseContext {
     pub fn new() -> Self {
         let (pending_resources, set_pending_resources) = create_signal(0);
         let pending_serializable_resources = create_rw_signal(0);
-        let has_local_only = store_value(true);
+        let local_status = store_value(None);
         let should_block = store_value(false);
         Self {
             pending_resources,
             set_pending_resources,
             pending_serializable_resources,
-            has_local_only,
+            local_status,
             should_block,
         }
     }
@@ -117,11 +132,29 @@ impl SuspenseContext {
     pub fn increment(&self, serializable: bool) {
         let setter = self.set_pending_resources;
         let serializable_resources = self.pending_serializable_resources;
-        let has_local_only = self.has_local_only;
+        let local_status = self.local_status;
         setter.update(|n| *n += 1);
         if serializable {
             serializable_resources.update(|n| *n += 1);
-            has_local_only.set_value(false);
+            local_status.update_value(|status| {
+                *status = Some(match status {
+                    None => LocalStatus::SerializableOnly,
+                    Some(LocalStatus::LocalOnly) => LocalStatus::LocalOnly,
+                    Some(LocalStatus::Mixed) => LocalStatus::Mixed,
+                    Some(LocalStatus::SerializableOnly) => {
+                        LocalStatus::SerializableOnly
+                    }
+                });
+            });
+        } else {
+            local_status.update_value(|status| {
+                *status = Some(match status {
+                    None => LocalStatus::LocalOnly,
+                    Some(LocalStatus::LocalOnly) => LocalStatus::LocalOnly,
+                    Some(LocalStatus::Mixed) => LocalStatus::Mixed,
+                    Some(LocalStatus::SerializableOnly) => LocalStatus::Mixed,
+                });
+            });
         }
     }