chore: expose internals of SerializedDataId and SsrSharedContext to allow creating custom hydration contexts (#3145)

hydration_context/src/lib.rs

@@ -44,6 +44,18 @@ pub type PinnedStream<T> = Pin<Box<dyn Stream<Item = T> + Send + Sync>>;
 /// from the server to the client.
 pub struct SerializedDataId(usize);
 
+impl SerializedDataId {
+    /// Create a new instance of [`SerializedDataId`].
+    pub fn new(id: usize) -> Self {
+        SerializedDataId(id)
+    }
+
+    /// Consume into the inner usize identifier.
+    pub fn into_inner(self) -> usize {
+        self.0
+    }
+}
+
 impl From<SerializedDataId> for ErrorId {
     fn from(value: SerializedDataId) -> Self {
         value.0.into()

hydration_context/src/ssr.rs

@@ -58,6 +58,27 @@ impl SsrSharedContext {
             ..Default::default()
         }
     }
+
+    /// Consume the data buffers, awaiting all async resources,
+    /// returning both sync and async buffers.
+    /// Useful to implement custom hydration contexts.
+    ///
+    /// WARNING: this will clear the internal buffers, it should only be called once.
+    /// A second call would return an empty `vec![]`.
+    pub async fn consume_buffers(&self) -> Vec<(SerializedDataId, String)> {
+        let sync_data = mem::take(&mut *self.sync_buf.write().or_poisoned());
+        let async_data = mem::take(&mut *self.async_buf.write().or_poisoned());
+
+        let mut all_data = Vec::new();
+        for resolved in sync_data {
+            all_data.push((resolved.0, resolved.1));
+        }
+        for (id, fut) in async_data {
+            let data = fut.await;
+            all_data.push((id, data));
+        }
+        all_data
+    }
 }
 
 impl Debug for SsrSharedContext {