return undeliverable messages received by servers (#781)

Summary:
Pull Request resolved: #781

while working on D79478197, mariusae noticed that there are circumstances where undeliverable messages get returned to mailbox servers like routers for example. handling returns of such messages to their original senders has been a long standing gap (the stock handling until now has been to log and continue) which this diff goes towards addressing.

this diff improves undeliverable message handling in  `hyperactor_mesh`. the main change is replacing the router-specific return handler with a generic `server_return_handle` that works with any `MailboxServer`. when messages can't be delivered, the system now attempts to return them to the sender and tries also to deal with nested undeliverable messages should, for whatever reason, that come up.

i use an environment variable `HYPERACTOR_MESH_ROUTER_NO_GLOBAL_FALLBACK` (eschewing the attempt at config in D79725191) for router failure injection in the newly added `test_router_undeliverable_return`. this test verifies supervision events are generated when routers fail.

Reviewed By: mariusae

Differential Revision: D79739939

fbshipit-source-id: 8f9d48954dd26b47fd54837220235008184000c3

Author: shayne-fletcher

hyperactor/src/mailbox.rs

@@ -2399,6 +2399,45 @@ 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)]

hyperactor_mesh/src/actor_mesh.rs

@@ -623,16 +623,39 @@ pub(crate) mod test_util {
     #[async_trait]
     impl Handler<Echo> for ProxyActor {
         async fn handle(&mut self, cx: &Context<Self>, message: Echo) -> Result<(), anyhow::Error> {
-            let actor = self.actor_mesh.get(0).unwrap();
-
-            // For now, we reply directly to the client.
-            // We will support directly wiring up the meshes later.
-            let (tx, mut rx) = cx.open_port();
-
-            actor.send(cx, Echo(message.0, tx.bind()))?;
-            message.1.send(cx, rx.recv().await.unwrap())?;
+            if std::env::var("HYPERACTOR_MESH_ROUTER_NO_GLOBAL_FALLBACK").is_err() {
+                // test_proxy_mesh
+
+                let actor = self.actor_mesh.get(0).unwrap();
+
+                // For now, we reply directly to the client.
+                // We will support directly wiring up the meshes later.
+                let (tx, mut rx) = cx.open_port();
+
+                actor.send(cx, Echo(message.0, tx.bind()))?;
+                message.1.send(cx, rx.recv().await.unwrap())?;
+
+                Ok(())
+            } else {
+                // test_router_undeliverable_return
+
+                let actor: ActorRef<_> = self.actor_mesh.get(0).unwrap();
+                let (tx, mut rx) = cx.open_port::<String>();
+                actor.send(cx, Echo(message.0, tx.bind()))?;
+
+                use tokio::time::Duration;
+                use tokio::time::timeout;
+                #[allow(clippy::disallowed_methods)]
+                match timeout(Duration::from_secs(1), rx.recv()).await {
+                    Ok(_) => message
+                        .1
+                        .send(cx, "the impossible happened".to_owned())
+                        .unwrap(),
+                    _ => (),
+                }
 
-            Ok(())
+                Ok(())
+            }
         }
     }
 }
@@ -677,8 +700,6 @@ mod tests {
                 use $crate::alloc::AllocSpec;
                 use $crate::alloc::Allocator;
 
-                hyperactor_telemetry::initialize_logging(hyperactor::clock::ClockKind::default());
-
                 use ndslice::extent;
 
                 let alloc = $allocator
@@ -693,7 +714,12 @@ mod tests {
                 let proxy_actor = actor_mesh.get(0).unwrap();
                 let (tx, mut rx) = actor_mesh.open_port::<String>();
                 proxy_actor.send(proc_mesh.client(), Echo("hello!".to_owned(), tx.bind())).unwrap();
-                assert_eq!(rx.recv().await.unwrap(), "hello!");
+
+                #[allow(clippy::disallowed_methods)]
+                match tokio::time::timeout(tokio::time::Duration::from_secs(3), rx.recv()).await {
+                    Ok(msg) => assert_eq!(&msg.unwrap(), "hello!"),
+                    Err(_) =>  assert!(false),
+                }
             }
 
             #[tokio::test]
@@ -1287,6 +1313,62 @@ mod tests {
                 assert_eq!(event.actor_id.name(), &actor_mesh.name);
             }
         }
+
+        // Set this test only for `mod process` because it relies on a
+        // trick to emulate router failure that only works when using
+        // non-local allocators.
+        #[cfg(fbcode_build)]
+        #[tokio::test]
+        async fn test_router_undeliverable_return() {
+            // Test that an undeliverable message received by a
+            // router results in actor mesh supervision events.
+            use ndslice::extent;
+
+            use super::test_util::*;
+            use crate::alloc::AllocSpec;
+            use crate::alloc::Allocator;
+
+            let alloc = process_allocator()
+                .allocate(AllocSpec {
+                    extent: extent! { replica = 1 },
+                    constraints: Default::default(),
+                })
+                .await
+                .unwrap();
+
+            // SAFETY: Not multithread safe.
+            unsafe { std::env::set_var("HYPERACTOR_MESH_ROUTER_NO_GLOBAL_FALLBACK", "1") };
+
+            let mut proc_mesh = ProcMesh::allocate(alloc).await.unwrap();
+            let mut proc_events = proc_mesh.events().unwrap();
+            let mut actor_mesh: RootActorMesh<'_, ProxyActor> =
+                { proc_mesh.spawn("proxy", &()).await.unwrap() };
+            let mut actor_events = actor_mesh.events().unwrap();
+
+            let proxy_actor = actor_mesh.get(0).unwrap();
+            let (tx, mut rx) = actor_mesh.open_port::<String>();
+            proxy_actor
+                .send(proc_mesh.client(), Echo("hello!".to_owned(), tx.bind()))
+                .unwrap();
+
+            #[allow(clippy::disallowed_methods)]
+            match tokio::time::timeout(tokio::time::Duration::from_secs(3), rx.recv()).await {
+                Ok(_) => panic!("the impossible happened"),
+                Err(_) => {
+                    assert_matches!(
+                        proc_events.next().await.unwrap(),
+                        ProcEvent::Crashed(0, reason) if reason.contains("undeliverable")
+                    );
+                    assert_eq!(
+                        actor_events.next().await.unwrap().actor_id.name(),
+                        &actor_mesh.name
+                    );
+                }
+            }
+
+            // SAFETY: Not multithread safe.
+            unsafe { std::env::remove_var("HYPERACTOR_MESH_ROUTER_NO_GLOBAL_FALLBACK") };
+        }
     }
 
     mod sim {

hyperactor_mesh/src/alloc/remoteprocess.rs

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

hyperactor_mesh/src/proc_mesh.rs

@@ -37,6 +37,7 @@ 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;
@@ -211,7 +212,7 @@ impl ProcMesh {
         }
         router
             .clone()
-            .serve(router_rx, mailbox::custom_monitored_return_handle("router"));
+            .serve(router_rx, server_return_handle(router.clone()));
 
         // 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
@@ -225,10 +226,9 @@ impl ProcMesh {
             client_proc_id.clone(),
             BoxedMailboxSender::new(router.clone()),
         );
-        client_proc.clone().serve(
-            client_rx,
-            mailbox::custom_monitored_return_handle("client proc"),
-        );
+        client_proc
+            .clone()
+            .serve(client_rx, server_return_handle(client_proc.clone()));
         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

@@ -208,8 +208,17 @@ impl MeshAgentMessageHandler for MeshAgent {
         // Wire up the local proc to the global (process) router. This ensures that child
         // meshes are reachable from any actor created by this mesh.
         let client = MailboxClient::new(channel::dial(forwarder)?);
-        let default = super::global_router().fallback(client.into_boxed());
-        let router = DialMailboxRouter::new_with_default(default.into_boxed());
+
+        // `HYPERACTOR_MESH_ROUTER_CONFIG_NO_GLOBAL_FALLBACK` may be
+        // set as a means of failure injection in the testing of
+        // supervision codepaths.
+        let router = if std::env::var("HYPERACTOR_MESH_ROUTER_NO_GLOBAL_FALLBACK").is_err() {
+            let default = super::global_router().fallback(client.into_boxed());
+            DialMailboxRouter::new_with_default(default.into_boxed())
+        } else {
+            DialMailboxRouter::new_with_default(client.into_boxed())
+        };
+
         for (proc_id, addr) in address_book {
             router.bind(proc_id.into(), addr);
         }