Fixing implementation details of maps and tasks

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

Author: mxgrey

src/channel.rs

@@ -103,7 +103,7 @@ impl InnerChannel {
     }
 }
 
-type ChannelItem = Box<dyn FnOnce(&mut World, &mut OperationRoster) + Send>;
+pub(crate) type ChannelItem = Box<dyn FnOnce(&mut World, &mut OperationRoster) + Send>;
 
 #[derive(Resource)]
 pub(crate) struct ChannelQueue {

src/disposal.rs

@@ -47,6 +47,10 @@ impl Disposal {
         ServiceUnavailable { service, for_node }.into()
     }
 
+    pub fn task_despawned(task: Entity, node: Entity) -> Disposal {
+        TaskDespawned { task, node }.into()
+    }
+
     pub fn branching(
         branched_at_node: Entity,
         disposed_for_target: Entity,
@@ -85,6 +89,11 @@ pub enum DisposalCause {
     /// is the unavailable service.
     ServiceUnavailable(ServiceUnavailable),
 
+    /// An entity that was managing the execution of a task was despawned,
+    /// causing the task to be cancelled and making it impossible to deliver a
+    /// response.
+    TaskDespawned(TaskDespawned),
+
     /// An output was disposed because a mutex was poisoned.
     PoisonedMutex(PoisonedMutexDisposal),
 
@@ -187,6 +196,21 @@ impl From<ServiceUnavailable> for DisposalCause {
     }
 }
 
+/// A variant of [`DisposalCause`]
+#[derive(Debug)]
+pub struct TaskDespawned {
+    /// The entity that was managing the task
+    pub task: Entity,
+    /// The node that the task was spawned by
+    pub node: Entity,
+}
+
+impl From<TaskDespawned> for DisposalCause {
+    fn from(value: TaskDespawned) -> Self {
+        Self::TaskDespawned(value)
+    }
+}
+
 /// A variant of [`DisposalCause`]
 #[derive(Debug)]
 pub struct PoisonedMutexDisposal {

src/flush.rs

@@ -26,7 +26,7 @@ use smallvec::SmallVec;
 
 use crate::{
     ChannelQueue, WakeQueue, OperationRoster, ServiceHook, InputReady,
-    Cancel, DroppedPromiseQueue, UnusedTarget, ServiceLifecycle, ServiceLifecycleChannel,
+    Cancel, UnusedTarget, ServiceLifecycle, ServiceLifecycleChannel,
     OperationRequest,
     execute_operation, dispose_for_despawned_service,
 };

src/map.rs

@@ -36,16 +36,16 @@ pub trait AsMap<M> {
 
 /// A trait that all different ways of defining a Blocking Map must funnel into.
 pub(crate) trait CallBlockingMap<Request, Response> {
-    fn call(self, input: BlockingMap<Request>) -> Response;
+    fn call(&mut self, input: BlockingMap<Request>) -> Response;
 }
 
 impl<F, Request, Response> CallBlockingMap<Request, Response> for MapDef<F>
 where
-    F: FnOnce(BlockingMap<Request>) -> Response + 'static + Send + Sync,
+    F: FnMut(BlockingMap<Request>) -> Response + 'static + Send + Sync,
     Request: 'static + Send + Sync,
     Response: 'static + Send + Sync,
 {
-    fn call(self, request: BlockingMap<Request>) -> Response {
+    fn call(&mut self, request: BlockingMap<Request>) -> Response {
         (self.0)(request)
     }
 }
@@ -111,25 +111,25 @@ pub struct BlockingMapAdapter<F>(F);
 
 impl<F, Request, Response> CallBlockingMap<Request, Response> for BlockingMapAdapter<F>
 where
-    F: FnOnce(Request) -> Response,
+    F: FnMut(Request) -> Response,
 {
-    fn call(self, BlockingMap{ request }: BlockingMap<Request>) -> Response {
+    fn call(&mut self, BlockingMap{ request }: BlockingMap<Request>) -> Response {
         (self.0)(request)
     }
 }
 
 pub(crate) trait CallAsyncMap<Request, Task, Streams: StreamPack> {
-    fn call(self, input: AsyncMap<Request, Streams>) -> Task;
+    fn call(&mut self, input: AsyncMap<Request, Streams>) -> Task;
 }
 
 impl<F, Request, Task, Streams> CallAsyncMap<Request, Task, Streams> for MapDef<F>
 where
-    F: FnOnce(AsyncMap<Request, Streams>) -> Task + 'static + Send + Sync,
+    F: FnMut(AsyncMap<Request, Streams>) -> Task + 'static + Send + Sync,
     Request: 'static + Send + Sync,
     Task: 'static + Send + Sync,
     Streams: StreamPack,
 {
-    fn call(self, input: AsyncMap<Request, Streams>) -> Task {
+    fn call(&mut self, input: AsyncMap<Request, Streams>) -> Task {
         (self.0)(input)
     }
 }
@@ -198,10 +198,10 @@ pub struct AsyncMapAdapter<F>(F);
 
 impl<F, Request, Task> CallAsyncMap<Request, Task, ()> for AsyncMapAdapter<F>
 where
-    F: FnOnce(Request) -> Task + 'static + Send + Sync,
+    F: FnMut(Request) -> Task + 'static + Send + Sync,
     Task: Future + 'static + Send + Sync,
 {
-    fn call(self, AsyncMap{ request, .. }: AsyncMap<Request, ()>) -> Task {
+    fn call(&mut self, AsyncMap{ request, .. }: AsyncMap<Request, ()>) -> Task {
         (self.0)(request)
     }
 }

src/operation.rs

@@ -61,9 +61,6 @@ pub(crate) use operate_service::*;
 mod operate_task;
 pub(crate) use operate_task::*;
 
-mod receive;
-pub(crate) use receive::*;
-
 mod scope;
 pub use scope::*;
 
@@ -265,9 +262,6 @@ pub(crate) struct Blocker {
     pub(crate) serve_next: fn(Blocker, &mut World, &mut OperationRoster),
 }
 
-#[derive(Component)]
-pub(crate) struct BlockerStorage(pub(crate) Option<Blocker>);
-
 pub enum OperationError {
     Broken(Option<Backtrace>),
     NotReady,

src/operation/operate_map.rs

@@ -50,7 +50,7 @@ where
     pub(crate) fn new(target: Entity, f: F) -> Self {
         Self {
             storage: BlockingMapStorage {
-                f,
+                f: Some(f),
                 _ignore: Default::default(),
             },
             target: SingleTargetStorage::new(target),
@@ -60,7 +60,7 @@ where
 
 #[derive(Component)]
 struct BlockingMapStorage<F, Request, Response> {
-    f: F,
+    f: Option<F>,
     _ignore: std::marker::PhantomData<(Request, Response)>,
 }
 
@@ -87,11 +87,13 @@ where
         let mut source_mut = world.get_entity_mut(source).or_broken()?;
         let target = source_mut.get::<SingleTargetStorage>().or_broken()?.0;
         let Input { session, data: request } = source_mut.take_input::<Request>()?;
-        let map = source_mut.take::<BlockingMapStorage<F, Request, Response>>().or_broken()?.f;
-        let mut target_mut = world.get_entity_mut(target).or_broken()?;
+        let mut map = source_mut.take::<BlockingMapStorage<F, Request, Response>>().or_broken()?;
+        let mut f = map.f.take().or_broken()?;
 
-        let response = map.call(BlockingMap { request });
-        target_mut.give_input(session, response, roster)?;
+        let response = f.call(BlockingMap { request });
+        map.f = Some(f);
+
+        world.get_entity_mut(target).or_broken()?.give_input(session, response, roster)?;
         Ok(())
     }
 
@@ -130,7 +132,7 @@ where
     pub(crate) fn new(target: Entity, f: F) -> Self {
         Self {
             storage: AsyncMapStorage {
-                f,
+                f: Some(f),
                 _ignore: Default::default(),
             },
             target: SingleTargetStorage::new(target),
@@ -140,7 +142,7 @@ where
 
 #[derive(Component)]
 struct AsyncMapStorage<F, Request, Task, Streams> {
-    f: F,
+    f: Option<F>,
     _ignore: std::marker::PhantomData<(Request, Task, Streams)>,
 }
 
@@ -171,15 +173,19 @@ where
         let mut source_mut = world.get_entity_mut(source).or_broken()?;
         let Input { session, data: request } = source_mut.take_input::<Request>()?;
         let target = source_mut.get::<SingleTargetStorage>().or_broken()?.0;
-        let map = source_mut.take::<AsyncMapStorage<F, Request, Task, Streams>>().or_broken()?.f;
+        let mut f = source_mut.get_mut::<AsyncMapStorage<F, Request, Task, Streams>>().or_broken()?
+            .f.take().or_broken()?;
 
-        let channel = InnerChannel::new(source, session, sender);
+        let channel = InnerChannel::new(source, session, sender.clone());
         let channel = channel.into_specific(&world)?;
 
-        let task = AsyncComputeTaskPool::get().spawn(map.call(AsyncMap { request, channel }));
+        let task = AsyncComputeTaskPool::get().spawn(f.call(AsyncMap { request, channel }));
+        world.get_entity_mut(source).or_broken()?
+            .get_mut::<AsyncMapStorage<F, Request, Task, Streams>>().or_broken()?
+            .f = Some(f);
 
         let task_source = world.spawn(()).id();
-        OperateTask::new(session, source, target, task, None)
+        OperateTask::new(task_source, session, source, target, task, None, sender)
             .setup(OperationSetup { source: task_source, world });
         roster.queue(task_source);
         Ok(())