Fix unreachability detection for streams out of async nodes

mxgrey · mxgrey · commit ac7467bd80a5 · 2024-07-17T13:03:02.000+08:00
Signed-off-by: Michael X. Grey &lt;grey@openrobotics.org&gt;
diff --git a/src/callback.rs b/src/callback.rs
@@ -129,7 +129,7 @@ impl<'a> CallbackRequest<'a> {
         Ok(())
     }
 
-    fn give_task<Task: Future + 'static + Send>(
+    fn give_task<Task: Future + 'static + Send, Streams: StreamPack>(
         &mut self,
         session: Entity,
         task: Task,
@@ -140,7 +140,7 @@ impl<'a> CallbackRequest<'a> {
         let sender = self.world.get_resource_or_insert_with(|| ChannelQueue::new()).sender.clone();
         let task = AsyncComputeTaskPool::get().spawn(task);
         let task_id = self.world.spawn(()).id();
-        OperateTask::new(task_id, session, self.source, self.target, task, None, sender)
+        OperateTask::<_, Streams>::new(task_id, session, self.source, self.target, task, None, sender)
             .add(self.world, self.roster);
         Ok(())
     }
@@ -244,7 +244,7 @@ where
         }, &mut input.world);
         self.system.apply_deferred(&mut input.world);
 
-        input.give_task(session, task)
+        input.give_task::<_, Streams>(session, task)
     }
 }
 
@@ -361,7 +361,7 @@ where
             let task = (self)(AsyncCallback {
                 request, streams, channel, source: input.source, session,
             });
-            input.give_task(session, task)
+            input.give_task::<_, Streams>(session, task)
         };
         Callback::new(MapCallback { callback })
     }
diff --git a/src/impulse/map.rs b/src/impulse/map.rs
@@ -182,8 +182,9 @@ where
         }));
 
         let task_source = world.spawn(()).id();
-        OperateTask::new(task_source, session, source, target, task, None, sender)
-            .add(world, roster);
+        OperateTask::<_, Streams>::new(
+            task_source, session, source, target, task, None, sender,
+        ).add(world, roster);
         Ok(())
     }
 }
diff --git a/src/operation/operate_map.rs b/src/operation/operate_map.rs
@@ -199,7 +199,7 @@ where
             .f = Some(f);
 
         let task_source = world.spawn(()).id();
-        OperateTask::new(
+        OperateTask::<_, Streams>::new(
             task_source, session, source, target, task, None, sender,
         ).add(world, roster);
         Ok(())
diff --git a/src/operation/operate_task.rs b/src/operation/operate_task.rs
@@ -39,7 +39,7 @@ use crate::{
     OperationRoster, Blocker, ManageInput, ChannelQueue, UnhandledErrors,
     OperationSetup, OperationRequest, OperationResult, Operation, AddOperation,
     OrBroken, OperationCleanup, ChannelItem, OperationError, Broken, ScopeStorage,
-    OperationReachability, ReachabilityResult, emit_disposal, Disposal,
+    OperationReachability, ReachabilityResult, emit_disposal, Disposal, StreamPack,
 };
 
 struct JobWaker {
@@ -70,7 +70,7 @@ impl WakeQueue {
 }
 
 #[derive(Component)]
-pub(crate) struct OperateTask<Response: 'static + Send + Sync> {
+pub(crate) struct OperateTask<Response: 'static + Send + Sync, Streams: StreamPack> {
     source: Entity,
     session: Entity,
     node: Entity,
@@ -81,9 +81,10 @@ pub(crate) struct OperateTask<Response: 'static + Send + Sync> {
     disposal: Option<Disposal>,
     being_cleaned: bool,
     finished_normally: bool,
+    _ignore: std::marker::PhantomData<Streams>,
 }
 
-impl<Response: 'static + Send + Sync> OperateTask<Response> {
+impl<Response: 'static + Send + Sync, Streams: StreamPack> OperateTask<Response, Streams> {
     pub(crate) fn new(
         source: Entity,
         session: Entity,
@@ -104,6 +105,7 @@ impl<Response: 'static + Send + Sync> OperateTask<Response> {
             disposal: None,
             being_cleaned: false,
             finished_normally: false,
+            _ignore: Default::default(),
         }
     }
 
@@ -121,7 +123,11 @@ impl<Response: 'static + Send + Sync> OperateTask<Response> {
     }
 }
 
-impl<Response: 'static + Send + Sync> Drop for OperateTask<Response> {
+impl<Response, Streams> Drop for OperateTask<Response, Streams>
+where
+    Response: 'static + Send + Sync,
+    Streams: StreamPack,
+{
     fn drop(&mut self) {
         if self.finished_normally {
             // The task finished normally so no special action needs to be taken
@@ -157,7 +163,11 @@ impl<Response: 'static + Send + Sync> Drop for OperateTask<Response> {
     }
 }
 
-impl<Response: 'static + Send + Sync> Operation for OperateTask<Response> {
+impl<Response, Streams> Operation for OperateTask<Response, Streams>
+where
+    Response: 'static + Send + Sync,
+    Streams: StreamPack,
+{
     fn setup(self, OperationSetup { source, world }: OperationSetup) -> OperationResult {
         let wake_queue = world.get_resource_or_insert_with(|| WakeQueue::new());
         let waker = Arc::new(JobWaker {
@@ -172,7 +182,7 @@ impl<Response: 'static + Send + Sync> Operation for OperateTask<Response> {
             .insert((
                 self,
                 JobWakerStorage(waker),
-                StopTask(stop_task::<Response>),
+                StopTask(stop_task::<Response, Streams>),
             ))
             .set_parent(node);
 
@@ -192,7 +202,7 @@ impl<Response: 'static + Send + Sync> Operation for OperateTask<Response> {
         let mut source_mut = world.get_entity_mut(source).or_not_ready()?;
         // If the task has been stopped / cancelled then OperateTask will have
         // been removed, even if it has not despawned yet.
-        let mut operation = source_mut.get_mut::<OperateTask<Response>>().or_not_ready()?;
+        let mut operation = source_mut.get_mut::<OperateTask<Response, Streams>>().or_not_ready()?;
         if operation.being_cleaned {
             // The operation is being cleaned up, so the task will not be
             // available and there will be nothing for us to do here. We should
@@ -228,13 +238,35 @@ impl<Response: 'static + Send + Sync> Operation for OperateTask<Response> {
                 // 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;
+                world.get_mut::<OperateTask<Response, Streams>>(source).or_broken()?.finished_normally = true;
                 cleanup_task::<Response>(session, source, node, unblock, being_cleaned, world, roster);
+
+                if Streams::has_streams() {
+                    if let Some(scope) = world.get::<ScopeStorage>(node) {
+                        // When an async task with any number of streams >= 1 is
+                        // finished, we should always do a disposal notification
+                        // to force a reachability check. Normally there are
+                        // specific events that prompt us to check reachability,
+                        // but if a reachability test occurred while the async
+                        // node was running and the reachability depends on a
+                        // stream which may or may not have been emitted, then
+                        // the reachability test may have concluded with a false
+                        // positive, and it needs to be rechecked now that the
+                        // async node has finished.
+                        //
+                        // TODO(@mxgrey): Make this more efficient, e.g. only
+                        // trigger this disposal if we detected that a
+                        // reachability test happened while this task was
+                        // running.
+                        roster.disposed(scope.get(), session);
+                    }
+                }
+
                 r?;
             }
             Poll::Pending => {
                 // Task is still running
-                if let Some(mut operation) = world.get_mut::<OperateTask<Response>>(source) {
+                if let Some(mut operation) = world.get_mut::<OperateTask<Response, Streams>>(source) {
                     operation.task = Some(task);
                     operation.blocker = unblock;
                     world.entity_mut(source).insert(JobWakerStorage(waker));
@@ -249,7 +281,7 @@ impl<Response: 'static + Send + Sync> Operation for OperateTask<Response> {
                             || ChannelQueue::default()
                         ).sender.clone();
 
-                        let operation = OperateTask::new(
+                        let operation = OperateTask::<_, Streams>::new(
                             source, session, node, target, task, unblock, sender,
                         );
 
@@ -267,7 +299,7 @@ impl<Response: 'static + Send + Sync> Operation for OperateTask<Response> {
         let session = clean.session;
         let source = clean.source;
         let mut source_mut = clean.world.get_entity_mut(source).or_broken()?;
-        let mut operation = source_mut.get_mut::<OperateTask<Response>>().or_broken()?;
+        let mut operation = source_mut.get_mut::<OperateTask<Response, Streams>>().or_broken()?;
         operation.being_cleaned = true;
         let node = operation.node;
         let task = operation.task.take();
@@ -292,7 +324,7 @@ impl<Response: 'static + Send + Sync> Operation for OperateTask<Response> {
     fn is_reachable(reachability: OperationReachability) -> ReachabilityResult {
         let session = reachability.world
             .get_entity(reachability.source).or_broken()?
-            .get::<OperateTask<Response>>().or_broken()?.session;
+            .get::<OperateTask<Response, Streams>>().or_broken()?.session;
         Ok(session == reachability.session)
     }
 }
@@ -351,13 +383,13 @@ fn cleanup_task<Response>(
 #[derive(Component, Clone, Copy)]
 pub(crate) struct StopTask(pub(crate) fn(OperationRequest, Disposal) -> OperationResult);
 
-fn stop_task<Response: 'static + Send + Sync>(
+fn stop_task<Response: 'static + Send + Sync, Streams: StreamPack>(
     OperationRequest { source, world, .. }: OperationRequest,
     disposal: Disposal,
 ) -> OperationResult {
     let mut operation = world
         .get_entity_mut(source).or_broken()?
-        .take::<OperateTask<Response>>().or_broken()?;
+        .take::<OperateTask<Response, Streams>>().or_broken()?;
 
     operation.disposal = Some(disposal);
     drop(operation);
diff --git a/src/service/async_srv.rs b/src/service/async_srv.rs
@@ -238,8 +238,9 @@ where
 
     let task = AsyncComputeTaskPool::get().spawn(job);
 
-    OperateTask::new(task_id, session, source, target, task, blocker, sender)
-        .add(world, roster);
+    OperateTask::<_, Streams>::new(
+        task_id, session, source, target, task, blocker, sender,
+    ).add(world, roster);
     Ok(())
 }
 
diff --git a/src/stream.rs b/src/stream.rs
@@ -291,6 +291,9 @@ pub trait StreamPack: 'static + Send + Sync {
         world: &mut World,
         roster: &mut OperationRoster,
     ) -> OperationResult;
+
+    /// Are there actually any streams in the pack?
+    fn has_streams() -> bool;
 }
 
 impl<T: Stream> StreamPack for T {
@@ -396,6 +399,10 @@ impl<T: Stream> StreamPack for T {
 
         Ok(())
     }
+
+    fn has_streams() -> bool {
+        true
+    }
 }
 
 impl StreamPack for () {
@@ -470,6 +477,10 @@ impl StreamPack for () {
     ) -> OperationResult {
         Ok(())
     }
+
+    fn has_streams() -> bool {
+        false
+    }
 }
 
 impl<T1: StreamPack> StreamPack for (T1,) {
@@ -545,6 +556,10 @@ impl<T1: StreamPack> StreamPack for (T1,) {
         T1::process_buffer(buffer, source, session, unused, world, roster)?;
         Ok(())
     }
+
+    fn has_streams() -> bool {
+        T1::has_streams()
+    }
 }
 
 impl<T1: StreamPack, T2: StreamPack> StreamPack for (T1, T2) {
@@ -635,6 +650,11 @@ impl<T1: StreamPack, T2: StreamPack> StreamPack for (T1, T2) {
         T2::process_buffer(buffer.1, source, session, unused, world, roster)?;
         Ok(())
     }
+
+    fn has_streams() -> bool {
+        T1::has_streams()
+        || T2::has_streams()
+    }
 }
 
 impl<T1: StreamPack, T2: StreamPack, T3: StreamPack> StreamPack for (T1, T2, T3) {
@@ -733,6 +753,12 @@ impl<T1: StreamPack, T2: StreamPack, T3: StreamPack> StreamPack for (T1, T2, T3)
         T3::process_buffer(buffer.2, source, session, unused, world, roster)?;
         Ok(())
     }
+
+    fn has_streams() -> bool {
+        T1::has_streams()
+        || T2::has_streams()
+        || T3::has_streams()
+    }
 }
 
 /// Used by [`ServiceDiscovery`](crate::ServiceDiscovery) to filter services
diff --git a/src/workflow.rs b/src/workflow.rs
@@ -382,9 +382,10 @@ mod tests {
     }
 
     #[test]
-    fn test_reachability_with_unused_streams() {
+    fn test_stream_reachability() {
         let mut context = TestingContext::minimal_plugins();
 
+        // Test for streams from a blocking node
         let workflow = context.spawn_io_workflow(|scope, builder| {
             let stream_node = builder.create_map(|_: BlockingMap<(), StreamOf<u32>>| {
                 // Do nothing. The purpose of this node is to just return without
@@ -405,5 +406,26 @@ mod tests {
         context.run_with_conditions(&mut promise, Duration::from_secs(2));
         assert!(promise.peek().is_cancelled());
         assert!(context.no_unhandled_errors());
+
+        // Test for streams from an async node
+        let workflow = context.spawn_io_workflow(|scope, builder| {
+            let stream_node = builder.create_map(|_: AsyncMap<(), StreamOf<u32>>| {
+                async { /* Do nothing */}
+            });
+
+            builder.connect(scope.input, stream_node.input);
+            stream_node.streams.chain(builder)
+                .inner()
+                .map_block(|value| 2 * value)
+                .connect(scope.terminate);
+        });
+
+        let mut promise = context.command(|commands| {
+            commands.request((), workflow).take_response()
+        });
+
+        context.run_with_conditions(&mut promise, Duration::from_secs(2));
+        assert!(promise.peek().is_cancelled());
+        assert!(context.no_unhandled_errors());
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -182,8 +182,9 @@ where`
`182`	`182`	`}));`
`183`	`183`
`184`	`184`	`let task_source = world.spawn(()).id();`
`185`		`- OperateTask::new(task_source, session, source, target, task, None, sender)`
`186`		`- .add(world, roster);`
	`185`	`+ OperateTask::<_, Streams>::new(`
	`186`	`+ task_source, session, source, target, task, None, sender,`
	`187`	`+ ).add(world, roster);`
`187`	`188`	`Ok(())`
`188`	`189`	`}`
`189`	`190`	`}`
Original file line number	Diff line number	Diff line change
`@@ -238,8 +238,9 @@ where`
`238`	`238`
`239`	`239`	`let task = AsyncComputeTaskPool::get().spawn(job);`
`240`	`240`
`241`		`- OperateTask::new(task_id, session, source, target, task, blocker, sender)`
`242`		`- .add(world, roster);`
	`241`	`+ OperateTask::<_, Streams>::new(`
	`242`	`+ task_id, session, source, target, task, blocker, sender,`
	`243`	`+ ).add(world, roster);`
`243`	`244`	`Ok(())`
`244`	`245`	`}`
`245`	`246`
Original file line number	Diff line number	Diff line change
`@@ -291,6 +291,9 @@ pub trait StreamPack: 'static + Send + Sync {`
`291`	`291`	`world: &mut World,`
`292`	`292`	`roster: &mut OperationRoster,`
`293`	`293`	`) -> OperationResult;`
	`294`	`+`
	`295`	`+ /// Are there actually any streams in the pack?`
	`296`	`+ fn has_streams() -> bool;`
`294`	`297`	`}`
`295`	`298`
`296`	`299`	`impl<T: Stream> StreamPack for T {`
`@@ -396,6 +399,10 @@ impl<T: Stream> StreamPack for T {`
`396`	`399`
`397`	`400`	`Ok(())`
`398`	`401`	`}`
	`402`	`+`
	`403`	`+ fn has_streams() -> bool {`
	`404`	`+ true`
	`405`	`+ }`
`399`	`406`	`}`
`400`	`407`
`401`	`408`	`impl StreamPack for () {`
`@@ -470,6 +477,10 @@ impl StreamPack for () {`
`470`	`477`	`) -> OperationResult {`
`471`	`478`	`Ok(())`
`472`	`479`	`}`
	`480`	`+`
	`481`	`+ fn has_streams() -> bool {`
	`482`	`+ false`
	`483`	`+ }`
`473`	`484`	`}`
`474`	`485`
`475`	`486`	`impl<T1: StreamPack> StreamPack for (T1,) {`
`@@ -545,6 +556,10 @@ impl<T1: StreamPack> StreamPack for (T1,) {`
`545`	`556`	`T1::process_buffer(buffer, source, session, unused, world, roster)?;`
`546`	`557`	`Ok(())`
`547`	`558`	`}`
	`559`	`+`
	`560`	`+ fn has_streams() -> bool {`
	`561`	`+ T1::has_streams()`
	`562`	`+ }`
`548`	`563`	`}`
`549`	`564`
`550`	`565`	`impl<T1: StreamPack, T2: StreamPack> StreamPack for (T1, T2) {`
`@@ -635,6 +650,11 @@ impl<T1: StreamPack, T2: StreamPack> StreamPack for (T1, T2) {`
`635`	`650`	`T2::process_buffer(buffer.1, source, session, unused, world, roster)?;`
`636`	`651`	`Ok(())`
`637`	`652`	`}`
	`653`	`+`
	`654`	`+ fn has_streams() -> bool {`
	`655`	`+ T1::has_streams()`
	`656`	`+ \|\| T2::has_streams()`
	`657`	`+ }`
`638`	`658`	`}`
`639`	`659`
`640`	`660`	`impl<T1: StreamPack, T2: StreamPack, T3: StreamPack> StreamPack for (T1, T2, T3) {`
`@@ -733,6 +753,12 @@ impl<T1: StreamPack, T2: StreamPack, T3: StreamPack> StreamPack for (T1, T2, T3)`
`733`	`753`	`T3::process_buffer(buffer.2, source, session, unused, world, roster)?;`
`734`	`754`	`Ok(())`
`735`	`755`	`}`
	`756`	`+`
	`757`	`+ fn has_streams() -> bool {`
	`758`	`+ T1::has_streams()`
	`759`	`+ \|\| T2::has_streams()`
	`760`	`+ \|\| T3::has_streams()`
	`761`	`+ }`
`736`	`762`	`}`
`737`	`763`
`738`	`764`	/// Used by [`ServiceDiscovery`](crate::ServiceDiscovery) to filter services