proc_mesh: supervise undeliverable messages for mesh client (#383)

Summary:
Pull Request resolved: #383

undeliverable messages returned to a proc mesh client are now transformed into actor supervision events, which are in turn mapped to `ProcEvent::Crashed` events.

supervision is handled via `supervise_undeliverable_messages`, and direct access to the proc mesh client undeliverable message port receiver is removed.

Reviewed By: vidhyav

Differential Revision: D77555707

fbshipit-source-id: 2ec3643324a3592654f186f4d09afd5f62402d6f

Author: shayne-fletcher

hyperactor/src/mailbox.rs

@@ -124,6 +124,7 @@ mod undeliverable;
 pub use undeliverable::Undeliverable;
 pub use undeliverable::UndeliverableMessageError;
 pub use undeliverable::monitored_return_handle; // TODO: Audit
+pub use undeliverable::supervise_undeliverable_messages;
 /// For [`MailboxAdminMessage`], a message type for mailbox administration.
 pub mod mailbox_admin_message;
 pub use mailbox_admin_message::MailboxAdminMessage;

hyperactor/src/mailbox/undeliverable.rs

@@ -16,13 +16,15 @@ use crate::ActorId;
 use crate::Message;
 use crate::PortId;
 use crate::RemoteMessage;
+use crate::actor::ActorStatus;
 use crate::id;
 use crate::mailbox::DeliveryError;
 use crate::mailbox::MailboxSender;
 use crate::mailbox::MessageEnvelope;
 use crate::mailbox::PortHandle;
 use crate::mailbox::PortReceiver;
 use crate::mailbox::UndeliverableMailboxSender;
+use crate::supervision::ActorSupervisionEvent;
 
 /// An undeliverable `M`-typed message (in practice `M` is
 /// [MessageEnvelope]).
@@ -130,3 +132,33 @@ impl UndeliverableMessageError {
         }
     }
 }
+
+/// Spawns a task that listens for undeliverable messages and posts a
+/// corresponding `ActorSupervisionEvent` to the given supervision
+/// port.
+///
+/// The `mailbox_id` identifies the source mailbox for context in
+/// supervision events.
+pub fn supervise_undeliverable_messages(
+    mailbox_id: ActorId,
+    supervision_port: PortHandle<ActorSupervisionEvent>,
+    mut rx: PortReceiver<Undeliverable<MessageEnvelope>>,
+) {
+    tokio::spawn(async move {
+        while let Ok(Undeliverable(mut envelope)) = rx.recv().await {
+            envelope.try_set_error(DeliveryError::BrokenLink(
+                "message returned to undeliverable port".to_string(),
+            ));
+            if supervision_port
+                .send(ActorSupervisionEvent::new(
+                    mailbox_id.clone(),
+                    ActorStatus::Failed(format!("message not delivered: {}", envelope)),
+                ))
+                .is_err()
+            {
+                UndeliverableMailboxSender
+                    .post(envelope.clone(), /*unused*/ monitored_return_handle())
+            }
+        }
+    });
+}

hyperactor_mesh/src/actor_mesh.rs

@@ -505,9 +505,9 @@ mod tests {
     use hyperactor::ProcId;
     use hyperactor::WorldId;
     use hyperactor::attrs::Attrs;
-    use hyperactor::mailbox::Undeliverable;
 
     use super::*;
+    use crate::proc_mesh::ProcEvent;
 
     // These tests are parametric over allocators.
     #[macro_export]
@@ -783,18 +783,18 @@ mod tests {
 
                 let name = alloc.name().to_string();
                 let mut mesh = ProcMesh::allocate(alloc).await.unwrap();
-                let mut undeliverable_rx = mesh.client_undeliverable_receiver().take()
-                    .expect("client_undeliverable_receiver should be available");
+                let mut events = mesh.events().unwrap();
 
                 // Send a message to a non-existent actor (the proc however exists).
                 let unmonitored_reply_to = mesh.client().open_port::<usize>().0.bind();
                 let bad_actor = ActorRef::<TestActor>::attest(ActorId(ProcId(WorldId(name.clone()), 0), "foo".into(), 0));
                 bad_actor.send(mesh.client(), GetRank(true, unmonitored_reply_to)).unwrap();
 
                 // The message will be returned!
-                let Undeliverable(msg) = undeliverable_rx.recv().await.unwrap();
-                assert_eq!(mesh.client().actor_id(), msg.sender());
-                assert_eq!(&bad_actor.actor_id().port_id(GetRank::port()), msg.dest());
+                assert_matches!(
+                    events.next().await.unwrap(),
+                    ProcEvent::Crashed(0, reason) if reason.contains("failed: message not delivered")
+                );
 
                 // TODO: Stop the proc.
             }
@@ -870,7 +870,6 @@ mod tests {
             let monkey = alloc.chaos_monkey();
             let mut mesh = ProcMesh::allocate(alloc).await.unwrap();
             let mut events = mesh.events().unwrap();
-            let mut undeliverable_msg_rx = mesh.client_undeliverable_receiver().take().unwrap();
 
             let ping_pong_actor_params = PingPongActorParams::new(
                 PortRef::attest_message_port(mesh.client().actor_id()),
@@ -901,8 +900,10 @@ mod tests {
             .unwrap();
 
             // The message will be returned!
-            let Undeliverable(msg) = undeliverable_msg_rx.recv().await.unwrap();
-            assert_eq!(msg.sender(), mesh.client().actor_id());
+            assert_matches!(
+                events.next().await.unwrap(),
+                ProcEvent::Crashed(0, reason) if reason.contains("failed: message not delivered")
+            );
 
             // Get 'pong' to send 'ping' a message. Since 'ping's
             // mailbox is stopped, the send will timeout and fail.
@@ -914,11 +915,9 @@ mod tests {
             .unwrap();
 
             // The message will be returned!
-            let Undeliverable(msg) = undeliverable_msg_rx.recv().await.unwrap();
-            assert_eq!(msg.sender(), pong.actor_id());
-            assert_eq!(
-                msg.dest(),
-                &ping.actor_id().port_id(PingPongMessage::port())
+            assert_matches!(
+                events.next().await.unwrap(),
+                ProcEvent::Crashed(0, reason) if reason.contains("failed: message not delivered")
             );
         }
 
@@ -938,10 +937,6 @@ mod tests {
 
             let stop = alloc.stopper();
             let mut mesh = ProcMesh::allocate(alloc).await.unwrap();
-            let mut undeliverable_rx = mesh
-                .client_undeliverable_receiver()
-                .take()
-                .expect("client_undeliverable_receiver should be available");
             let mut events = mesh.events().unwrap();
 
             let actor_mesh = mesh
@@ -970,15 +965,10 @@ mod tests {
                 .cast(sel!(*), GetRank(false, reply_handle.bind()))
                 .unwrap();
 
-            // The message will be returned.
-            let Undeliverable(msg) = undeliverable_rx.recv().await.unwrap();
-            assert_eq!(
-                msg.sender(),
-                &ActorId(
-                    ProcId(actor_mesh.world_id().clone(), 0),
-                    "comm".to_owned(),
-                    0
-                )
+            // The message will be returned!
+            assert_matches!(
+                events.next().await.unwrap(),
+                ProcEvent::Crashed(0, reason) if reason.contains("failed: message not delivered")
             );
 
             // Stop the mesh.

hyperactor_mesh/src/proc_mesh.rs

@@ -82,7 +82,6 @@ pub struct ProcMesh {
     #[allow(dead_code)] // will be used in subsequent diff
     client_proc: Proc,
     client: Mailbox,
-    client_undeliverable_receiver: Option<PortReceiver<Undeliverable<MessageEnvelope>>>,
     comm_actors: Vec<ActorRef<CommActor>>,
     world_id: WorldId,
 }
@@ -211,16 +210,25 @@ impl ProcMesh {
 
         // TODO: No actor bound to "supervisor" yet.
         let supervisor = client_proc.attach("supervisor")?;
-        let (supervison_port, supervision_events) = supervisor.open_port();
+        let (supervision_port, supervision_events) =
+            supervisor.open_port::<ActorSupervisionEvent>();
 
         // Now, configure the full mesh, so that the local agents are
-        // wired up to our router. Bind an undeliverable message port
-        // in the client and return the port receiver.
-        // No actor bound to this "client" yet
+        // wired up to our router.
+        // TODO: No actor bound to "client" yet.
         let client = client_proc.attach("client")?;
+        // Bind an undeliverable message port in the client.
         let (undeliverable_messages, client_undeliverable_receiver) =
             client.open_port::<Undeliverable<MessageEnvelope>>();
         undeliverable_messages.bind_to(Undeliverable::<MessageEnvelope>::port());
+        // Monitor undeliverable messages from the client and emit
+        // corresponding actor supervision events via the supervision
+        // port.
+        hyperactor::mailbox::supervise_undeliverable_messages(
+            client.actor_id().clone(),
+            supervision_port.clone(),
+            client_undeliverable_receiver,
+        );
 
         // Map of procs -> channel addresses
         let address_book: HashMap<_, _> = running
@@ -235,7 +243,7 @@ impl ProcMesh {
                     &client,
                     rank,
                     router_channel_addr.clone(),
-                    supervison_port.bind(),
+                    supervision_port.bind(),
                     address_book.clone(),
                     config_handle.bind(),
                 )
@@ -304,7 +312,6 @@ impl ProcMesh {
                 .collect(),
             client_proc,
             client,
-            client_undeliverable_receiver: Some(client_undeliverable_receiver),
             comm_actors,
             world_id,
         })
@@ -399,22 +406,6 @@ impl ProcMesh {
         &self.client
     }
 
-    /// Returns a mutable reference to the client mailbox's
-    /// undeliverable message port receiver.
-    ///
-    /// This allows the caller to extract the
-    /// `PortReceiver<Undeliverable<MessageEnvelope>>` by calling
-    /// `.take()` on the returned `Option`, transferring ownership of
-    /// the receiver.
-    ///
-    /// Typically used to access the port bound by
-    /// `ProcMesh::allocate`.
-    pub fn client_undeliverable_receiver(
-        &mut self,
-    ) -> &mut Option<PortReceiver<Undeliverable<MessageEnvelope>>> {
-        &mut self.client_undeliverable_receiver
-    }
-
     pub fn client_proc(&self) -> &Proc {
         &self.client_proc
     }
@@ -508,6 +499,11 @@ impl ProcEvents {
                 Ok(event) = self.event_state.supervision_events.recv() => {
                     let (actor_id, actor_status) = event.clone().into_inner();
                     let Some(rank) = self.ranks.get(actor_id.proc_id()) else {
+                        let client_proc = ProcId(WorldId(format!("{}_manager", self.event_state.alloc.world_id().name())), 0);
+                        let client = client_proc.actor_id("client", 0);
+                        if client == actor_id {
+                            break Some(ProcEvent::Crashed(0, actor_status.to_string()));
+                        }
                         tracing::warn!("received supervision event for unmapped actor {}", actor_id);
                         continue;
                     };