Switch actor-loop supervision channel to a raw tokio mpsc (#3834)

Summary:
Pull Request resolved: #3834

The supervision channel between an actor and its child does not need to be a monarch channel. A raw tokio mpsc channel is good enough.

The main benefit of using the raw tokio mpsc is that we do not need to worry about the additional complexity brought from the monarch channel, especially the sequence number generation.

In addition, it make the semantics clear. The raw tokio mpsc means the sender is within the same process.

Reviewed By: mariusae

Differential Revision: D104848635

fbshipit-source-id: b5f4777a6927ac97caa11783d7edd7a0bad8ae77

Author: pzhan9

hyperactor/src/proc.rs

@@ -523,7 +523,7 @@ pub struct ActorInstance<A: Actor> {
     /// Handle to the actor (used for lifecycle control and port access).
     pub handle: ActorHandle<A>,
     /// Supervision events delivered to this actor.
-    pub supervision: PortReceiver<ActorSupervisionEvent>,
+    pub supervision: mpsc::UnboundedReceiver<ActorSupervisionEvent>,
     /// Control signals for the actor.
     pub signal: PortReceiver<Signal>,
     /// Primary work queue for handler dispatch.
@@ -1814,7 +1814,10 @@ impl<A: Actor> Drop for InstanceState<A> {
 pub struct InstanceReceivers<A: Actor> {
     /// Signal and supervision receivers for the actor loop. `None`
     /// for detached/client instances that don't run an actor loop.
-    actor_loop: Option<(PortReceiver<Signal>, PortReceiver<ActorSupervisionEvent>)>,
+    actor_loop: Option<(
+        PortReceiver<Signal>,
+        mpsc::UnboundedReceiver<ActorSupervisionEvent>,
+    )>,
     /// Work queue for dispatching messages to actor handlers.
     work: mpsc::UnboundedReceiver<WorkCell<A>>,
     /// Introspect message receiver for the dedicated introspect task.
@@ -1854,10 +1857,10 @@ impl<A: Actor> Instance<A> {
             None
         } else {
             let (signal_port, signal_receiver) = mailbox.open_port::<Signal>();
-            let (supervision_port, supervision_receiver) =
-                mailbox.open_port::<ActorSupervisionEvent>();
+            let (supervision_tx, supervision_receiver) =
+                mpsc::unbounded_channel::<ActorSupervisionEvent>();
             Some((
-                (signal_port, supervision_port),
+                (signal_port, supervision_tx),
                 (signal_receiver, supervision_receiver),
             ))
         };
@@ -2308,7 +2311,10 @@ impl<A: Actor> Instance<A> {
     async fn serve(
         mut self,
         mut actor: A,
-        actor_loop_receivers: (PortReceiver<Signal>, PortReceiver<ActorSupervisionEvent>),
+        actor_loop_receivers: (
+            PortReceiver<Signal>,
+            mpsc::UnboundedReceiver<ActorSupervisionEvent>,
+        ),
         mut work_rx: mpsc::UnboundedReceiver<WorkCell<A>>,
     ) {
         let result = self
@@ -2366,7 +2372,7 @@ impl<A: Actor> Instance<A> {
         if let Some(parent) = self.inner.cell.maybe_unlink_parent() {
             if let Some(event) = event {
                 // Parent exists, failure should be propagated to the parent.
-                parent.send_supervision_event_or_crash(&self, event);
+                parent.send_supervision_event_or_crash(event);
             }
             // TODO: we should get rid of this signal, and use *only* supervision events for
             // the purpose of conveying lifecycle changes
@@ -2400,7 +2406,10 @@ impl<A: Actor> Instance<A> {
     async fn run_actor_tree(
         &mut self,
         actor: &mut A,
-        mut actor_loop_receivers: (PortReceiver<Signal>, PortReceiver<ActorSupervisionEvent>),
+        mut actor_loop_receivers: (
+            PortReceiver<Signal>,
+            mpsc::UnboundedReceiver<ActorSupervisionEvent>,
+        ),
         work_rx: &mut mpsc::UnboundedReceiver<WorkCell<A>>,
     ) -> Result<String, ActorError> {
         // It is okay to catch all panics here, because we are in a tokio task,
@@ -2519,7 +2528,10 @@ impl<A: Actor> Instance<A> {
     async fn run(
         &mut self,
         actor: &mut A,
-        actor_loop_receivers: &mut (PortReceiver<Signal>, PortReceiver<ActorSupervisionEvent>),
+        actor_loop_receivers: &mut (
+            PortReceiver<Signal>,
+            mpsc::UnboundedReceiver<ActorSupervisionEvent>,
+        ),
         work_rx: &mut mpsc::UnboundedReceiver<WorkCell<A>>,
     ) -> Result<String, ActorError> {
         let (signal_receiver, supervision_event_receiver) = actor_loop_receivers;
@@ -2572,7 +2584,7 @@ impl<A: Actor> Instance<A> {
                     let _ = ACTOR_MESSAGE_HANDLER_DURATION.start(metric_pairs);
                     let work = work.expect("inconsistent work queue state");
                     if let Err(err) = work.handle(actor, self).await {
-                        for supervision_event in supervision_event_receiver.drain() {
+                        while let Ok(supervision_event) = supervision_event_receiver.try_recv() {
                             self.handle_supervision_event(actor, supervision_event).await?;
                         }
                         let kind = ActorErrorKind::processing(err);
@@ -2582,7 +2594,7 @@ impl<A: Actor> Instance<A> {
                         });
                     }
                 }
-                Ok(supervision_event) = supervision_event_receiver.recv() => {
+                Some(supervision_event) = supervision_event_receiver.recv() => {
                     self.handle_supervision_event(actor, supervision_event).await?;
                 }
             }
@@ -2971,7 +2983,10 @@ struct InstanceCellState {
     proc: Proc,
 
     /// Control port handles to the actor loop, if one is running.
-    actor_loop: Option<(PortHandle<Signal>, PortHandle<ActorSupervisionEvent>)>,
+    actor_loop: Option<(
+        PortHandle<Signal>,
+        mpsc::UnboundedSender<ActorSupervisionEvent>,
+    )>,
 
     /// An observer that stores the current status of the actor.
     status: watch::Receiver<ActorStatus>,
@@ -3120,7 +3135,10 @@ impl InstanceCell {
         actor_id: ActorAddr,
         actor_type: ActorType,
         proc: Proc,
-        actor_loop: Option<(PortHandle<Signal>, PortHandle<ActorSupervisionEvent>)>,
+        actor_loop: Option<(
+            PortHandle<Signal>,
+            mpsc::UnboundedSender<ActorSupervisionEvent>,
+        )>,
         status: watch::Receiver<ActorStatus>,
         parent: Option<InstanceCell>,
         ports: Arc<dyn Any + Send + Sync>,
@@ -3229,14 +3247,10 @@ impl InstanceCell {
     /// Note that "let it crash" is the default behavior when a supervision event
     /// cannot be delivered upstream. It is the upstream's responsibility to
     /// detect and handle crashes.
-    pub fn send_supervision_event_or_crash(
-        &self,
-        child_cx: &impl context::Actor, // context of the child who sends the event.
-        event: ActorSupervisionEvent,
-    ) {
+    pub fn send_supervision_event_or_crash(&self, event: ActorSupervisionEvent) {
         match &self.inner.actor_loop {
-            Some((_, supervision_port)) => {
-                if let Err(err) = supervision_port.send(child_cx, event.clone()) {
+            Some((_, supervision_tx)) => {
+                if let Err(err) = supervision_tx.send(event.clone()) {
                     if !event.is_error() {
                         // Normal lifecycle events (e.g. clean stop) that fail to
                         // send are silently dropped. This happens when a child

hyperactor_mesh/src/global_context.rs

@@ -177,7 +177,7 @@ pub struct GlobalClientActor {
     /// The root client is a monitor, so it should process these
     /// events without crashing on routine routing/delivery failures
     /// it observes.
-    supervision_rx: PortReceiver<ActorSupervisionEvent>,
+    supervision_rx: mpsc::UnboundedReceiver<ActorSupervisionEvent>,
     /// Primary work queue for handler dispatch.
     ///
     /// Any bound handler message (e.g. `MeshFailure`,
@@ -195,7 +195,7 @@ impl GlobalClientActor {
                     work = self.work_rx.recv() => {
                         let work = work.expect("inconsistent work queue state");
                         if let Err(err) = work.handle(&mut self, instance).await {
-                            for supervision_event in self.supervision_rx.drain() {
+                            while let Ok(supervision_event) = self.supervision_rx.try_recv() {
                                 instance.handle_supervision_event(&mut self, supervision_event).await
                                     .expect("GlobalClientActor::handle_supervision_event is infallible");
                             }
@@ -209,7 +209,7 @@ impl GlobalClientActor {
                     _ = self.signal_rx.recv() => {
                         // TODO: do we need any signal handling for the root client?
                     }
-                    Ok(supervision_event) = self.supervision_rx.recv() => {
+                    Some(supervision_event) = self.supervision_rx.recv() => {
                         instance.handle_supervision_event(&mut self, supervision_event).await
                             .expect("GlobalClientActor::handle_supervision_event is infallible");
                     }

hyperactor_mesh/src/testing.rs

@@ -182,7 +182,7 @@ fn process_name(pid: u32) -> Option<String> {
 #[derive(Debug)]
 pub struct TestRootClient {
     signal_rx: PortReceiver<Signal>,
-    supervision_rx: PortReceiver<ActorSupervisionEvent>,
+    supervision_rx: mpsc::UnboundedReceiver<ActorSupervisionEvent>,
     work_rx: mpsc::UnboundedReceiver<WorkCell<Self>>,
 }
 
@@ -206,7 +206,7 @@ impl TestRootClient {
                     work = self.work_rx.recv() => {
                         let work = work.expect("inconsistent work queue state");
                         if let Err(err) = work.handle(&mut self, instance).await {
-                            for supervision_event in self.supervision_rx.drain() {
+                            while let Ok(supervision_event) = self.supervision_rx.try_recv() {
                                 if let Err(err) = instance.handle_supervision_event(&mut self, supervision_event).await {
                                     break 'messages err;
                                 }
@@ -221,7 +221,7 @@ impl TestRootClient {
                     _ = self.signal_rx.recv() => {
                         // TODO: do we need any signal handling for the root client?
                     }
-                    Ok(supervision_event) = self.supervision_rx.recv() => {
+                    Some(supervision_event) = self.supervision_rx.recv() => {
                         if let Err(err) = instance.handle_supervision_event(&mut self, supervision_event).await {
                             break err;
                         }

monarch_hyperactor/src/actor.rs

@@ -762,7 +762,7 @@ impl PythonActor {
                             // exiting the loop.
                             // Else, continue handling messages.
                             if let Err(err) = instance.handle_supervision_event(&mut actor, supervision_event).await {
-                                for supervision_event in supervision_rx.drain() {
+                                while let Ok(supervision_event) = supervision_rx.try_recv() {
                                     if let Err(err) = instance.handle_supervision_event(&mut actor, supervision_event).await {
                                         break 'messages Some(err);
                                     }
@@ -787,7 +787,7 @@ impl PythonActor {
                             Err(err) => break Some(err),
                         }
                     }
-                    Ok(supervision_event) = supervision_rx.recv() => {
+                    Some(supervision_event) = supervision_rx.recv() => {
                         if let Err(err) = instance.handle_supervision_event(&mut actor, supervision_event).await {
                             break Some(err);
                         }

monarch_hyperactor/src/actor_mesh.rs

@@ -864,7 +864,7 @@ mod tests {
     #[derive(Debug)]
     struct TestClient {
         signal_rx: PortReceiver<Signal>,
-        supervision_rx: PortReceiver<ActorSupervisionEvent>,
+        supervision_rx: mpsc::UnboundedReceiver<ActorSupervisionEvent>,
         work_rx: mpsc::UnboundedReceiver<WorkCell<Self>>,
     }
 
@@ -895,7 +895,7 @@ mod tests {
                             }
                         }
                         _ = self.signal_rx.recv() => {}
-                        Ok(event) = self.supervision_rx.recv() => {
+                        Some(event) = self.supervision_rx.recv() => {
                             let _ = instance
                                 .handle_supervision_event(&mut self, event)
                                 .await;