[hyperactor] mesh: implement local proc bypasss

This change allows procs local to a host to bypass the host multiplexer. This should improve multicast delivery, as local comm actors do not need to traverse the same (frontend) bottleneck.

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

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

ghstack-source-id: 314347091
Pull Request resolved: #1442

Author: mariusae

hyperactor/src/mailbox.rs

@@ -164,7 +164,8 @@ pub type Data = Vec<u8>;
     Deserialize,
     Named,
     Clone,
-    PartialEq
+    PartialEq,
+    Eq
 )]
 pub enum DeliveryError {
     /// The destination address is not reachable.

hyperactor/src/mailbox/undeliverable.rs

@@ -32,6 +32,13 @@ use crate::supervision::ActorSupervisionEvent;
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Named)]
 pub struct Undeliverable<M: Message>(pub M);
 
+impl<M: Message> Undeliverable<M> {
+    /// Return the inner M-typed message.
+    pub fn into_inner(self) -> M {
+        self.0
+    }
+}
+
 // Port handle and receiver for undeliverable messages.
 pub(crate) fn new_undeliverable_port() -> (
     PortHandle<Undeliverable<MessageEnvelope>>,

hyperactor_mesh/src/bootstrap.rs

@@ -15,6 +15,7 @@ use std::future;
 use std::io;
 use std::io::Write;
 use std::os::unix::process::ExitStatusExt;
+use std::path::Path;
 use std::path::PathBuf;
 use std::process::Stdio;
 use std::sync::Arc;
@@ -34,6 +35,7 @@ use hyperactor::ProcId;
 use hyperactor::attrs::Attrs;
 use hyperactor::channel;
 use hyperactor::channel::ChannelAddr;
+use hyperactor::channel::ChannelError;
 use hyperactor::channel::ChannelTransport;
 use hyperactor::channel::Rx;
 use hyperactor::channel::Tx;
@@ -48,11 +50,14 @@ use hyperactor::host::HostError;
 use hyperactor::host::ProcHandle;
 use hyperactor::host::ProcManager;
 use hyperactor::host::TerminateSummary;
+use hyperactor::mailbox::IntoBoxedMailboxSender;
+use hyperactor::mailbox::MailboxClient;
 use hyperactor::mailbox::MailboxServer;
 use hyperactor::proc::Proc;
 use libc::c_int;
 use serde::Deserialize;
 use serde::Serialize;
+use tempfile::TempDir;
 use tokio::process::Child;
 use tokio::process::Command;
 use tokio::sync::oneshot;
@@ -65,6 +70,8 @@ use crate::v1;
 use crate::v1::host_mesh::mesh_agent::HostAgentMode;
 use crate::v1::host_mesh::mesh_agent::HostMeshAgent;
 
+mod mailbox;
+
 declare_attrs! {
     /// If enabled (default), bootstrap child processes install
     /// `PR_SET_PDEATHSIG(SIGKILL)` so the kernel reaps them if the
@@ -213,6 +220,10 @@ pub enum Bootstrap {
         backend_addr: ChannelAddr,
         /// The callback address used to indicate successful spawning.
         callback_addr: ChannelAddr,
+        /// Directory for storing proc socket files. Procs place their sockets
+        /// in this directory, so that they can be looked up by other procs
+        /// for direct transfer.
+        socket_dir_path: PathBuf,
         /// Optional config snapshot (`hyperactor::config::Attrs`)
         /// captured by the parent. If present, the child installs it
         /// as the `Runtime` layer so the parent's effective config
@@ -325,6 +336,7 @@ impl Bootstrap {
                 proc_id,
                 backend_addr,
                 callback_addr,
+                socket_dir_path,
                 config,
             } => {
                 if let Some(attrs) = config {
@@ -344,15 +356,39 @@ impl Bootstrap {
                     eprintln!("(bootstrap) PDEATHSIG disabled via config");
                 }
 
-                let result =
-                    host::spawn_proc(proc_id, backend_addr, callback_addr, |proc| async move {
-                        ProcMeshAgent::boot_v1(proc).await
-                    })
-                    .await;
-                match result {
-                    Ok(_proc) => halt().await,
-                    Err(e) => e.into(),
-                }
+                let (local_addr, name) = ok!(proc_id
+                    .as_direct()
+                    .ok_or_else(|| anyhow::anyhow!("invalid proc id type")));
+                // TODO provide a direct way to construct these
+                let serve_addr = format!("unix:{}", socket_dir_path.join(name).display());
+                let serve_addr = serve_addr.parse().unwrap();
+
+                // The following is a modified host::spawn_proc to support direct
+                // dialing between local procs: 1) we bind each proc to a deterministic
+                // address in socket_dir_path; 2) we use LocalProcDialer to dial these
+                // addresses for local procs.
+                let proc_sender = mailbox::LocalProcDialer::new(
+                    local_addr.clone(),
+                    socket_dir_path,
+                    ok!(MailboxClient::dial(backend_addr)),
+                );
+
+                let proc = Proc::new(proc_id.clone(), proc_sender.into_boxed());
+
+                let agent_handle = ok!(ProcMeshAgent::boot_v1(proc.clone())
+                    .await
+                    .map_err(|e| HostError::AgentSpawnFailure(proc_id, e)));
+
+                // Finally serve the proc on the same transport as the backend address,
+                // and call back.
+                let (proc_addr, proc_rx) = ok!(channel::serve(serve_addr));
+                proc.clone().serve(proc_rx);
+                ok!(ok!(channel::dial(callback_addr))
+                    .send((proc_addr, agent_handle.bind::<ProcMeshAgent>()))
+                    .await
+                    .map_err(ChannelError::from));
+
+                halt().await
             }
             Bootstrap::Host {
                 addr,
@@ -370,7 +406,7 @@ impl Bootstrap {
                     Some(command) => command,
                     None => ok!(BootstrapCommand::current()),
                 };
-                let manager = BootstrapProcManager::new(command);
+                let manager = BootstrapProcManager::new(command).unwrap();
                 let (host, _handle) = ok!(Host::serve(manager, addr).await);
                 let addr = host.addr().clone();
                 let host_mesh_agent = ok!(host
@@ -1400,6 +1436,11 @@ pub struct BootstrapProcManager {
     /// exclusively in the [`Drop`] impl to send `SIGKILL` without
     /// needing async context.
     pid_table: Arc<std::sync::Mutex<HashMap<ProcId, u32>>>,
+
+    /// Directory for storing proc socket files. Procs place their sockets
+    /// in this directory, so that they can be looked up by other procs
+    /// for direct transfer.
+    socket_dir: TempDir,
 }
 
 impl Drop for BootstrapProcManager {
@@ -1449,12 +1490,13 @@ impl BootstrapProcManager {
     /// This is the general entry point when you want to manage procs
     /// backed by a specific binary path (e.g. a bootstrap
     /// trampoline).
-    pub(crate) fn new(command: BootstrapCommand) -> Self {
-        Self {
+    pub(crate) fn new(command: BootstrapCommand) -> Result<Self, io::Error> {
+        Ok(Self {
             command,
             children: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
             pid_table: Arc::new(std::sync::Mutex::new(HashMap::new())),
-        }
+            socket_dir: tempfile::tempdir()?,
+        })
     }
 
     /// The bootstrap command used to launch processes.
@@ -1626,6 +1668,7 @@ impl ProcManager for BootstrapProcManager {
             proc_id: proc_id.clone(),
             backend_addr,
             callback_addr,
+            socket_dir_path: self.socket_dir.path().to_owned(),
             config: Some(cfg),
         };
         let mut cmd = Command::new(&self.command.program);
@@ -2060,6 +2103,7 @@ mod tests {
                 proc_id: id!(foo[0]),
                 backend_addr: ChannelAddr::any(ChannelTransport::Tcp),
                 callback_addr: ChannelAddr::any(ChannelTransport::Unix),
+                socket_dir_path: PathBuf::from("notexist"),
                 config: None,
             },
         ];
@@ -2117,13 +2161,16 @@ mod tests {
         attrs[MESH_TAIL_LOG_LINES] = 123;
         attrs[MESH_BOOTSTRAP_ENABLE_PDEATHSIG] = false;
 
+        let socket_dir = tempfile::tempdir().unwrap();
+
         // Proc case
         {
             let original = Bootstrap::Proc {
                 proc_id: id!(foo[42]),
                 backend_addr: ChannelAddr::any(ChannelTransport::Unix),
                 callback_addr: ChannelAddr::any(ChannelTransport::Unix),
                 config: Some(attrs.clone()),
+                socket_dir_path: socket_dir.path().to_owned(),
             };
             let env_str = original.to_env_safe_string().expect("encode bootstrap");
             let decoded = Bootstrap::from_env_safe_string(&env_str).expect("decode bootstrap");
@@ -2163,14 +2210,13 @@ mod tests {
         use std::process::Stdio;
 
         use tokio::process::Command;
-        use tokio::time::Duration;
 
         // Manager; program path is irrelevant for this test.
         let command = BootstrapCommand {
             program: PathBuf::from("/bin/true"),
             ..Default::default()
         };
-        let manager = BootstrapProcManager::new(command);
+        let manager = BootstrapProcManager::new(command).unwrap();
 
         // Spawn a long-running child process (sleep 30) with
         // kill_on_drop(true).
@@ -2550,7 +2596,7 @@ mod tests {
             program: PathBuf::from("/bin/true"),
             ..Default::default()
         };
-        let manager = BootstrapProcManager::new(command);
+        let manager = BootstrapProcManager::new(command).unwrap();
 
         // Spawn a fast-exiting child.
         let mut cmd = Command::new("/bin/true");
@@ -2584,7 +2630,7 @@ mod tests {
             program: PathBuf::from("/bin/sleep"),
             ..Default::default()
         };
-        let manager = BootstrapProcManager::new(command);
+        let manager = BootstrapProcManager::new(command).unwrap();
 
         // Spawn a process that will live long enough to kill.
         let mut cmd = Command::new("/bin/sleep");
@@ -2701,7 +2747,8 @@ mod tests {
         let manager = BootstrapProcManager::new(BootstrapCommand {
             program: PathBuf::from("/bin/true"),
             ..Default::default()
-        });
+        })
+        .unwrap();
         let unknown = ProcId::Direct(ChannelAddr::any(ChannelTransport::Unix), "nope".into());
         assert!(manager.status(&unknown).await.is_none());
     }
@@ -2711,7 +2758,8 @@ mod tests {
         let manager = BootstrapProcManager::new(BootstrapCommand {
             program: PathBuf::from("/bin/sleep"),
             ..Default::default()
-        });
+        })
+        .unwrap();
 
         // Long-ish child so it's alive while we "steal" it.
         let mut cmd = Command::new("/bin/sleep");
@@ -2750,7 +2798,8 @@ mod tests {
         let manager = BootstrapProcManager::new(BootstrapCommand {
             program: PathBuf::from("/bin/sleep"),
             ..Default::default()
-        });
+        })
+        .unwrap();
 
         let mut cmd = Command::new("/bin/sleep");
         cmd.arg("5").stdout(Stdio::null()).stderr(Stdio::null());
@@ -3125,7 +3174,7 @@ mod tests {
             .unwrap();
         let (instance, _handle) = root.instance("client").unwrap();
 
-        let mgr = BootstrapProcManager::new(BootstrapCommand::test());
+        let mgr = BootstrapProcManager::new(BootstrapCommand::test()).unwrap();
         let (proc_id, backend_addr) = make_proc_id_and_backend_addr(&instance, "t_term").await;
         let handle = mgr
             .spawn(proc_id.clone(), backend_addr.clone())
@@ -3181,7 +3230,7 @@ mod tests {
             .unwrap();
         let (instance, _handle) = root.instance("client").unwrap();
 
-        let mgr = BootstrapProcManager::new(BootstrapCommand::test());
+        let mgr = BootstrapProcManager::new(BootstrapCommand::test()).unwrap();
 
         // Proc identity + host backend channel the child will dial.
         let (proc_id, backend_addr) = make_proc_id_and_backend_addr(&instance, "t_kill").await;
@@ -3380,7 +3429,8 @@ mod tests {
         let manager = BootstrapProcManager::new(BootstrapCommand {
             program: std::path::PathBuf::from("/bin/true"), // unused in this test
             ..Default::default()
-        });
+        })
+        .unwrap();
         manager.spawn_exit_monitor(proc_id.clone(), handle.clone());
 
         // Await terminal status and assert on exit code and stderr