Fixed task cleanup race conditions

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

Author: mxgrey

src/flush.rs

@@ -36,7 +36,7 @@ use crate::{
     UnusedTarget, ServiceLifecycle, ServiceLifecycleChannel, MiscellaneousFailure,
     OperationRequest, ImpulseLifecycleChannel, AddImpulse, Finished, OperationCleanup,
     UnhandledErrors, UnusedTargetDrop, ValidateScopeReachability, OperationError,
-    execute_operation, dispose_for_despawned_service,
+    execute_operation, awaken_task, dispose_for_despawned_service,
 };
 
 #[derive(Resource, Default)]
@@ -93,6 +93,11 @@ pub fn flush_impulses(
             garbage_cleanup(world, &mut roster);
         }
 
+        while let Some(source) = roster.awake.pop_front() {
+            awaken_task(OperationRequest { source, world, roster: &mut roster });
+            garbage_cleanup(world, &mut roster);
+        }
+
         collect_from_channels(new_service_query, world, &mut roster);
     }
 }
@@ -152,7 +157,7 @@ fn collect_from_channels(
         .receiver
         .try_iter()
     {
-        roster.queue(wakeable);
+        roster.awake(wakeable);
     }
 
     let mut unused_targets_state: SystemState<Query<(Entity, &Detached), With<UnusedTarget>>> =

src/operation.rs

@@ -182,6 +182,9 @@ pub(crate) struct UnusedTarget;
 pub struct OperationRoster {
     /// Operation sources that should be triggered
     pub(crate) queue: VecDeque<Entity>,
+    /// 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 cancelled
     pub(crate) cancel: VecDeque<Cancel>,
     /// Async services that should pull their next item
@@ -201,6 +204,10 @@ impl OperationRoster {
         self.queue.push_back(source);
     }
 
+    pub fn awake(&mut self, source: Entity) {
+        self.awake.push_back(source);
+    }
+
     pub fn cancel(&mut self, source: Cancel) {
         self.cancel.push_back(source);
     }
@@ -218,8 +225,11 @@ impl OperationRoster {
     }
 
     pub fn is_empty(&self) -> bool {
-        self.queue.is_empty() && self.cancel.is_empty()
-        && self.unblock.is_empty() && self.disposed.is_empty()
+        self.queue.is_empty()
+        && self.awake.is_empty()
+        && self.cancel.is_empty()
+        && self.unblock.is_empty()
+        && self.disposed.is_empty()
         && self.cleanup_finished.is_empty()
     }
 
@@ -626,6 +636,15 @@ pub fn execute_operation(request: OperationRequest) {
     operator(request);
 }
 
+pub fn awaken_task(request: OperationRequest) {
+    let Some(operator) = request.world.get::<OperationExecuteStorage>(request.source) else {
+        // If the task is not available, we just accept that it has despawned.
+        return;
+    };
+    let operator = operator.0;
+    operator(request);
+}
+
 fn perform_operation<Op: Operation>(
     OperationRequest { source, world, roster }: OperationRequest
 ) {

src/operation/operate_task.rs

@@ -123,10 +123,8 @@ impl<Response: 'static + Send + Sync> OperateTask<Response> {
 
 impl<Response: 'static + Send + Sync> Drop for OperateTask<Response> {
     fn drop(&mut self) {
-        println!(" ==== DROPPING TASK: {:?}:\n{:?}", self.source, backtrace::Backtrace::new());
         if self.finished_normally {
             // The task finished normally so no special action needs to be taken
-            dbg!(self.source);
             return;
         }
 
@@ -187,8 +185,20 @@ impl<Response: 'static + Send + Sync> Operation for OperateTask<Response> {
     fn execute(
         OperationRequest { source, world, roster }: OperationRequest
     ) -> OperationResult {
-        let mut source_mut = world.get_entity_mut(source).or_broken()?;
-        let mut operation = source_mut.get_mut::<OperateTask<Response>>().or_broken()?;
+        // It's possible for a task to get into the roster after it has despawned
+        // so we'll just exit early when that happens. However this should not
+        // actually be possible because the OperationExecuteStorage must still be
+        // accessible in order to be inside this function.
+        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()?;
+        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
+            // simply return immediately.
+            return Ok(());
+        }
         let mut task = operation.task.take().or_broken()?;
         let target = operation.target;
         let session = operation.session;
@@ -215,8 +225,9 @@ impl<Response: 'static + Send + Sync> Operation for OperateTask<Response> {
         ) {
             Poll::Ready(result) => {
                 // Task has finished
-                world.get_mut::<OperateTask<Response>>(source).or_broken()?.finished_normally = true;
+                println!(" == READY: {source:?}");
                 let r = world.entity_mut(target).give_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?;
             }
@@ -264,7 +275,9 @@ impl<Response: 'static + Send + Sync> Operation for OperateTask<Response> {
         dbg!(source, node);
         AsyncComputeTaskPool::get().spawn(async move {
             if let Some(task) = task {
+                dbg!(source, node);
                 task.cancel().await;
+                dbg!(source, node);
             }
             if let Err(err) = sender.send(Box::new(move |world: &mut World, roster: &mut OperationRoster| {
                 cleanup_task::<Response>(session, source, node, unblock, true, world, roster);
@@ -316,6 +329,7 @@ fn cleanup_task<Response>(
             if being_cleaned && cleanup_ready {
                 // We are notifying about the cleanup on behalf of the node that
                 // created this task, so we set initialize as source: node
+                dbg!(source, node);
                 let mut cleanup = OperationCleanup {
                     source: node, session, world, roster
                 };
@@ -331,6 +345,8 @@ fn cleanup_task<Response>(
     if world.get_entity(source).is_some() {
         world.despawn(source);
     }
+
+    roster.purge(source);
 }
 
 #[derive(Component, Clone, Copy)]

src/testing.rs

@@ -24,7 +24,7 @@ pub use bevy::{
     ecs::system::CommandQueue,
 };
 
-pub use std::time::Duration;
+pub use std::time::{Duration, Instant};
 
 use crate::{
     Promise, Service, InAsyncService, InBlockingService, UnhandledErrors,
@@ -262,8 +262,15 @@ pub struct WaitRequest<Value> {
 #[cfg(test)]
 pub async fn wait<Value>(request: WaitRequest<Value>) -> Value {
     use async_std::future;
-    let never = future::pending::<()>();
-    let _ = future::timeout(request.duration, never);
+    let start = Instant::now();
+    let mut elapsed = start.elapsed();
+    while elapsed < request.duration {
+        let never = future::pending::<()>();
+        let timeout = request.duration - elapsed;
+        dbg!(request.duration, elapsed, timeout);
+        let _ = future::timeout(timeout, never).await;
+        elapsed = start.elapsed();
+    }
     request.value
 }
 

src/workflow.rs

@@ -343,14 +343,14 @@ mod tests {
                 |chain: Chain<f64>| chain
                     .map_block(|t| WaitRequest { duration: Duration::from_secs_f64(10.0*t), value: 10.0*t })
                     .map(|r: AsyncMap<WaitRequest<f64>>| {
-                        dbg!(r.source);
+                        dbg!(r.source, r.request.value);
                         wait(r.request)
                     })
                     .connect(scope.terminate),
                 |chain: Chain<f64>| chain
                     .map_block(|t| WaitRequest { duration: Duration::from_secs_f64(t/100.0), value: t/100.0 })
                     .map(|r: AsyncMap<WaitRequest<f64>>| {
-                        dbg!(r.source);
+                        dbg!(r.source, r.request.value);
                         wait(r.request)
                     })
                     .connect(scope.terminate),