work on Axum integration and on error boundaries

hydration_context/src/csr.rs

@@ -65,4 +65,7 @@ impl SharedContext for CsrSharedContext {
         _error: any_error::Error,
     ) {
     }
+
+    #[inline(always)]
+    fn seal_errors(&self, boundary_id: &SerializedDataId) {}
 }

hydration_context/src/hydrate.rs

@@ -141,6 +141,9 @@ impl SharedContext for HydrateSharedContext {
     ) {
     }
 
+    #[inline(always)]
+    fn seal_errors(&self, boundary_id: &SerializedDataId) {}
+
     fn take_errors(&self) -> Vec<(SerializedDataId, ErrorId, Error)> {
         self.errors.clone()
     }

hydration_context/src/lib.rs

@@ -104,6 +104,12 @@ pub trait SharedContext: Debug {
     /// Returns the set of errors that have been registered with a particular boundary.
     fn errors(&self, boundary_id: &SerializedDataId) -> Vec<(ErrorId, Error)>;
 
+    /// "Seals" an error boundary, preventing further errors from being registered for it.
+    ///
+    /// This can be used in streaming SSR scenarios in which the final state of the error boundary
+    /// can only be known after the initial state is hydrated.
+    fn seal_errors(&self, boundary_id: &SerializedDataId);
+
     /// Registers an error with the context to be shared from server to client.
     fn register_error(
         &self,

hydration_context/src/ssr.rs

@@ -7,6 +7,7 @@ use futures::{
 };
 use or_poisoned::OrPoisoned;
 use std::{
+    collections::{HashMap, HashSet},
     fmt::{Debug, Write},
     mem,
     sync::{
@@ -23,6 +24,7 @@ pub struct SsrSharedContext {
     sync_buf: RwLock<Vec<ResolvedData>>,
     async_buf: RwLock<Vec<(SerializedDataId, PinnedFuture<String>)>>,
     errors: Arc<RwLock<Vec<(SerializedDataId, ErrorId, Error)>>>,
+    sealed_error_boundaries: Arc<RwLock<HashSet<SerializedDataId>>>,
 }
 
 impl SsrSharedContext {
@@ -110,19 +112,23 @@ impl SharedContext for SsrSharedContext {
             .into_iter()
             .map(|(id, data)| {
                 let errors = Arc::clone(&self.errors);
+                let sealed = Arc::clone(&self.sealed_error_boundaries);
                 async move {
                     let data = data.await;
                     let data = data.replace('<', "\\u003c");
                     let mut val =
                         format!("__RESOLVED_RESOURCES[{}] = {:?};", id.0, data);
+                    let sealed = sealed.read().or_poisoned();
                     for error in mem::take(&mut *errors.write().or_poisoned()) {
-                        write!(
-                            val,
-                            "__SERIALIZED_ERRORS.push([{}, {}, {:?}]);",
-                            error.0 .0,
-                            error.1,
-                            error.2.to_string()
-                        );
+                        if !sealed.contains(&error.0) {
+                            write!(
+                                val,
+                                "__SERIALIZED_ERRORS.push([{}, {}, {:?}]);",
+                                error.0 .0,
+                                error.1,
+                                error.2.to_string()
+                            );
+                        }
                     }
                     val
                 }
@@ -181,6 +187,13 @@ impl SharedContext for SsrSharedContext {
     fn take_errors(&self) -> Vec<(SerializedDataId, ErrorId, Error)> {
         mem::take(&mut *self.errors.write().or_poisoned())
     }
+
+    fn seal_errors(&self, boundary_id: &SerializedDataId) {
+        self.sealed_error_boundaries
+            .write()
+            .or_poisoned()
+            .insert(boundary_id.clone());
+    }
 }
 
 #[derive(Debug)]