Add proper handling of node storages to ParallelFlow

Adds merging of child storages into parent storage

Author: HANDZCZ

src/flows/parallel_flow/chain_run/poll.rs

@@ -12,7 +12,7 @@ pub trait ChainPollParallel<Output>: Send {
         self: Pin<&mut Self>,
         cx: &mut Context<'_>,
         tail_ready: bool,
-        storage: &mut Storage,
+        storage_acc: &mut Vec<Storage>,
     ) -> Poll<Output>;
 }
 
@@ -28,19 +28,19 @@ where
         self: Pin<&mut Self>,
         cx: &mut Context<'_>,
         tail_ready: bool,
-        storage: &mut Storage,
+        storage_acc: &mut Vec<Storage>,
     ) -> Poll<Result<(HeadOutput, TailOutput), Error>> {
         let (head, tail) = unsafe { self.get_unchecked_mut() };
         let (head, mut tail) = unsafe { (Pin::new_unchecked(head), Pin::new_unchecked(tail)) };
         let tail_ready = tail.as_mut().poll(cx).is_ready() && tail_ready;
 
-        let Poll::Ready(res) = ChainPollParallel::poll(head, cx, tail_ready, storage) else {
+        let Poll::Ready(res) = ChainPollParallel::poll(head, cx, tail_ready, storage_acc) else {
             return Poll::Pending;
         };
         match res {
             Ok(head_out) => match tail.take_output().unwrap() {
                 Ok((tail_out, node_storage)) => {
-                    // TODO: merge storage
+                    storage_acc.push(node_storage);
                     Poll::Ready(Ok((head_out, tail_out)))
                 }
                 Err(e) => Poll::Ready(Err(e)),
@@ -60,13 +60,13 @@ where
         self: Pin<&mut Self>,
         cx: &mut Context<'_>,
         tail_ready: bool,
-        storage: &mut Storage,
+        storage_acc: &mut Vec<Storage>,
     ) -> Poll<Result<(HeadOutput,), Error>> {
         let mut head = unsafe { Pin::new_unchecked(&mut self.get_unchecked_mut().0) };
         if head.as_mut().poll(cx).is_ready() && tail_ready {
             match head.take_output().unwrap() {
                 Ok((output, node_storage)) => {
-                    // TODO: merge storage
+                    storage_acc.push(node_storage);
                     Poll::Ready(Ok((output,)))
                 }
                 Err(e) => Poll::Ready(Err(e)),

src/flows/parallel_flow/chain_run/run.rs

@@ -20,7 +20,12 @@ where
 {
     async fn run_with_storage(&self, input: Input, storage: &mut Storage) -> Result<Output, Error> {
         let fut_chain = self.spawn_with_storage(input, storage.new_gen());
+        let mut storage_acc = Vec::with_capacity(U::NUM_FUTURES);
         let mut fut_chain = pin!(fut_chain);
-        poll_fn(move |cx| ChainPollParallel::poll(fut_chain.as_mut(), cx, true, storage)).await
+        let res =
+            poll_fn(|cx| ChainPollParallel::poll(fut_chain.as_mut(), cx, true, &mut storage_acc))
+                .await;
+        storage.merge(&mut storage_acc);
+        res
     }
 }

src/flows/parallel_flow/chain_run/spawn.rs

@@ -8,6 +8,7 @@ use crate::{
 
 pub trait ChainSpawn<Input, Error, HeadOut, T> {
     type ChainOut;
+    const NUM_FUTURES: usize;
 
     fn spawn_with_storage(
         &self,
@@ -45,6 +46,7 @@ where
     Error: Send,
 {
     type ChainOut = Result<(HeadOut, NodeOutputStruct<TailNodeOutType>), Error>;
+    const NUM_FUTURES: usize = Head::NUM_FUTURES + 1;
 
     fn spawn_with_storage(
         &self,
@@ -86,6 +88,7 @@ where
     Error: Send,
 {
     type ChainOut = Result<(NodeOutputStruct<HeadNodeOutType>,), Error>;
+    const NUM_FUTURES: usize = 1;
 
     fn spawn_with_storage(
         &self,