[hyperactor_mesh]: comm: handle undeliverable messages in casts

provide `CommActor` with a custom `handle_undeliverable_message` implementation that routes delivery failures back to the original cast sender.

Differential Revision: [D77398378](https://our.internmc.facebook.com/intern/diff/D77398378/)

ghstack-source-id: 292942224
Pull Request resolved: #361

Author: Shayne Fletcher

hyperactor/src/mailbox/mod.rs

@@ -248,6 +248,11 @@ impl MessageEnvelope {
         &self.dest
     }
 
+    /// The message headers.
+    pub fn headers(&self) -> &Attrs {
+        &self.headers
+    }
+
     /// Tells whether this is a signal message.
     pub fn is_signal(&self) -> bool {
         self.dest.index() == Signal::port()
@@ -2141,8 +2146,15 @@ impl MailboxSender for WeakMailboxRouter {
     }
 }
 
-/// A serializable [`MailboxRouter`]. It keeps a serializable address book so that
-/// the mailbox sender can be recovered.
+/// A dynamic mailbox router that supports remote delivery.
+///
+/// `DialMailboxRouter` maintains a runtime address book mapping
+/// references to `ChannelAddr`s. It holds a cache of active
+/// connections and forwards messages to the appropriate
+/// `MailboxClient`.
+///
+/// Messages sent to unknown destinations are routed to the `default`
+/// sender, if present.
 #[derive(Debug, Clone)]
 pub struct DialMailboxRouter {
     address_book: Arc<RwLock<BTreeMap<Reference, ChannelAddr>>>,

hyperactor_mesh/Cargo.toml

@@ -40,7 +40,6 @@ buck-resources = "1"
 dir-diff = "0.3"
 maplit = "1.0"
 timed_test = { version = "0.0.0", path = "../timed_test" }
-tracing-test = { version = "0.2.3", features = ["no-env-filter"] }
 
 [lints]
 rust = { unexpected_cfgs = { check-cfg = ["cfg(fbcode_build)"], level = "warn" } }

hyperactor_mesh/src/actor_mesh.rs

@@ -974,12 +974,8 @@ mod tests {
             );
         }
 
-        // The intent is to emulate the behaviors of the Python
-        // interaction of T225230867 "process hangs when i send
-        // messages to a dead actor".
-        #[tracing_test::traced_test]
         #[tokio::test]
-        async fn test_behaviors_on_actor_error() {
+        async fn test_cast_failure() {
             use crate::alloc::ProcStopReason;
             use crate::proc_mesh::ProcEvent;
             use crate::sel;
@@ -994,6 +990,10 @@ 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
@@ -1016,14 +1016,22 @@ mod tests {
                 ProcEvent::Crashed(0, reason) if reason.contains("intentional error!")
             );
 
-            // Uncomment this to cause an infinite hang.
-            /*
-            let (reply_handle, mut reply_receiver) = actor_mesh.open_port();
-                actor_mesh
-                    .cast(sel!(*), GetRank(false, reply_handle.bind()))
-                    .unwrap();
-            let rank = reply_receiver.recv().await.unwrap();
-            */
+            // Cast the message.
+            let (reply_handle, _) = actor_mesh.open_port();
+            actor_mesh
+                .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
+                )
+            );
 
             // Stop the mesh.
             stop();

hyperactor_mesh/src/comm/mod.rs

@@ -20,9 +20,14 @@ use hyperactor::ActorRef;
 use hyperactor::Handler;
 use hyperactor::Instance;
 use hyperactor::Named;
+use hyperactor::PortRef;
 use hyperactor::WorldId;
 use hyperactor::attrs::Attrs;
+use hyperactor::attrs::declare_attrs;
 use hyperactor::data::Serialized;
+use hyperactor::mailbox::DeliveryError;
+use hyperactor::mailbox::Undeliverable;
+use hyperactor::mailbox::UndeliverableMessageError;
 use hyperactor::reference::UnboundPort;
 use ndslice::Slice;
 use ndslice::selection::routing::RoutingFrame;
@@ -145,6 +150,10 @@ impl CommActorMode {
     }
 }
 
+declare_attrs! {
+    pub attr CAST_ORIGINATING_SENDER: ActorId;
+}
+
 #[async_trait]
 impl Actor for CommActor {
     type Params = CommActorParams;
@@ -156,6 +165,46 @@ impl Actor for CommActor {
             mode: Default::default(),
         })
     }
+
+    // This is an override of the default actor behavior.
+    async fn handle_undeliverable_message(
+        &mut self,
+        this: &Instance<Self>,
+        undelivered: hyperactor::mailbox::Undeliverable<hyperactor::mailbox::MessageEnvelope>,
+    ) -> Result<(), anyhow::Error> {
+        let Undeliverable(mut message_envelope) = undelivered;
+
+        // 1. Case delivery failure at a "forwarding" step.
+        if let Ok(ForwardMessage { message, .. }) =
+            message_envelope.deserialized::<ForwardMessage>()
+        {
+            let sender = message.sender();
+            let return_port = PortRef::attest_message_port(sender);
+            return_port
+                .send(this, Undeliverable(message_envelope.clone()))
+                .map_err(|err| {
+                    message_envelope
+                        .try_set_error(DeliveryError::BrokenLink(format!("send failure: {err}")));
+                    UndeliverableMessageError::return_failure(&message_envelope)
+                })?;
+            return Ok(());
+        }
+
+        // 2. Case delivery failure at a "deliver here" step.
+        if let Some(sender) = message_envelope.headers().get(CAST_ORIGINATING_SENDER) {
+            let return_port = PortRef::attest_message_port(sender);
+            return_port
+                .send(this, Undeliverable(message_envelope.clone()))
+                .map_err(|err| {
+                    message_envelope
+                        .try_set_error(DeliveryError::BrokenLink(format!("send failure: {err}")));
+                    UndeliverableMessageError::return_failure(&message_envelope)
+                })?;
+            return Ok(());
+        }
+
+        unreachable!()
+    }
 }
 
 impl CommActor {
@@ -181,6 +230,7 @@ impl CommActor {
         seq: usize,
         last_seqs: &mut HashMap<usize, usize>,
     ) -> Result<()> {
+        tracing::info!("sf: comm handle_message sender {:#?}", sender);
         // Split ports, if any, and update message with new ports. In this
         // way, children actors will reply to this comm actor's ports, instead
         // of to the original ports provided by parent.
@@ -202,14 +252,21 @@ impl CommActor {
                 *r = CastRank(mode.self_rank(this.self_id()));
                 Ok(())
             })?;
+
+            // Preserve the original sender in the headers so that if
+            // it turns out the message is returned undeliverable, we
+            // can recover it.
+            let mut headers = Attrs::new();
+            headers.set(CAST_ORIGINATING_SENDER, message.sender().clone());
+
             // TODO(pzhang) split reply ports so children can reply to this comm
             // actor instead of parent.
             this.post(
                 this.self_id()
                     .proc_id()
                     .actor_id(message.dest_port().actor_name(), 0)
                     .port_id(message.dest_port().port()),
-                Attrs::new(),
+                headers,
                 Serialized::serialize(message.data())?,
             );
         }

hyperactor_mesh/src/comm/multicast.rs

@@ -75,6 +75,10 @@ impl CastMessageEnvelope {
         }
     }
 
+    pub(crate) fn sender(&self) -> &ActorId {
+        &self.sender
+    }
+
     pub(crate) fn dest_port(&self) -> &DestinationPort {
         &self.dest_port
     }