Add tests and fixes for streams

Signed-off-by: Michael X. Grey <grey@openrobotics.org>

Author: mxgrey

src/callback.rs

@@ -179,7 +179,7 @@ struct BlockingCallbackSystem<Request, Response, Streams: StreamPack> {
     initialized: bool,
 }
 
-impl<Request, Response, Streams> CallbackTrait<Request, Response, ()> for BlockingCallbackSystem<Request, Response, Streams>
+impl<Request, Response, Streams> CallbackTrait<Request, Response, Streams> for BlockingCallbackSystem<Request, Response, Streams>
 where
     Request: 'static + Send + Sync,
     Response: 'static + Send + Sync,
@@ -263,15 +263,16 @@ pub trait AsCallback<M> {
     fn as_callback(self) -> Callback<Self::Request, Self::Response, Self::Streams>;
 }
 
-impl<Request, Response, M, Sys> AsCallback<BlockingCallbackMarker<(Request, Response, M)>> for Sys
+impl<Request, Response, Streams, M, Sys> AsCallback<BlockingCallbackMarker<(Request, Response, Streams, M)>> for Sys
 where
-    Sys: IntoSystem<BlockingCallback<Request>, Response, M>,
+    Sys: IntoSystem<BlockingCallback<Request, Streams>, Response, M>,
     Request: 'static + Send + Sync,
     Response: 'static + Send + Sync,
+    Streams: StreamPack,
 {
     type Request = Request;
     type Response = Response;
-    type Streams = ();
+    type Streams = Streams;
 
     fn as_callback(self) -> Callback<Self::Request, Self::Response, Self::Streams> {
         Callback::new(BlockingCallbackSystem {

src/channel.rs

@@ -31,8 +31,11 @@ use crate::{
 
 #[derive(Clone)]
 pub struct Channel<Streams: StreamPack = ()> {
+    /// Stream channels that will let you send stream information. This will
+    /// usually be a [`StreamChannel`] or a (possibly nested) tuple of
+    /// `StreamChannel`s, whichever matches the [`StreamPack`] description.
+    pub streams: Streams::Channel,
     inner: Arc<InnerChannel>,
-    streams: Streams::Channel,
     _ignore: std::marker::PhantomData<Streams>,
 }
 
@@ -67,13 +70,6 @@ impl<Streams: StreamPack> Channel<Streams> {
 
         promise
     }
-
-    /// Get stream channels that will let you send stream information. This will
-    /// usually be one [`StreamChannel`] or a (possibly nested) tuple of
-    /// `StreamChannel`s, whichever matches the [`StreamPack`] description.
-    pub fn streams(&self) -> &Streams::Channel {
-        &self.streams
-    }
 }
 
 #[derive(Clone)]

src/flush.rs

@@ -125,6 +125,8 @@ fn collect_from_channels(
         (item)(world, roster);
     }
 
+    roster.process_deferals();
+
     world.get_resource_or_insert_with(|| ServiceLifecycleChannel::new());
     world.resource_scope::<ServiceLifecycleChannel, ()>(|world, lifecycles| {
         // Clean up the dangling requests of any services that have been despawned.

src/impulse/taken.rs

@@ -92,7 +92,6 @@ impl<T: 'static + Send + Sync> Impulsive for TakenStream<T> {
         let Input { data, .. } = source_mut.take_input::<T>()?;
         let stream = source_mut.get::<TakenStream<T>>().or_broken()?;
         stream.sender.send(data).ok();
-        source_mut.despawn_recursive();
         Ok(())
     }
 }

src/input.rs

@@ -84,13 +84,26 @@ impl<T: 'static + Send + Sync> InputBundle<T> {
 }
 
 pub trait ManageInput {
+    /// Give an input to this node. The node will be queued up to immediately
+    /// process the input.
     fn give_input<T: 'static + Send + Sync>(
         &mut self,
         session: Entity,
         data: T,
         roster: &mut OperationRoster,
     ) -> Result<(), OperationError>;
 
+    /// Same as [`Self::give_input`], but the wakeup for this node will be
+    /// deferred until after the [`ChannelQueue`](crate::ChannelQueue) is flushed.
+    /// This is used for async output to ensure that all async operations are
+    /// finished being processed before the final output gets processed.
+    fn defer_input<T: 'static + Send + Sync>(
+        &mut self,
+        session: Entity,
+        data: T,
+        roster: &mut OperationRoster,
+    ) -> Result<(), OperationError>;
+
     /// Give an input to this node without flagging it in the roster. This
     /// should not generally be used. It's only for special cases where we know
     /// the node will be manually run after giving this input. It's marked
@@ -181,6 +194,17 @@ impl<'w> ManageInput for EntityMut<'w> {
         Ok(())
     }
 
+    fn defer_input<T: 'static + Send + Sync>(
+        &mut self,
+        session: Entity,
+        data: T,
+        roster: &mut OperationRoster,
+    ) -> Result<(), OperationError> {
+        unsafe { self.sneak_input(session, data)?; }
+        roster.defer(self.id());
+        Ok(())
+    }
+
     unsafe fn sneak_input<T: 'static + Send + Sync>(
         &mut self,
         session: Entity,

src/lib.rs

@@ -128,7 +128,7 @@ pub struct BlockingService<Request, Streams: StreamPack = ()> {
 }
 
 /// Use this to reduce bracket noise when you need `In<BlockingService<R>>`.
-pub type InBlockingService<Request, Streams = ()> = In<BlockingService<Request, Streams>>;
+pub type BlockingServiceInput<Request, Streams = ()> = In<BlockingService<Request, Streams>>;
 
 /// Use AsyncService to indicate that your system is an async [`Service`]. Being
 /// async means it must return a [`Future<Output=Response>`](std::future::Future)
@@ -152,8 +152,8 @@ pub struct AsyncService<Request, Streams: StreamPack = ()> {
     pub session: Entity,
 }
 
-/// Use this to reduce backet noise when you need `In<AsyncService<R, S>>`.
-pub type InAsyncService<Request, Streams = ()> = In<AsyncService<Request, Streams>>;
+/// Use this to reduce backet noise when you need `In<`[`AsyncService<R, S>`]`>`.
+pub type AsyncServiceInput<Request, Streams = ()> = In<AsyncService<Request, Streams>>;
 
 /// Use BlockingCallback to indicate that your system is meant to define a
 /// blocking [`Callback`]. Callbacks are different from services because they are
@@ -174,6 +174,9 @@ pub struct BlockingCallback<Request, Streams: StreamPack = ()> {
     pub session: Entity,
 }
 
+/// Use this to reduce bracket noise when you need `In<`[`BlockingCallback<R, S>`]`>`.
+pub type BlockingCallbackInput<Request, Streams = ()> = In<BlockingCallback<Request, Streams>>;
+
 /// Use AsyncCallback to indicate that your system or function is meant to define
 /// an async [`Callback`]. An async callback is not associated with any entity,
 /// and it must return a [`Future<Output=Response>`](std::future::Future) that
@@ -192,6 +195,9 @@ pub struct AsyncCallback<Request, Streams: StreamPack = ()> {
     pub session: Entity,
 }
 
+/// Use this to reduce bracket noise when you need `In<`[`AsyncCallback<R, S>`]`>`.
+pub type AsyncCallbackInput<Request, Streams = ()> = In<AsyncCallback<Request, Streams>>;
+
 /// Use `BlockingMap`` to indicate that your function is a blocking map. A map
 /// is not associated with any entity, and it cannot be a Bevy System. These
 /// restrictions allow them to be processed more efficiently.
@@ -207,6 +213,9 @@ pub struct BlockingMap<Request, Streams: StreamPack = ()> {
     pub session: Entity,
 }
 
+/// Use this to reduce the bracket noise when you need `In<`[`BlockingMap<R, S>`]`>`.
+pub type BlockingMapInput<Request, Streams = ()> = In<BlockingMap<Request, Streams>>;
+
 /// Use AsyncMap to indicate that your function is an async map. A Map is not
 /// associated with any entity, and it cannot be a Bevy System. These
 /// restrictions allow them to be processed more efficiently.
@@ -226,3 +235,6 @@ pub struct AsyncMap<Request, Streams: StreamPack = ()> {
     /// The unique session ID for the workflow
     pub session: Entity,
 }
+
+/// Use this to reduce bracket noise when you need `In<`[`AsyncMap<R, S>`]`>`.
+pub type AsyncMapInput<Request, Streams = ()> = In<AsyncMap<Request, Streams>>;

src/operation.rs

@@ -185,6 +185,10 @@ pub struct OperationRoster {
     /// Tasks that should be awoken. If the task is already despawned, then
     /// it should not be considered an error.
     pub(crate) awake: VecDeque<Entity>,
+    /// Operation sources that should be triggered after the next ChannelQueue
+    /// flush. This is for the final outputs of polled tasks, to make sure their
+    /// stream data gets flushed before their final output is flushed.
+    pub(crate) deferred_queue: VecDeque<Entity>,
     /// Operation sources that should be cancelled
     pub(crate) cancel: VecDeque<Cancel>,
     /// Async services that should pull their next item
@@ -208,6 +212,10 @@ impl OperationRoster {
         self.awake.push_back(source);
     }
 
+    pub fn defer(&mut self, source: Entity) {
+        self.deferred_queue.push_back(source);
+    }
+
     pub fn cancel(&mut self, source: Cancel) {
         self.cancel.push_back(source);
     }
@@ -227,6 +235,7 @@ impl OperationRoster {
     pub fn is_empty(&self) -> bool {
         self.queue.is_empty()
         && self.awake.is_empty()
+        && self.deferred_queue.is_empty()
         && self.cancel.is_empty()
         && self.unblock.is_empty()
         && self.disposed.is_empty()
@@ -246,6 +255,13 @@ impl OperationRoster {
     pub fn purge(&mut self, target: Entity) {
         self.queue.retain(|e| *e != target);
     }
+
+    /// Move all items from the deferred queue into the immediate queue
+    pub fn process_deferals(&mut self) {
+        for e in self.deferred_queue.drain(..) {
+            self.queue.push_back(e);
+        }
+    }
 }
 
 /// Notify the scope manager that a disposal took place. This will prompt the

src/operation/operate_task.rs

@@ -224,8 +224,10 @@ impl<Response: 'static + Send + Sync> Operation for OperateTask<Response> {
             &mut Context::from_waker(&waker_ref(&waker))
         ) {
             Poll::Ready(result) => {
-                // Task has finished
-                let r = world.entity_mut(target).give_input(session, result, roster);
+                // Task has finished. We will defer its input until after the
+                // ChannelQueue has been processed so that any streams from this
+                // task will be delivered before the final output.
+                let r = world.entity_mut(target).defer_input(session, result, roster);
                 world.get_mut::<OperateTask<Response>>(source).or_broken()?.finished_normally = true;
                 cleanup_task::<Response>(session, source, node, unblock, being_cleaned, world, roster);
                 r?;

src/service.rs

@@ -365,7 +365,7 @@ where
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::{BlockingService, InBlockingService, AsyncService, InAsyncService};
+    use crate::{BlockingService, BlockingServiceInput, AsyncService, AsyncServiceInput};
     use bevy::{
         prelude::*,
         ecs::world::EntityMut,
@@ -485,7 +485,7 @@ mod tests {
     }
 
     fn sys_async_service(
-        In(AsyncService{ request, .. }): InAsyncService<String>,
+        In(AsyncService{ request, .. }): AsyncServiceInput<String>,
         people: Query<&TestPeople>,
     ) -> impl Future<Output=u64> {
         let mut matching_people = Vec::new();
@@ -507,7 +507,7 @@ mod tests {
     }
 
     fn sys_blocking_service(
-        In(BlockingService{ request, provider, .. }): InBlockingService<String>,
+        In(BlockingService{ request, provider, .. }): BlockingServiceInput<String>,
         people: Query<&TestPeople>,
         multipliers: Query<&Multiplier>,
     ) -> u64 {

src/service/async_srv.rs

@@ -16,7 +16,7 @@
 */
 
 use crate::{
-    AsyncService, InAsyncService, IntoService, ServiceTrait, ServiceBundle, ServiceRequest,
+    AsyncService, AsyncServiceInput, IntoService, ServiceTrait, ServiceBundle, ServiceRequest,
     InnerChannel, ChannelQueue, OperationRoster, Blocker,
     StreamPack, ServiceBuilder, ChooseAsyncServiceDelivery, OperationRequest,
     OperationError, OrBroken, ManageInput, Input, OperateTask,
@@ -358,6 +358,6 @@ where
     }
 }
 
-fn peel_async<Request>(In(AsyncService { request, .. }): InAsyncService<Request>) -> Request {
+fn peel_async<Request>(In(AsyncService { request, .. }): AsyncServiceInput<Request>) -> Request {
     request
 }