mailbox: MailboxServers now handle undeliverables intrinsically (#791)

Summary:
Pull Request resolved: #791

D79739939 established that mailbox servers can handle undeliverable messages by forwarding them back to the sender. that diff improved return-handling, but the `return_handle` was still passed explicitly to `.serve()`.

however, inspecting how these return handles are constructed reveals that their creation can-and should-be handled internally by `.serve()` itself. this eliminates the need for `return_handle: PortHandle<Undeliverable<MessageEnvelope>>` in the `MailboxServer` interface altogether.

this change simplifies the interface and unlocks widespread cleanup and better default handling of undeliverable messages throughout the system.

Reviewed By: mariusae

Differential Revision: D79812085

fbshipit-source-id: c6c6a24cd48db1e662cc572f0f7069f8ed59b290

Author: shayne-fletcher

controller/src/lib.rs

@@ -638,7 +638,6 @@ mod tests {
     use hyperactor::mailbox::MailboxServer;
     use hyperactor::mailbox::PortHandle;
     use hyperactor::mailbox::PortReceiver;
-    use hyperactor::mailbox::monitored_return_handle;
     use hyperactor::message::IndexedErasedUnbound;
     use hyperactor::proc::Proc;
     use hyperactor::reference::GangId;
@@ -1545,9 +1544,7 @@ mod tests {
         let (local_proc_message_port, local_proc_message_receiver) = local_proc_mbox.open_port();
         local_proc_message_port.bind();
 
-        let _local_proc_serve_handle = local_proc_mbox
-            .clone()
-            .serve(local_proc_rx, monitored_return_handle());
+        let _local_proc_serve_handle = local_proc_mbox.clone().serve(local_proc_rx);
         (
             world_id,
             local_proc_id,

hyperactor/src/mailbox.rs

@@ -838,6 +838,43 @@ impl Future for MailboxServerHandle {
     }
 }
 
+// A `MailboxServer` (such as a router) can can receive a message
+// that couldn't reach its destination. We can use the fact that
+// servers are `MailboxSender`s to attempt to forward them back to
+// their senders.
+fn server_return_handle<T: MailboxServer>(server: T) -> PortHandle<Undeliverable<MessageEnvelope>> {
+    let (return_handle, mut rx) = undeliverable::new_undeliverable_port();
+
+    tokio::task::spawn(async move {
+        while let Ok(Undeliverable(mut envelope)) = rx.recv().await {
+            if let Ok(Undeliverable(e)) = envelope.deserialized::<Undeliverable<MessageEnvelope>>()
+            {
+                // A non-returnable undeliverable.
+                UndeliverableMailboxSender.post(e, monitored_return_handle());
+                continue;
+            }
+            envelope.try_set_error(DeliveryError::BrokenLink(
+                "message was undeliverable".to_owned(),
+            ));
+            server.post(
+                MessageEnvelope::new(
+                    envelope.sender().clone(),
+                    PortRef::<Undeliverable<MessageEnvelope>>::attest_message_port(
+                        envelope.sender(),
+                    )
+                    .port_id()
+                    .clone(),
+                    Serialized::serialize(&Undeliverable(envelope)).unwrap(),
+                    Attrs::new(),
+                ),
+                monitored_return_handle(),
+            );
+        }
+    });
+
+    return_handle
+}
+
 /// Serve a port on the provided [`channel::Rx`]. This dispatches all
 /// channel messages directly to the port.
 pub trait MailboxServer: MailboxSender + Clone + Sized + 'static {
@@ -847,8 +884,41 @@ pub trait MailboxServer: MailboxSender + Clone + Sized + 'static {
     fn serve(
         self,
         mut rx: impl channel::Rx<MessageEnvelope> + Send + 'static,
-        return_handle: PortHandle<Undeliverable<MessageEnvelope>>,
     ) -> MailboxServerHandle {
+        // A `MailboxServer` can can receive a message that couldn't
+        // reach its destination. We can use the fact that servers are
+        // `MailboxSender`s to attempt to forward them back to their
+        // senders.
+        let (return_handle, mut undeliverable_rx) = undeliverable::new_undeliverable_port();
+        let server = self.clone();
+        tokio::task::spawn(async move {
+            while let Ok(Undeliverable(mut envelope)) = undeliverable_rx.recv().await {
+                if let Ok(Undeliverable(e)) =
+                    envelope.deserialized::<Undeliverable<MessageEnvelope>>()
+                {
+                    // A non-returnable undeliverable.
+                    UndeliverableMailboxSender.post(e, monitored_return_handle());
+                    continue;
+                }
+                envelope.try_set_error(DeliveryError::BrokenLink(
+                    "message was undeliverable".to_owned(),
+                ));
+                server.post(
+                    MessageEnvelope::new(
+                        envelope.sender().clone(),
+                        PortRef::<Undeliverable<MessageEnvelope>>::attest_message_port(
+                            envelope.sender(),
+                        )
+                        .port_id()
+                        .clone(),
+                        Serialized::serialize(&Undeliverable(envelope)).unwrap(),
+                        Attrs::new(),
+                    ),
+                    monitored_return_handle(),
+                );
+            }
+        });
+
         let (stopped_tx, mut stopped_rx) = watch::channel(false);
         let join_handle = tokio::spawn(async move {
             let mut detached = false;
@@ -2399,45 +2469,6 @@ impl MailboxSender for DialMailboxRouter {
     }
 }
 
-/// A `MailboxServer` (such as a router) can can receive a message
-/// that couldn't reach its destination. We can use the fact that
-/// servers are `MailboxSender`s to attempt to forward them back to
-/// their senders.
-pub fn server_return_handle<T: MailboxServer>(
-    server: T,
-) -> PortHandle<Undeliverable<MessageEnvelope>> {
-    let (return_handle, mut rx) = undeliverable::new_undeliverable_port();
-
-    tokio::task::spawn(async move {
-        while let Ok(Undeliverable(mut envelope)) = rx.recv().await {
-            if let Ok(Undeliverable(e)) = envelope.deserialized::<Undeliverable<MessageEnvelope>>()
-            {
-                // A non-returnable undeliverable.
-                UndeliverableMailboxSender.post(e, monitored_return_handle());
-                continue;
-            }
-            envelope.try_set_error(DeliveryError::BrokenLink(
-                "message was undeliverable".to_owned(),
-            ));
-            server.post(
-                MessageEnvelope::new(
-                    envelope.sender().clone(),
-                    PortRef::<Undeliverable<MessageEnvelope>>::attest_message_port(
-                        envelope.sender(),
-                    )
-                    .port_id()
-                    .clone(),
-                    Serialized::serialize(&Undeliverable(envelope)).unwrap(),
-                    Attrs::new(),
-                ),
-                monitored_return_handle(),
-            );
-        }
-    });
-
-    return_handle
-}
-
 /// A MailboxSender that reports any envelope as undeliverable due to
 /// routing failure.
 #[derive(Debug)]
@@ -2657,7 +2688,7 @@ mod tests {
     async fn test_local_client_server() {
         let mbox = Mailbox::new_detached(id!(test[0].actor0));
         let (tx, rx) = channel::local::new();
-        let serve_handle = mbox.clone().serve(rx, monitored_return_handle());
+        let serve_handle = mbox.clone().serve(rx);
         let client = MailboxClient::new(tx);
 
         let (port, receiver) = mbox.open_once_port::<u64>();
@@ -2688,7 +2719,7 @@ mod tests {
             .unwrap();
         let tx = dial::<MessageEnvelope>(src_to_dst).unwrap();
         let mbox = Mailbox::new_detached(id!(test[0].actor0));
-        let serve_handle = mbox.clone().serve(rx, monitored_return_handle());
+        let serve_handle = mbox.clone().serve(rx);
         let client = MailboxClient::new(tx);
         let (port, receiver) = mbox.open_once_port::<u64>();
         let port = port.bind();
@@ -2798,7 +2829,7 @@ mod tests {
             let (addr, rx) = channel::serve(ChannelAddr::any(ChannelTransport::Local))
                 .await
                 .unwrap();
-            let handle = (*mbox).clone().serve(rx, monitored_return_handle());
+            let handle = (*mbox).clone().serve(rx);
             handles.push(handle);
 
             eprintln!("{}: {}", mbox.actor_id(), addr);

hyperactor_mesh/src/alloc.rs

@@ -459,9 +459,7 @@ pub(crate) mod testing {
             .unwrap();
         let router =
             DialMailboxRouter::new_with_default((UndeliverableMailboxSender {}).into_boxed());
-        router
-            .clone()
-            .serve(router_rx, mailbox::monitored_return_handle());
+        router.clone().serve(router_rx);
 
         let client_proc_id = ProcId(WorldId("test_stuck".to_string()), 0);
         let (client_proc_addr, client_rx) =
@@ -470,9 +468,7 @@ pub(crate) mod testing {
             client_proc_id.clone(),
             BoxedMailboxSender::new(router.clone()),
         );
-        client_proc
-            .clone()
-            .serve(client_rx, mailbox::monitored_return_handle());
+        client_proc.clone().serve(client_rx);
         router.bind(client_proc_id.clone().into(), client_proc_addr);
         (
             router,

hyperactor_mesh/src/alloc/local.rs

@@ -187,7 +187,7 @@ impl Alloc for LocalAlloc {
                     };
 
                     // Undeliverable messages get forwarded to the mesh agent.
-                    let handle = proc.clone().serve(proc_rx, mesh_agent.port());
+                    let handle = proc.clone().serve(proc_rx);
 
                     self.procs.insert(
                         rank,

hyperactor_mesh/src/alloc/remoteprocess.rs

@@ -35,7 +35,6 @@ use hyperactor::clock::RealClock;
 use hyperactor::config;
 use hyperactor::mailbox::DialMailboxRouter;
 use hyperactor::mailbox::MailboxServer;
-use hyperactor::mailbox::server_return_handle;
 use hyperactor::reference::Reference;
 use hyperactor::serde_json;
 use mockall::automock;
@@ -284,9 +283,7 @@ impl RemoteProcessAllocator {
                 }
             };
         let router = DialMailboxRouter::new();
-        let mailbox_handle = router
-            .clone()
-            .serve(forwarder_rx, server_return_handle(router.clone()));
+        let mailbox_handle = router.clone().serve(forwarder_rx);
         tracing::info!("started forwarder on: {}", forwarder_addr);
 
         // Check if we need to write TORCH_ELASTIC_CUSTOM_HOSTNAMES_LIST_FILE

hyperactor_mesh/src/bootstrap.rs

@@ -168,8 +168,7 @@ pub async fn bootstrap() -> anyhow::Error {
                     let (proc, mesh_agent) = MeshAgent::bootstrap(proc_id.clone()).await?;
                     let (proc_addr, proc_rx) =
                         channel::serve(ChannelAddr::any(listen_transport)).await?;
-                    // Undeliverable messages get forwarded to the mesh agent.
-                    let handle = proc.clone().serve(proc_rx, mesh_agent.port());
+                    let handle = proc.clone().serve(proc_rx);
                     drop(handle); // linter appeasement; it is safe to drop this future
                     tx.send(Process2Allocator(
                         bootstrap_index,

hyperactor_mesh/src/logging.rs

@@ -840,7 +840,6 @@ mod tests {
     use hyperactor::channel::ChannelTx;
     use hyperactor::channel::Tx;
     use hyperactor::id;
-    use hyperactor::mailbox;
     use hyperactor::mailbox::BoxedMailboxSender;
     use hyperactor::mailbox::DialMailboxRouter;
     use hyperactor::mailbox::MailboxServer;
@@ -858,8 +857,7 @@ mod tests {
             .await
             .unwrap();
         let proc = Proc::new(id!(client[0]), BoxedMailboxSender::new(router.clone()));
-        proc.clone()
-            .serve(client_rx, mailbox::monitored_return_handle());
+        proc.clone().serve(client_rx);
         router.bind(id!(client[0]).into(), proc_addr.clone());
         let client = proc.attach("client").unwrap();
 

hyperactor_mesh/src/proc_mesh.rs

@@ -37,7 +37,6 @@ use hyperactor::mailbox::MailboxServer;
 use hyperactor::mailbox::MessageEnvelope;
 use hyperactor::mailbox::PortReceiver;
 use hyperactor::mailbox::Undeliverable;
-use hyperactor::mailbox::server_return_handle;
 use hyperactor::metrics;
 use hyperactor::proc::Proc;
 use hyperactor::reference::ProcId;
@@ -210,9 +209,7 @@ impl ProcMesh {
             // Work around for Allocs that have more than one world.
             world_ids.insert(proc_id.world_id().clone());
         }
-        router
-            .clone()
-            .serve(router_rx, server_return_handle(router.clone()));
+        router.clone().serve(router_rx);
 
         // Set up a client proc for the mesh itself, so that we can attach ourselves
         // to it, and communicate with the agents. We wire it into the same router as
@@ -226,9 +223,7 @@ impl ProcMesh {
             client_proc_id.clone(),
             BoxedMailboxSender::new(router.clone()),
         );
-        client_proc
-            .clone()
-            .serve(client_rx, server_return_handle(client_proc.clone()));
+        client_proc.clone().serve(client_rx);
         router.bind(client_proc_id.clone().into(), client_proc_addr.clone());
 
         // Bind this router to the global router, to enable cross-mesh routing.

hyperactor_mesh/src/proc_mesh/mesh_agent.rs

@@ -37,14 +37,12 @@ use hyperactor::channel::ChannelAddr;
 use hyperactor::clock::Clock;
 use hyperactor::clock::RealClock;
 use hyperactor::mailbox::BoxedMailboxSender;
-use hyperactor::mailbox::DeliveryError;
 use hyperactor::mailbox::DialMailboxRouter;
 use hyperactor::mailbox::IntoBoxedMailboxSender;
 use hyperactor::mailbox::MailboxClient;
 use hyperactor::mailbox::MailboxSender;
 use hyperactor::mailbox::MessageEnvelope;
 use hyperactor::mailbox::Undeliverable;
-use hyperactor::mailbox::UndeliverableMessageError;
 use hyperactor::proc::Proc;
 use hyperactor::supervision::ActorSupervisionEvent;
 use serde::Deserialize;
@@ -162,30 +160,6 @@ impl Actor for MeshAgent {
         self.proc.set_supervision_coordinator(this.port())?;
         Ok(())
     }
-
-    // This is an override of the default actor behavior.
-    async fn handle_undeliverable_message(
-        &mut self,
-        cx: &Instance<Self>,
-        undelivered: Undeliverable<MessageEnvelope>,
-    ) -> Result<(), anyhow::Error> {
-        let Undeliverable(ref envelope) = undelivered;
-        tracing::debug!("took charge of a message not delivered: {}", envelope);
-
-        let sender = envelope.sender().clone();
-        if cx.self_id() == &sender {
-            anyhow::bail!(UndeliverableMessageError::delivery_failure(envelope));
-        }
-
-        let mut envelope = envelope.clone();
-        let return_port = PortRef::attest_message_port(&sender);
-        return_port.send(cx, undelivered).map_err(|err| {
-            envelope.try_set_error(DeliveryError::BrokenLink(format!("send failure: {err}")));
-            UndeliverableMessageError::return_failure(&envelope)
-        })?;
-
-        Ok(())
-    }
 }
 
 #[async_trait]

hyperactor_multiprocess/src/proc_actor.rs

@@ -44,7 +44,6 @@ use hyperactor::mailbox::MailboxAdminMessageHandler;
 use hyperactor::mailbox::MailboxClient;
 use hyperactor::mailbox::MailboxServer;
 use hyperactor::mailbox::MailboxServerHandle;
-use hyperactor::mailbox::monitored_return_handle;
 use hyperactor::mailbox::open_port;
 use hyperactor::proc::ActorLedgerSnapshot;
 use hyperactor::proc::Proc;
@@ -419,7 +418,7 @@ impl ProcActor {
         lifecycle_mode: ProcLifecycleMode,
     ) -> Result<BootstrappedProc, anyhow::Error> {
         let (local_addr, rx) = channel::serve(listen_addr).await?;
-        let mailbox_handle = proc.clone().serve(rx, monitored_return_handle());
+        let mailbox_handle = proc.clone().serve(rx);
         let (state_tx, mut state_rx) = watch::channel(ProcState::AwaitingJoin);
 
         let handle = match proc