[hyperactor]: mailbox: CanSend for Mailbox use mailbox port for returns

Pull Request resolved: #303

the `CanSend` for `Mailbox` requires an undeliverable message port be self-bound.  legacy tests were updated to be compliant with the new requirement.

removing `monitored_return_handle()` from `CanSend` is a clear improvement.

ghstack-source-id: 291502157

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

**NOTE FOR REVIEWERS**: This PR has internal Meta-specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D76926674/)!

Author: Shayne Fletcher

controller/src/lib.rs

@@ -662,6 +662,7 @@ mod tests {
         let (worker, worker_ref, mut worker_rx) = proc
             .attach_actor::<WorkerActor, WorkerMessage>("worker")
             .unwrap();
+
         IndexedErasedUnbound::<WorkerMessage>::bind_for_test_only(worker_ref.clone(), &worker)
             .unwrap();
 

hyperactor/src/mailbox/mod.rs

@@ -1082,6 +1082,33 @@ impl Mailbox {
         MailboxError::new(self.state.actor_id.clone(), err)
     }
 
+    /// Look up the `UnboundedPortSender` for the port associated with
+    /// the message type `M`, as defined by `<M as Named>::port()`.
+    ///
+    /// This performs a dynamic downcast to recover the original
+    /// sender type. Returns `None` if the port is not bound or if the
+    /// type does not match.
+    ///
+    /// # Panics
+    /// Panics in debug mode if the port is bound but the stored
+    /// `port_id` does not match the expected one.
+    pub(crate) fn lookup_sender<M: RemoteMessage>(&self) -> Option<UnboundedPortSender<M>> {
+        let port_index = M::port();
+        self.state.ports.get(&port_index).and_then(|boxed| {
+            boxed
+                .as_any()
+                .downcast_ref::<UnboundedSender<M>>()
+                .map(|s| {
+                    debug_assert_eq!(
+                        s.port_id,
+                        self.actor_id().port_id(port_index),
+                        "port_id mismatch in downcasted UnboundedSender"
+                    );
+                    s.sender.clone()
+                })
+        })
+    }
+
     fn bind<M: RemoteMessage>(&self, handle: &PortHandle<M>) -> PortRef<M> {
         assert_eq!(
             handle.mailbox.actor_id(),
@@ -1210,8 +1237,22 @@ impl MailboxSender for Mailbox {
 
 impl cap::sealed::CanSend for Mailbox {
     fn post(&self, dest: PortId, data: Serialized) {
+        let return_handle = self
+            .lookup_sender::<Undeliverable<MessageEnvelope>>()
+            .map_or_else(
+                || {
+                    let bt = std::backtrace::Backtrace::capture();
+                    tracing::warn!(
+                        actor_id = ?self.actor_id(),
+                        backtrace = ?bt,
+                        "Mailbox attempted to post a message without binding Undeliverable<MessageEnvelope>"
+                    );
+                    monitored_return_handle()
+                },
+                |sender| PortHandle::new(self.clone(), u64::MAX, sender),
+            );
         let envelope = MessageEnvelope::new(self.actor_id().clone(), dest, data);
-        MailboxSender::post(self, envelope, monitored_return_handle());
+        MailboxSender::post(self, envelope, return_handle);
     }
 }
 
@@ -1592,6 +1633,8 @@ pub struct SerializedSenderError {
 ///   - It abstracts over [`Port`]s and [`OncePort`]s, by dynamically tracking the
 ///     validity of the underlying port.
 trait SerializedSender: Send + Sync {
+    fn as_any(&self) -> &dyn Any;
+
     /// Send a serialized message. SerializedSender will deserialize the
     /// message (failing if it fails to deserialize), and then send the
     /// resulting message on the underlying port.
@@ -1603,7 +1646,7 @@ trait SerializedSender: Send + Sync {
 }
 
 /// A sender to an M-typed unbounded port.
-enum UnboundedPortSender<M: Message> {
+pub(crate) enum UnboundedPortSender<M: Message> {
     /// Send directly to the mpsc queue.
     Mpsc(mpsc::UnboundedSender<M>),
     /// Use the provided function to enqueue the item.
@@ -1675,6 +1718,10 @@ impl<M: Message> Clone for UnboundedSender<M> {
 }
 
 impl<M: RemoteMessage> SerializedSender for UnboundedSender<M> {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
     fn send_serialized(&self, serialized: Serialized) -> Result<bool, SerializedSenderError> {
         match serialized.deserialized() {
             Ok(message) => {
@@ -1757,6 +1804,10 @@ impl<M: Message> Clone for OnceSender<M> {
 }
 
 impl<M: RemoteMessage> SerializedSender for OnceSender<M> {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
     fn send_serialized(&self, serialized: Serialized) -> Result<bool, SerializedSenderError> {
         match serialized.deserialized() {
             Ok(message) => self.send_once(message).map_err(|e| SerializedSenderError {
@@ -1781,6 +1832,10 @@ struct UntypedUnboundedSender {
 }
 
 impl SerializedSender for UntypedUnboundedSender {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
     fn send_serialized(&self, serialized: Serialized) -> Result<bool, SerializedSenderError> {
         (self.sender)(serialized).map_err(|(data, err)| SerializedSenderError {
             data,

hyperactor_mesh/src/actor_mesh.rs

@@ -541,7 +541,6 @@ mod tests {
     use hyperactor::ProcId;
     use hyperactor::WorldId;
     use hyperactor::id;
-    use hyperactor::mailbox::MessageEnvelope;
     use hyperactor::mailbox::Undeliverable;
     use hyperactor::message::Bind;
     use hyperactor::message::Unbind;
@@ -821,12 +820,13 @@ mod tests {
                     .unwrap();
 
                 let name = alloc.name().to_string();
-                let mesh = ProcMesh::allocate(alloc).await.unwrap();
-                let unmonitored_reply_to = mesh.client().open_port::<usize>().0.bind();
-                let (undeliverable_messages, mut undeliverable_rx) = mesh.client().open_port::<Undeliverable<MessageEnvelope>>();
-                undeliverable_messages.bind_to(Undeliverable::<MessageEnvelope>::port());
+                let mut mesh = ProcMesh::allocate(alloc).await.unwrap();
+                let mut undeliverable_rx = mesh
+                    .take_client_undeliverable_receiver()
+                    .expect("client_undeliverable_receiver should be available");
 
                 // 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();
 
@@ -870,10 +870,12 @@ 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.take_client_undeliverable_receiver().unwrap();
 
-            let (undeliverable_msg_tx, mut undeliverable_msg_rx) = mesh.client().open_port();
-            let ping_pong_actor_params =
-                PingPongActorParams::new(undeliverable_msg_tx.bind(), None);
+            let ping_pong_actor_params = PingPongActorParams::new(
+                PortRef::attest_message_port(mesh.client().actor_id()),
+                None,
+            );
             let actor_mesh: RootActorMesh<PingPongActor> = mesh
                 .spawn::<PingPongActor>("ping-pong", &ping_pong_actor_params)
                 .await

hyperactor_mesh/src/proc_mesh/mod.rs

@@ -15,6 +15,7 @@ use async_trait::async_trait;
 use hyperactor::Actor;
 use hyperactor::ActorRef;
 use hyperactor::Mailbox;
+use hyperactor::Named;
 use hyperactor::RemoteMessage;
 use hyperactor::WorldId;
 use hyperactor::actor::RemoteActor;
@@ -28,7 +29,9 @@ use hyperactor::mailbox::BoxedMailboxSender;
 use hyperactor::mailbox::DialMailboxRouter;
 use hyperactor::mailbox::MailboxRouter;
 use hyperactor::mailbox::MailboxServer;
+use hyperactor::mailbox::MessageEnvelope;
 use hyperactor::mailbox::PortReceiver;
+use hyperactor::mailbox::Undeliverable;
 use hyperactor::proc::Proc;
 use hyperactor::reference::ProcId;
 use hyperactor::reference::Reference;
@@ -75,6 +78,7 @@ 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>>,
 }
 
@@ -201,9 +205,13 @@ impl ProcMesh {
         let supervisor = client_proc.attach("supervisor")?;
         let (supervison_port, supervision_events) = supervisor.open_port();
 
-        // Now, configure the full mesh, so that the local agents are wired up to
-        // our router.
+        // 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.
         let client = client_proc.attach("client")?;
+        let (undeliverable_messages, client_undeliverable_receiver) =
+            client.open_port::<Undeliverable<MessageEnvelope>>();
+        undeliverable_messages.bind_to(Undeliverable::<MessageEnvelope>::port());
 
         // Map of procs -> channel addresses
         let address_book: HashMap<_, _> = running
@@ -285,6 +293,7 @@ impl ProcMesh {
                 .collect(),
             client_proc,
             client,
+            client_undeliverable_receiver: Some(client_undeliverable_receiver),
             comm_actors,
         })
     }
@@ -371,6 +380,14 @@ impl ProcMesh {
         &self.client
     }
 
+    /// Used to get the `PortReceiver<Undeliverable<MessageEnvelope>>`
+    /// bound in the client mailbox.
+    pub fn take_client_undeliverable_receiver(
+        &mut self,
+    ) -> Option<PortReceiver<Undeliverable<MessageEnvelope>>> {
+        self.client_undeliverable_receiver.take()
+    }
+
     pub fn client_proc(&self) -> &Proc {
         &self.client_proc
     }